[OMNIML-6029] [4/5] llm_sparsity: load the model from local disk, downloading it first if needed

examples/llm_sparsity/attention_sparsity/hf_sa.py

@@ -27,7 +27,7 @@
 
 import modelopt.torch.opt as mto
 import modelopt.torch.sparsity.attention_sparsity as mtsa
-from modelopt.torch.export import export_hf_checkpoint
+from modelopt.torch.export import ensure_local_checkpoint, export_hf_checkpoint
 from modelopt.torch.sparsity.attention_sparsity.config import (
     SKIP_SOFTMAX_CALIB,
     SKIP_SOFTMAX_CALIB_SPARSE24,
@@ -146,14 +146,18 @@ def main(args):
     np.random.seed(RAND_SEED)
     launch_memory_monitor()
 
-    print(f"Loading model: {args.pyt_ckpt_path}")
+    # --pyt_ckpt_path stays as given. Later steps read the local copy of the whole checkpoint in
+    # local_checkpoint_path, since exporters read the source's files from local disk only;
+    # hub_model_id is the Hub ID, or None for a local path.
+    args.hub_model_id, args.local_checkpoint_path = ensure_local_checkpoint(args.pyt_ckpt_path)
+    print(f"Loading model: {args.local_checkpoint_path}")
 
     # No need to specify attn_implementation here — mtsa.sparsify() sets it
     # automatically ("eager" for pytorch backend, "modelopt_triton" for triton).
     model = AutoModelForCausalLM.from_pretrained(
-        args.pyt_ckpt_path, attn_implementation="eager", dtype="auto", device_map="auto"
+        args.local_checkpoint_path, attn_implementation="eager", dtype="auto", device_map="auto"
     )
-    tokenizer = AutoTokenizer.from_pretrained(args.pyt_ckpt_path)
+    tokenizer = AutoTokenizer.from_pretrained(args.local_checkpoint_path)
 
     # Set pad token if not set
     if tokenizer.pad_token is None:
@@ -232,10 +236,11 @@ def main(args):
         export_dir = Path(args.export_dir)
         export_dir.mkdir(parents=True, exist_ok=True)
 
+        # The export copies the source tokenizer files unchanged; the pad token set above is only
+        # for generation here.
         with torch.inference_mode():
             export_hf_checkpoint(model, export_dir=export_dir)
 
-        tokenizer.save_pretrained(export_dir)
         print(f"Model exported successfully to: {export_dir}")
 
 

examples/llm_sparsity/weight_sparsity/export_hf_ckpt.py

@@ -23,7 +23,7 @@
 
 import modelopt.torch.opt as mto
 import modelopt.torch.sparsity as mts
-from modelopt.torch.export import export_hf_checkpoint
+from modelopt.torch.export import ensure_local_checkpoint, export_hf_checkpoint
 
 DEFAULT_PAD_TOKEN = "[PAD]"
 
@@ -91,9 +91,15 @@ def main(args):
     random.seed(1234)
     np.random.seed(1234)
 
-    model = get_model(args.model_name_or_path, args.dtype, trust_remote_code=args.trust_remote_code)
+    # --model_name_or_path stays as given. Later steps read the local copy of the whole checkpoint
+    # in local_checkpoint_path, since exporters read the source's files from local disk only;
+    # hub_model_id is the Hub ID, or None for a local path.
+    args.hub_model_id, args.local_checkpoint_path = ensure_local_checkpoint(args.model_name_or_path)
+    model = get_model(
+        args.local_checkpoint_path, args.dtype, trust_remote_code=args.trust_remote_code
+    )
     tokenizer = get_tokenizer(
-        args.model_name_or_path, args.model_max_length, trust_remote_code=args.trust_remote_code
+        args.local_checkpoint_path, args.model_max_length, trust_remote_code=args.trust_remote_code
     )
 
     if tokenizer.pad_token is None:
@@ -114,6 +120,8 @@ def main(args):
     with torch.inference_mode():
         model = mts.export(model)
         export_hf_checkpoint(model, export_dir=args.output_dir)
+        # After the export, so it replaces the source tokenizer files the export copied: a pad token
+        # added above changes the vocabulary.
         tokenizer.save_pretrained(args.output_dir)