From 33e171ee71a5db8b012b9c220654df4247c85a03 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 30 Sep 2026 21:36:27 +0000 Subject: [PATCH 1/2] llm_sparsity: export from a local checkpoint The export now carries the source's non-model files only from local disk, so hf_sa.py and export_hf_ckpt.py make the model local first with ensure_local_checkpoint, keeping their argument as given and reading from hf_model_path. hf_sa.py no longer saves its tokenizer into the export: it sets pad_token to eos only for generation, so the source's files are the right ones. export_hf_ckpt.py keeps its save, after the export, because it can add a pad token and resize the embeddings. Signed-off-by: Shengliang Xu --- examples/llm_sparsity/attention_sparsity/hf_sa.py | 15 ++++++++++----- .../weight_sparsity/export_hf_ckpt.py | 12 +++++++++--- 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/examples/llm_sparsity/attention_sparsity/hf_sa.py b/examples/llm_sparsity/attention_sparsity/hf_sa.py index 1eacc6f18b3..f5bba9cdcf7 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 + # hf_model_path, since exporters read the source's files from local disk only; hf_model_name is + # the Hub ID, or None for a local path. + args.hf_model_name, args.hf_model_path = ensure_local_checkpoint(args.pyt_ckpt_path) + print(f"Loading model: {args.hf_model_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.hf_model_path, attn_implementation="eager", dtype="auto", device_map="auto" ) - tokenizer = AutoTokenizer.from_pretrained(args.pyt_ckpt_path) + tokenizer = AutoTokenizer.from_pretrained(args.hf_model_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..1babc63bd90 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,13 @@ 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 hf_model_path, since exporters read the source's files from local disk only; hf_model_name + # is the Hub ID, or None for a local path. + args.hf_model_name, args.hf_model_path = ensure_local_checkpoint(args.model_name_or_path) + model = get_model(args.hf_model_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.hf_model_path, args.model_max_length, trust_remote_code=args.trust_remote_code ) if tokenizer.pad_token is None: @@ -114,6 +118,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) From 737c34bf547acd236a8af511a350b929a8b820db Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 1 Oct 2026 01:00:44 +0000 Subject: [PATCH 2/2] llm_sparsity: store the local copy as local_checkpoint_path, the ID as hub_model_id As in hf_ptq: name the local checkpoint directory and the Hub identifier for what they are. Signed-off-by: Shengliang Xu --- examples/llm_sparsity/attention_sparsity/hf_sa.py | 12 ++++++------ .../llm_sparsity/weight_sparsity/export_hf_ckpt.py | 12 +++++++----- 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/examples/llm_sparsity/attention_sparsity/hf_sa.py b/examples/llm_sparsity/attention_sparsity/hf_sa.py index f5bba9cdcf7..b94238b8cb1 100644 --- a/examples/llm_sparsity/attention_sparsity/hf_sa.py +++ b/examples/llm_sparsity/attention_sparsity/hf_sa.py @@ -147,17 +147,17 @@ def main(args): launch_memory_monitor() # --pyt_ckpt_path stays as given. Later steps read the local copy of the whole checkpoint in - # hf_model_path, since exporters read the source's files from local disk only; hf_model_name is - # the Hub ID, or None for a local path. - args.hf_model_name, args.hf_model_path = ensure_local_checkpoint(args.pyt_ckpt_path) - print(f"Loading model: {args.hf_model_path}") + # 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.hf_model_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.hf_model_path) + tokenizer = AutoTokenizer.from_pretrained(args.local_checkpoint_path) # Set pad token if not set if tokenizer.pad_token is None: diff --git a/examples/llm_sparsity/weight_sparsity/export_hf_ckpt.py b/examples/llm_sparsity/weight_sparsity/export_hf_ckpt.py index 1babc63bd90..e494173e315 100644 --- a/examples/llm_sparsity/weight_sparsity/export_hf_ckpt.py +++ b/examples/llm_sparsity/weight_sparsity/export_hf_ckpt.py @@ -92,12 +92,14 @@ def main(args): np.random.seed(1234) # --model_name_or_path stays as given. Later steps read the local copy of the whole checkpoint - # in hf_model_path, since exporters read the source's files from local disk only; hf_model_name - # is the Hub ID, or None for a local path. - args.hf_model_name, args.hf_model_path = ensure_local_checkpoint(args.model_name_or_path) - model = get_model(args.hf_model_path, args.dtype, trust_remote_code=args.trust_remote_code) + # 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.hf_model_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: