diff --git a/examples/llm_sparsity/attention_sparsity/hf_sa.py b/examples/llm_sparsity/attention_sparsity/hf_sa.py index 1eacc6f18b3..b94238b8cb1 100644 --- a/examples/llm_sparsity/attention_sparsity/hf_sa.py +++ b/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}") diff --git a/examples/llm_sparsity/weight_sparsity/export_hf_ckpt.py b/examples/llm_sparsity/weight_sparsity/export_hf_ckpt.py index 7ef22c52eca..e494173e315 100644 --- a/examples/llm_sparsity/weight_sparsity/export_hf_ckpt.py +++ b/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)