From 90c448340622f0764812cd12a114a0d0c9666066 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 30 Sep 2026 21:36:06 +0000
Subject: [PATCH 01/11] Megatron export: copy non-model files from local disk
 only

The Megatron unified HF export copied every top-level non-safetensors file
of a local source -- *.bin weights included -- but only *.py of a Hub ID,
then regenerated the generation config, tokenizer and image processor.
Now it writes its config and copies the source's non-model files with
copy_non_model_files, from a local source only; a Hub ID gets a warning,
as in the HF exporters. That is a breaking change for callers passing a
Hub ID, who no longer get those files.

The megatron_bridge quantize, export, prune and distillation-export
scripts therefore make the HF model local first with
ensure_local_checkpoint, keeping their argument as given and reading from
hf_model_path. The tokenizer path recorded in saved Megatron checkpoints
stays the Hub ID, which remains valid on other hosts. prune_minitron and
the distillation VLM export -- which every rank reaches -- copy the
non-model files from the main rank only, and the distillation LLM export
now copies them too.

copy_non_safetensor_files_from_ckpt lost its last internal user and is
deprecated in favour of copy_non_model_files.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 CHANGELOG.rst                                 |  4 +-
 .../export_distilled_megatron_to_hf.py        | 21 +++++---
 .../export_quantized_megatron_to_hf.py        | 14 +++--
 examples/megatron_bridge/prune_minitron.py    | 20 +++++---
 examples/megatron_bridge/quantize.py          | 15 ++++--
 .../torch/export/unified_export_megatron.py   | 51 ++++---------------
 .../utils/plugins/hf_checkpoint_utils.py      |  9 ++++
 .../torch/utils/test_hf_checkpoint_utils.py   | 16 +++---
 8 files changed, 80 insertions(+), 70 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index f20a33f1c17..133d203e583 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -108,7 +108,7 @@ Changelog
 - Add ``--modelopt-*`` options to the ``examples/vllm_serve`` launcher for fakequant calibration and checkpoint reload. Point it at an exported fakequant model directory to reload automatically: Hugging Face exports use ``vllm_fq_modelopt_state.pth``, while Megatron-Core exports now include ``quant_recipe.yaml`` alongside ``quantizer_state.pth`` for reload.
 - Add a ``kd_loss_alpha`` option to HuggingFace ``KDTrainer`` to blend cross-entropy and knowledge distillation losses during training.
 - A tracked ``examples/hf_ptq/hf_ptq.py`` run now writes ``.experiment.json`` into ``--export_path`` and uploads the same file with the run, so a checkpoint on disk names the experiment and MLflow run id that produced it. The pointer is written only once the export completes, and an export that is not tracked removes one it would otherwise inherit from a reused ``--export_path`` or from a quantized source checkpoint.
-- ``export_hf_checkpoint`` and the vLLM fake-quant export now copy the source checkpoint's non-model files (tokenizer, processor, remote code, chat templates, ...) into the export unchanged. They read the source from local disk only: get a local copy of a Hugging Face Hub model with ``modelopt.torch.export.ensure_local_checkpoint`` and load it from there.
+- ``export_hf_checkpoint``, the vLLM fake-quant export and the Megatron unified HF export now copy the source checkpoint's non-model files (tokenizer, processor, remote code, chat templates, ...) into the export unchanged. They read the source from local disk only: get a local copy of a Hugging Face Hub model with ``modelopt.torch.export.ensure_local_checkpoint`` and load it from there.
 
 - Add Parallel Decoding Distillation (PDD) to ``modelopt.torch.fastgen`` with Qwen-Image training, distributed-checkpoint export, and PDD-2/4/8 inference. AutoModel remains an unmodified pinned runtime dependency.
 
@@ -120,6 +120,7 @@ Changelog
 - ``LogitsAndIntermediatesLossBalancer`` (Megatron distillation plugin) no longer rescales the distillation loss to the magnitude of the LM loss when the LM loss is included. The total is now the fixed convex combination ``(1 - alpha) * lm_loss + alpha * kd_loss`` with ``DistillationConfig.kd_loss_alpha`` (in [0, 1]), matching Megatron-LM's offline cached-logits KD. The default ``kd_loss_alpha=1.0`` skips the LM loss, as before. ``DistillationConfig.skip_lm_loss`` and ``kd_loss_scale`` are removed and raise ``ValueError`` if passed; set ``kd_loss_alpha`` instead. In ``examples/megatron_bridge/distill.py``, ``--kd_loss_alpha`` replaces ``--no_skip_lm_loss`` and ``--kd_loss_scale``.
 - The ``examples/vllm_serve`` fakequant launcher no longer supports vLLM 0.9.0. Upgrade to a version listed as tested in the example README.
 - The Megatron-Core DeepSeek-V4 indexer (``CSAIndexer``) is now a quantization module and persists its quantizer state in the checkpoint as ``indexer._extra_state``. A DeepSeek-V4 model quantized with an earlier release resumes from its ``torch_dist`` checkpoint only with a non-strict load (``--dist-ckpt-strictness log_unexpected`` in Megatron-LM) until it is saved again.
+- ``export_mcore_gpt_to_hf`` given a Hugging Face Hub model ID no longer writes the tokenizer, processor, generation config or remote-code files. Pass a local directory instead, e.g. from ``modelopt.torch.export.ensure_local_checkpoint``.
 - ``examples/hf_ptq`` no longer detects MTP layers by name. Weights the loader could not place -- an MTP head, an auxiliary tower -- are identified from Transformers' own accounting: the model is loaded with ``from_pretrained(..., output_loading_info=True)`` and the reported ``unexpected_keys`` (present in the checkpoint, not in the model's architecture) are recorded on the model and carried into the export unchanged. Everything the loader *did* place goes through the normal export path. This removes ``load_mtp_weights``, ``mtp_layer_prefixes_from_checkpoint`` and their support matrix of MTP storage conventions, along with ``_add_mtp_exclusions`` and the pre-quantization ``enable: False`` entries ``hf_ptq`` appended to the recipe's ``quant_cfg``. Two consequences: MTP layers now follow the recipe like any other module instead of being force-excluded by the script -- matching ``examples/megatron_bridge``, which has no MTP-specific code at all -- and ``quantization_config.ignore`` can no longer claim a layer is unquantized that the export in fact quantized. Recipes importing ``configs/ptq/units/default_disabled_quantizers`` still disable ``mtp.*``, so their behaviour is unchanged; a recipe omitting that unit will now quantize an MTP the model actually built.
 - ``examples/hf_ptq --vllm_fakequant_export`` now raises ``NotImplementedError`` when the checkpoint holds weights the model has no parameter for and a shard actually provides them (an MTP head, an auxiliary tower). The fake-quant exporter writes only model-backed state, so it would otherwise drop those weights silently -- and a fake-quant checkpoint is evaluated, where a missing head changes the score rather than failing loudly. Use the unified HF export, which carries them through. Buffers Transformers recomputes are not weights to lose: ``*.inv_freq`` is skipped even when a shard provides it, since older Llama/Mistral-lineage conversions do list it in the index and refusing an export over it would reject checkpoints that export correctly today. The check runs immediately after the model loads, not at export time, so an incompatible run fails before calibration rather than after it.
 - The ``modelopt.onnx.quantization.ort_patching`` module has been removed with no
@@ -173,6 +174,7 @@ Changelog
 - The ``evaluation`` agent skill no longer supports GDPVal: its task recipe, example config and Apptainer SIF helper are removed, and the shared NeMo Gym machinery it carried now lives at ``references/gym.md`` with the launcher wrapper renamed ``scripts/nel-gym.sh``. GDPVal is an AA-suite member, so an "AA" request now generates the ``aa/`` tasks only -- report per-task scores rather than an aggregate compared against a published AA Index.
 - DBRX model support is removed: ModelOpt no longer quantizes or exports DBRX (``DbrxForCausalLM``) checkpoints. Use ModelOpt 0.47 to quantize DBRX, or write your own plugin as described in the customized-model quantization guide.
 - Deprecate ``metadata.recipe_type`` in recipe YAML, for both single-file recipes and a directory recipe's ``metadata.yml``. A recipe now says what kind it is with a ``# modelopt-schema:`` comment naming its schema class, or by delegating to a recipe that does (single-file recipes only -- a directory recipe has no body to delegate through); ``recipe_type`` is still read and still honoured, so a recipe outside this repo keeps working unchanged, but new recipes should leave it out and every recipe shipped here has been converted. Where both are present they must agree, and so must a recipe and the recipe it delegates to -- a disagreement is an error rather than a silent preference.
+- ``modelopt.torch.export.copy_non_safetensor_files_from_ckpt`` is deprecated; use ``copy_non_model_files``.
 
 **Bug Fixes**
 
diff --git a/examples/megatron_bridge/export_distilled_megatron_to_hf.py b/examples/megatron_bridge/export_distilled_megatron_to_hf.py
index 83c2b595389..f32f99f9037 100644
--- a/examples/megatron_bridge/export_distilled_megatron_to_hf.py
+++ b/examples/megatron_bridge/export_distilled_megatron_to_hf.py
@@ -72,7 +72,7 @@
 from transformers import AutoConfig
 
 import modelopt.torch.utils.distributed as dist
-from modelopt.torch.export import copy_hf_ckpt_remote_code
+from modelopt.torch.export import copy_non_model_files, ensure_local_checkpoint
 from modelopt.torch.utils import print_args, print_rank_0
 from modelopt.torch.utils.mlflow import Tool, masked_args
 from modelopt.torch.utils.plugins.mbridge import (
@@ -168,6 +168,7 @@ def export_llm_to_hf(
     AutoConfig.from_pretrained(
         student_hf_path, trust_remote_code=trust_remote_code
     ).save_pretrained(hf_export_path)
+    copy_non_model_files(student_hf_path, hf_export_path)
 
 
 def save_vlm_to_hf(
@@ -194,7 +195,9 @@ def save_vlm_to_hf(
     )
     export_bridge.hf_pretrained.save_artifacts(hf_export_path)
     export_bridge.save_hf_weights([full_model], hf_export_path)
-    copy_hf_ckpt_remote_code(student_hf_path, hf_export_path)
+    # Every rank runs this function (save_hf_weights is collective); one writer is enough.
+    if dist.is_master():
+        copy_non_model_files(student_hf_path, hf_export_path)
 
 
 def get_args() -> argparse.Namespace:
@@ -268,6 +271,12 @@ def get_args() -> argparse.Namespace:
 
 
 def main(args: argparse.Namespace):
+    # --student_hf_path stays as given. Later steps read the local copy of the whole checkpoint in
+    # student_hf_model_path, since exporters read the source's files from local disk only;
+    # student_hf_model_name is the Hub ID, or None for a local path.
+    args.student_hf_model_name, args.student_hf_model_path = ensure_local_checkpoint(
+        args.student_hf_path
+    )
     checkpoint_export_paths: list[tuple[Path, Path]] = _get_checkpoint_export_paths(args)
     # This path drops quantization, so a QAD checkpoint would export silently unquantized.
     # ``has_modelopt_state`` ignores ``kd_loss``, so plain distillation still passes.
@@ -277,14 +286,14 @@ def main(args: argparse.Namespace):
             f"{quantized[0]} is quantized; this script exports full precision only and would drop "
             "the quantizers. Use export_quantized_megatron_to_hf.py instead."
         )
-    is_vlm = is_vlm_config(args.student_hf_path, trust_remote_code=args.trust_remote_code)
+    is_vlm = is_vlm_config(args.student_hf_model_path, trust_remote_code=args.trust_remote_code)
 
     if is_vlm:
         # Build the full VLM (vision tower / projector + original LM from HF), then overwrite the LM
         # with the distilled checkpoint weights, then export the assembled VLM.
         print_rank_0("Reassembling distilled VLM and exporting to HF format")
         _bridge, _provider, _model, full_model, _tokenizer = load_mbridge_model_from_hf(
-            hf_model_name_or_path=args.student_hf_path,
+            hf_model_name_or_path=args.student_hf_model_path,
             trust_remote_code=args.trust_remote_code,
             # Mirrors distill.py's unquantized branch
             moe_grouped_gemm=not args.no_moe_grouped_gemm,
@@ -309,7 +318,7 @@ def main(args: argparse.Namespace):
             save_vlm_to_hf(
                 full_model,
                 str(hf_export_path),
-                args.student_hf_path,
+                args.student_hf_model_path,
                 trust_remote_code=args.trust_remote_code,
             )
             record_exported_checkpoint(args, hf_export_path, dist.is_master())
@@ -327,7 +336,7 @@ def main(args: argparse.Namespace):
                 export_llm_to_hf(
                     megatron_path=str(megatron_path),
                     hf_export_path=str(hf_export_path),
-                    student_hf_path=args.student_hf_path,
+                    student_hf_path=args.student_hf_model_path,
                     template_hf=args.student_hf_model,
                     trust_remote_code=args.trust_remote_code,
                 )
diff --git a/examples/megatron_bridge/export_quantized_megatron_to_hf.py b/examples/megatron_bridge/export_quantized_megatron_to_hf.py
index a035dda679d..badc682eb31 100644
--- a/examples/megatron_bridge/export_quantized_megatron_to_hf.py
+++ b/examples/megatron_bridge/export_quantized_megatron_to_hf.py
@@ -52,7 +52,7 @@
 
 import modelopt.torch.quantization as mtq
 import modelopt.torch.utils.distributed as dist
-from modelopt.torch.export import export_mcore_gpt_to_hf
+from modelopt.torch.export import ensure_local_checkpoint, export_mcore_gpt_to_hf
 from modelopt.torch.export.plugins.vllm_fakequant_megatron import export_mcore_gpt_to_hf_vllm_fq
 from modelopt.torch.utils import print_args, print_rank_0
 from modelopt.torch.utils.mlflow import Tool, masked_args
@@ -158,16 +158,20 @@ def get_args() -> argparse.Namespace:
 
 
 def main(args: argparse.Namespace):
+    # --hf_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.hf_model_name_or_path)
     trust_remote_code = is_safe_repo(
-        trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_name_or_path
+        trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_path
     )
 
     # Build the model structure from HF
     _bridge, _provider, model, _unwrapped_model, _tokenizer = load_mbridge_model_from_hf(
-        hf_model_name_or_path=args.hf_model_name_or_path,
+        hf_model_name_or_path=args.hf_model_path,
         trust_remote_code=trust_remote_code,
         moe_grouped_gemm=use_moe_grouped_gemm(
-            args.hf_model_name_or_path,
+            args.hf_model_path,
             trust_remote_code=trust_remote_code,
             force_sequential=args.no_moe_grouped_gemm,
         ),
@@ -214,7 +218,7 @@ def main(args: argparse.Namespace):
     )
     export_fn(
         unwrapped_model,
-        args.hf_model_name_or_path,
+        args.hf_model_path,
         export_extra_modules=export_extra_modules,
         dtype=torch.bfloat16,
         export_dir=args.export_unified_hf_path,
diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py
index 0abaf810f3e..28bdbed5dfc 100644
--- a/examples/megatron_bridge/prune_minitron.py
+++ b/examples/megatron_bridge/prune_minitron.py
@@ -58,7 +58,7 @@
 import modelopt.torch.opt as mto
 import modelopt.torch.prune as mtp
 import modelopt.torch.utils.distributed as dist
-from modelopt.torch.export import copy_hf_ckpt_remote_code
+from modelopt.torch.export import copy_non_model_files, ensure_local_checkpoint
 from modelopt.torch.nas.plugins.megatron_model_stats import parse_main_layer_chars
 from modelopt.torch.utils import (
     get_supported_datasets,
@@ -438,8 +438,13 @@ def main(args: argparse.Namespace):
         warn_rank_0(f"\nPruned model already exists at {args.output_hf_path}. Exiting...")
         return
 
+    # --hf_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.hf_model_name_or_path)
+
     bridge, provider, model, unwrapped_model, tokenizer = load_mbridge_model_from_hf(
-        hf_model_name_or_path=args.hf_model_name_or_path,
+        hf_model_name_or_path=args.hf_model_path,
         trust_remote_code=args.trust_remote_code,
         provider_overrides={
             "tensor_model_parallel_size": 1,  # Tensor parallelism is not supported
@@ -510,7 +515,7 @@ def main(args: argparse.Namespace):
     # the full VLM forward over image-text pairs.
     if use_image_calib:
         processor = AutoProcessor.from_pretrained(
-            args.hf_model_name_or_path, trust_remote_code=args.trust_remote_code
+            args.hf_model_path, trust_remote_code=args.trust_remote_code
         )
         forward_loop = get_megatron_vlm_calibration_forward_loop(
             unwrapped_model,  # full VLM (vision encoder + projector + language model)
@@ -643,7 +648,8 @@ def score_func(m):
         bridge.save_megatron_model(
             model,
             args.output_megatron_path,
-            hf_tokenizer_path=args.hf_model_name_or_path,
+            # Recorded in the checkpoint's config: the Hub ID stays valid on other hosts.
+            hf_tokenizer_path=args.hf_model_name or args.hf_model_path,
             hf_tokenizer_kwargs=tokenizer_kwargs,
         )
         args.checkpoint_exported = True
@@ -752,7 +758,7 @@ def score_func(m):
             # from_hf_config can't infer it since AutoConfig consumes the kwarg.
             pruned_bridge.trust_remote_code = args.trust_remote_code
             pruned_bridge.save_hf_pretrained(
-                model, args.output_hf_path, source_path=args.hf_model_name_or_path
+                model, args.output_hf_path, source_path=args.hf_model_path
             )
         else:
             dummy_model_cls = AutoModelForImageTextToText if is_vlm else AutoModelForCausalLM
@@ -764,7 +770,9 @@ def score_func(m):
             )
             pruned_bridge.save_hf_weights(model, args.output_hf_path)
 
-        copy_hf_ckpt_remote_code(args.hf_model_name_or_path, args.output_hf_path)
+        # Every rank gets here (save_hf_weights is collective); one writer is enough.
+        if dist.is_master():
+            copy_non_model_files(args.hf_model_path, args.output_hf_path)
         args.checkpoint_exported = True
         print_rank_0(f"Saved pruned model to {args.output_hf_path} in HF checkpoint format")
 
diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py
index 44bcee9bf00..a30357dc567 100644
--- a/examples/megatron_bridge/quantize.py
+++ b/examples/megatron_bridge/quantize.py
@@ -76,6 +76,7 @@
     QUANT_CFG_CHOICES,
     RecipeSupersededAction,
 )
+from modelopt.torch.export import ensure_local_checkpoint
 from modelopt.torch.utils import print_args, print_rank_0, warn_rank_0
 from modelopt.torch.utils.dataset_utils import get_supported_datasets
 from modelopt.torch.utils.mlflow import Tool, masked_args, resolved_recipe_texts
@@ -331,18 +332,21 @@ def get_quant_config(args: argparse.Namespace) -> dict:
 
 
 def main(args: argparse.Namespace):
+    # --hf_model_name_or_path stays as given. Later steps read the local copy of the whole
+    # checkpoint in hf_model_path; 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.hf_model_name_or_path)
     trust_remote_code = is_safe_repo(
-        trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_name_or_path
+        trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_path
     )
 
     moe_grouped_gemm = use_moe_grouped_gemm(
-        args.hf_model_name_or_path,
+        args.hf_model_path,
         trust_remote_code=trust_remote_code,
         force_sequential=args.no_moe_grouped_gemm,
     )
 
     bridge, _provider, model, unwrapped_model, tokenizer = load_mbridge_model_from_hf(
-        hf_model_name_or_path=args.hf_model_name_or_path,
+        hf_model_name_or_path=args.hf_model_path,
         trust_remote_code=trust_remote_code,
         moe_grouped_gemm=moe_grouped_gemm,
         provider_overrides={
@@ -447,7 +451,7 @@ def forward_loop(_model=None):
         # VLMs: drive the full VLM forward on image-text pairs so the language model's quantizers
         # see vision-conditioned activations (we still quantize the LM only).
         processor = AutoProcessor.from_pretrained(
-            args.hf_model_name_or_path, trust_remote_code=trust_remote_code
+            args.hf_model_path, trust_remote_code=trust_remote_code
         )
         forward_loop = get_megatron_vlm_calibration_forward_loop(
             unwrapped_model,  # full VLM (vision encoder + projector + language model)
@@ -491,7 +495,8 @@ def forward_loop(_model=None):
     bridge.save_megatron_model(
         model,
         args.export_megatron_path,
-        hf_tokenizer_path=args.hf_model_name_or_path,
+        # Recorded in the checkpoint's config: the Hub ID stays valid on other hosts.
+        hf_tokenizer_path=args.hf_model_name or args.hf_model_path,
         hf_tokenizer_kwargs={"trust_remote_code": trust_remote_code},
     )
     args.checkpoint_exported = True
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index 6f3f128d1fc..7e227a9ddca 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -39,8 +39,7 @@
 from modelopt.torch.utils import import_plugin, warn_rank_0
 from modelopt.torch.utils.distributed import gather_state_dicts
 from modelopt.torch.utils.plugins.hf_checkpoint_utils import (
-    copy_hf_ckpt_remote_code,
-    copy_non_safetensor_files_from_ckpt,
+    copy_non_model_files,
     load_multimodal_components,
 )
 
@@ -83,7 +82,6 @@
 
 with import_plugin("transformers", verbose=False):
     import transformers
-    from transformers import AutoProcessor
 
 
 has_mcore = False
@@ -391,47 +389,20 @@ def save_pretrained(
 
         if is_last_stage_main_rank:
             if is_writer_rank:
+                self._hf_config.save_pretrained(save_directory)
+                # Everything else the source ships -- tokenizer, processor, generation config,
+                # remote code -- is carried over verbatim rather than regenerated, from local disk
+                # only: a Hub ID's cache holds just what loading needed.
                 if self._hf_pretrained_model_name is not None:
                     if os.path.isdir(self._hf_pretrained_model_name):
-                        copy_non_safetensor_files_from_ckpt(
-                            self._hf_pretrained_model_name, save_directory
-                        )
+                        copy_non_model_files(self._hf_pretrained_model_name, save_directory)
                     else:
-                        copy_hf_ckpt_remote_code(self._hf_pretrained_model_name, save_directory)
-                self._hf_config.save_pretrained(save_directory)
-                try:
-                    generation_config = transformers.GenerationConfig.from_pretrained(
-                        self._hf_pretrained_model_name,
-                        trust_remote_code=self.trust_remote_code,
-                    )
-                    # Pass it through unvalidated: save_pretrained rejects some shipped configs
-                    # (e.g. GLM-5.3-Flash sets top_p without do_sample) on newer transformers.
-                    generation_config.to_json_file(
-                        os.path.join(save_directory, "generation_config.json")
-                    )
-                except OSError:
-                    pass
-                # Hub-ID / None source: fetch tokenizer files via AutoTokenizer.
-                if self._hf_pretrained_model_name is None or not os.path.isdir(
-                    self._hf_pretrained_model_name
-                ):
-                    try:
-                        tokenizer = transformers.AutoTokenizer.from_pretrained(
-                            self._hf_pretrained_model_name,
-                            trust_remote_code=self.trust_remote_code,
+                        warn_rank_0(
+                            f"{self._hf_pretrained_model_name!r} is not a local directory, so its "
+                            "non-model files (tokenizer, processor, remote code, ...) were not "
+                            "copied into the export. Get a local copy first with "
+                            "modelopt.torch.export.ensure_local_checkpoint and pass that directory."
                         )
-                        tokenizer.save_pretrained(save_directory)
-                    except (OSError, TypeError, ValueError, ImportError):
-                        pass
-                try:
-                    # Load and save preprocessor config from the original model
-                    processor = AutoProcessor.from_pretrained(
-                        self._hf_pretrained_model_name, trust_remote_code=self.trust_remote_code
-                    )
-                    if hasattr(processor, "image_processor"):
-                        processor.image_processor.save_pretrained(save_directory)
-                except (OSError, ValueError, ImportError):
-                    pass
 
             # The live MTP export runs EP collectives, so every last-stage main rank joins it; the
             # collective-free copy from the source checkpoint only runs on the writer.
diff --git a/modelopt/torch/utils/plugins/hf_checkpoint_utils.py b/modelopt/torch/utils/plugins/hf_checkpoint_utils.py
index dc8138d4286..625d5141ea2 100644
--- a/modelopt/torch/utils/plugins/hf_checkpoint_utils.py
+++ b/modelopt/torch/utils/plugins/hf_checkpoint_utils.py
@@ -736,7 +736,16 @@ def copy_non_safetensor_files_from_ckpt(
 
     Returns:
         File names copied into ``dst``.
+
+    .. deprecated::
+        Use :func:`copy_non_model_files`, which also skips non-safetensors weights, keeps files the
+        export already wrote, and copies subdirectories.
     """
+    warnings.warn(
+        "copy_non_safetensor_files_from_ckpt is deprecated; use copy_non_model_files.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
     if not os.path.isdir(src):
         raise ValueError(f"Invalid source path: {src}. It should be a directory.")
     exclude_files = set(exclude_files or ())
diff --git a/tests/unit/torch/utils/test_hf_checkpoint_utils.py b/tests/unit/torch/utils/test_hf_checkpoint_utils.py
index 093e8e1cf66..23f23b0e83c 100644
--- a/tests/unit/torch/utils/test_hf_checkpoint_utils.py
+++ b/tests/unit/torch/utils/test_hf_checkpoint_utils.py
@@ -56,18 +56,20 @@ def test_copy_non_safetensor_files_from_ckpt_supports_additional_exclusions(tmp_
     (src_dir / "reasoning_parser.py").write_text("parser")
 
     default_dst = tmp_path / "default"
-    copy_non_safetensor_files_from_ckpt(src_dir, default_dst)
+    with pytest.deprecated_call():
+        copy_non_safetensor_files_from_ckpt(src_dir, default_dst)
     assert not (default_dst / "model.safetensors").exists()
     assert not (default_dst / "model.safetensors.index.json").exists()
     assert (default_dst / "pytorch_model.bin").exists()
     assert (default_dst / "stats.npy").exists()
 
     filtered_dst = tmp_path / "filtered"
-    copy_non_safetensor_files_from_ckpt(
-        src_dir,
-        filtered_dst,
-        exclude_patterns=("*.bin", "*.npy"),
-    )
+    with pytest.deprecated_call():
+        copy_non_safetensor_files_from_ckpt(
+            src_dir,
+            filtered_dst,
+            exclude_patterns=("*.bin", "*.npy"),
+        )
     assert (filtered_dst / "reasoning_parser.py").exists()
     assert not (filtered_dst / "pytorch_model.bin").exists()
     assert not (filtered_dst / "stats.npy").exists()
@@ -87,7 +89,7 @@ def copy2(source, *args, **kwargs):
         return original_copy2(source, *args, **kwargs)
 
     monkeypatch.setattr(hf_checkpoint_utils.shutil, "copy2", copy2)
-    with pytest.warns(UserWarning, match="bad.py"):
+    with pytest.deprecated_call(), pytest.warns(UserWarning, match="bad.py"):
         copied_files = copy_non_safetensor_files_from_ckpt(src_dir, tmp_path / "dst")
 
     assert copied_files == ["good.py"]

From 7e7f101acc4e528dc8e3aa93ebc123b40424a447 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 1 Oct 2026 01:00:13 +0000
Subject: [PATCH 02/11] megatron_bridge: 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, with a student_ prefix in the distillation export.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../export_distilled_megatron_to_hf.py         | 16 +++++++++-------
 .../export_quantized_megatron_to_hf.py         | 16 +++++++++-------
 examples/megatron_bridge/prune_minitron.py     | 18 ++++++++++--------
 examples/megatron_bridge/quantize.py           | 16 +++++++++-------
 4 files changed, 37 insertions(+), 29 deletions(-)

diff --git a/examples/megatron_bridge/export_distilled_megatron_to_hf.py b/examples/megatron_bridge/export_distilled_megatron_to_hf.py
index f32f99f9037..26c481c6de5 100644
--- a/examples/megatron_bridge/export_distilled_megatron_to_hf.py
+++ b/examples/megatron_bridge/export_distilled_megatron_to_hf.py
@@ -272,9 +272,9 @@ def get_args() -> argparse.Namespace:
 
 def main(args: argparse.Namespace):
     # --student_hf_path stays as given. Later steps read the local copy of the whole checkpoint in
-    # student_hf_model_path, since exporters read the source's files from local disk only;
-    # student_hf_model_name is the Hub ID, or None for a local path.
-    args.student_hf_model_name, args.student_hf_model_path = ensure_local_checkpoint(
+    # student_local_checkpoint_path, since exporters read the source's files from local disk only;
+    # student_hub_model_id is the Hub ID, or None for a local path.
+    args.student_hub_model_id, args.student_local_checkpoint_path = ensure_local_checkpoint(
         args.student_hf_path
     )
     checkpoint_export_paths: list[tuple[Path, Path]] = _get_checkpoint_export_paths(args)
@@ -286,14 +286,16 @@ def main(args: argparse.Namespace):
             f"{quantized[0]} is quantized; this script exports full precision only and would drop "
             "the quantizers. Use export_quantized_megatron_to_hf.py instead."
         )
-    is_vlm = is_vlm_config(args.student_hf_model_path, trust_remote_code=args.trust_remote_code)
+    is_vlm = is_vlm_config(
+        args.student_local_checkpoint_path, trust_remote_code=args.trust_remote_code
+    )
 
     if is_vlm:
         # Build the full VLM (vision tower / projector + original LM from HF), then overwrite the LM
         # with the distilled checkpoint weights, then export the assembled VLM.
         print_rank_0("Reassembling distilled VLM and exporting to HF format")
         _bridge, _provider, _model, full_model, _tokenizer = load_mbridge_model_from_hf(
-            hf_model_name_or_path=args.student_hf_model_path,
+            hf_model_name_or_path=args.student_local_checkpoint_path,
             trust_remote_code=args.trust_remote_code,
             # Mirrors distill.py's unquantized branch
             moe_grouped_gemm=not args.no_moe_grouped_gemm,
@@ -318,7 +320,7 @@ def main(args: argparse.Namespace):
             save_vlm_to_hf(
                 full_model,
                 str(hf_export_path),
-                args.student_hf_model_path,
+                args.student_local_checkpoint_path,
                 trust_remote_code=args.trust_remote_code,
             )
             record_exported_checkpoint(args, hf_export_path, dist.is_master())
@@ -336,7 +338,7 @@ def main(args: argparse.Namespace):
                 export_llm_to_hf(
                     megatron_path=str(megatron_path),
                     hf_export_path=str(hf_export_path),
-                    student_hf_path=args.student_hf_model_path,
+                    student_hf_path=args.student_local_checkpoint_path,
                     template_hf=args.student_hf_model,
                     trust_remote_code=args.trust_remote_code,
                 )
diff --git a/examples/megatron_bridge/export_quantized_megatron_to_hf.py b/examples/megatron_bridge/export_quantized_megatron_to_hf.py
index badc682eb31..08a58102ce7 100644
--- a/examples/megatron_bridge/export_quantized_megatron_to_hf.py
+++ b/examples/megatron_bridge/export_quantized_megatron_to_hf.py
@@ -159,19 +159,21 @@ def get_args() -> argparse.Namespace:
 
 def main(args: argparse.Namespace):
     # --hf_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.hf_model_name_or_path)
+    # 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.hf_model_name_or_path
+    )
     trust_remote_code = is_safe_repo(
-        trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_path
+        trust_remote_code=args.trust_remote_code, hf_path=args.local_checkpoint_path
     )
 
     # Build the model structure from HF
     _bridge, _provider, model, _unwrapped_model, _tokenizer = load_mbridge_model_from_hf(
-        hf_model_name_or_path=args.hf_model_path,
+        hf_model_name_or_path=args.local_checkpoint_path,
         trust_remote_code=trust_remote_code,
         moe_grouped_gemm=use_moe_grouped_gemm(
-            args.hf_model_path,
+            args.local_checkpoint_path,
             trust_remote_code=trust_remote_code,
             force_sequential=args.no_moe_grouped_gemm,
         ),
@@ -218,7 +220,7 @@ def main(args: argparse.Namespace):
     )
     export_fn(
         unwrapped_model,
-        args.hf_model_path,
+        args.local_checkpoint_path,
         export_extra_modules=export_extra_modules,
         dtype=torch.bfloat16,
         export_dir=args.export_unified_hf_path,
diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py
index 28bdbed5dfc..28e86c8baa1 100644
--- a/examples/megatron_bridge/prune_minitron.py
+++ b/examples/megatron_bridge/prune_minitron.py
@@ -439,12 +439,14 @@ def main(args: argparse.Namespace):
         return
 
     # --hf_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.hf_model_name_or_path)
+    # 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.hf_model_name_or_path
+    )
 
     bridge, provider, model, unwrapped_model, tokenizer = load_mbridge_model_from_hf(
-        hf_model_name_or_path=args.hf_model_path,
+        hf_model_name_or_path=args.local_checkpoint_path,
         trust_remote_code=args.trust_remote_code,
         provider_overrides={
             "tensor_model_parallel_size": 1,  # Tensor parallelism is not supported
@@ -515,7 +517,7 @@ def main(args: argparse.Namespace):
     # the full VLM forward over image-text pairs.
     if use_image_calib:
         processor = AutoProcessor.from_pretrained(
-            args.hf_model_path, trust_remote_code=args.trust_remote_code
+            args.local_checkpoint_path, trust_remote_code=args.trust_remote_code
         )
         forward_loop = get_megatron_vlm_calibration_forward_loop(
             unwrapped_model,  # full VLM (vision encoder + projector + language model)
@@ -649,7 +651,7 @@ def score_func(m):
             model,
             args.output_megatron_path,
             # Recorded in the checkpoint's config: the Hub ID stays valid on other hosts.
-            hf_tokenizer_path=args.hf_model_name or args.hf_model_path,
+            hf_tokenizer_path=args.hub_model_id or args.local_checkpoint_path,
             hf_tokenizer_kwargs=tokenizer_kwargs,
         )
         args.checkpoint_exported = True
@@ -758,7 +760,7 @@ def score_func(m):
             # from_hf_config can't infer it since AutoConfig consumes the kwarg.
             pruned_bridge.trust_remote_code = args.trust_remote_code
             pruned_bridge.save_hf_pretrained(
-                model, args.output_hf_path, source_path=args.hf_model_path
+                model, args.output_hf_path, source_path=args.local_checkpoint_path
             )
         else:
             dummy_model_cls = AutoModelForImageTextToText if is_vlm else AutoModelForCausalLM
@@ -772,7 +774,7 @@ def score_func(m):
 
         # Every rank gets here (save_hf_weights is collective); one writer is enough.
         if dist.is_master():
-            copy_non_model_files(args.hf_model_path, args.output_hf_path)
+            copy_non_model_files(args.local_checkpoint_path, args.output_hf_path)
         args.checkpoint_exported = True
         print_rank_0(f"Saved pruned model to {args.output_hf_path} in HF checkpoint format")
 
diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py
index a30357dc567..efdab627b36 100644
--- a/examples/megatron_bridge/quantize.py
+++ b/examples/megatron_bridge/quantize.py
@@ -333,20 +333,22 @@ def get_quant_config(args: argparse.Namespace) -> dict:
 
 def main(args: argparse.Namespace):
     # --hf_model_name_or_path stays as given. Later steps read the local copy of the whole
-    # checkpoint in hf_model_path; 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.hf_model_name_or_path)
+    # checkpoint in local_checkpoint_path; 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.hf_model_name_or_path
+    )
     trust_remote_code = is_safe_repo(
-        trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_path
+        trust_remote_code=args.trust_remote_code, hf_path=args.local_checkpoint_path
     )
 
     moe_grouped_gemm = use_moe_grouped_gemm(
-        args.hf_model_path,
+        args.local_checkpoint_path,
         trust_remote_code=trust_remote_code,
         force_sequential=args.no_moe_grouped_gemm,
     )
 
     bridge, _provider, model, unwrapped_model, tokenizer = load_mbridge_model_from_hf(
-        hf_model_name_or_path=args.hf_model_path,
+        hf_model_name_or_path=args.local_checkpoint_path,
         trust_remote_code=trust_remote_code,
         moe_grouped_gemm=moe_grouped_gemm,
         provider_overrides={
@@ -451,7 +453,7 @@ def forward_loop(_model=None):
         # VLMs: drive the full VLM forward on image-text pairs so the language model's quantizers
         # see vision-conditioned activations (we still quantize the LM only).
         processor = AutoProcessor.from_pretrained(
-            args.hf_model_path, trust_remote_code=trust_remote_code
+            args.local_checkpoint_path, trust_remote_code=trust_remote_code
         )
         forward_loop = get_megatron_vlm_calibration_forward_loop(
             unwrapped_model,  # full VLM (vision encoder + projector + language model)
@@ -496,7 +498,7 @@ def forward_loop(_model=None):
         model,
         args.export_megatron_path,
         # Recorded in the checkpoint's config: the Hub ID stays valid on other hosts.
-        hf_tokenizer_path=args.hf_model_name or args.hf_model_path,
+        hf_tokenizer_path=args.hub_model_id or args.local_checkpoint_path,
         hf_tokenizer_kwargs={"trust_remote_code": trust_remote_code},
     )
     args.checkpoint_exported = True

From d294c7d22042c30c3716682c10633e9ce5662b05 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 7 Oct 2026 08:12:32 +0000
Subject: [PATCH 03/11] Document that export_mcore_gpt_to_hf skips non-model
 files for a Hub model ID

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/unified_export_megatron.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index 7e227a9ddca..ee2dfcb24c2 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -127,6 +127,10 @@ class GPTModelExporter:
             pretrained model hosted inside a model repo on huggingface.co; or
             a *directory* containing model weights saved using
             [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
+            The source's non-model files (tokenizer, processor, generation config, remote code)
+            are copied into the export from a local directory only; for a model ID they are
+            skipped with a warning. Get a local copy first with
+            `modelopt.torch.export.ensure_local_checkpoint` to include them.
         export_extra_modules: If True, export extra modules like medusa_heads or
             eagle_module. Otherwise, only export the base model.
         dtype: The weights data type to export the unquantized layers.
@@ -2308,6 +2312,10 @@ def export_mcore_gpt_to_hf(
             pretrained model hosted inside a model repo on huggingface.co; or
             a *directory* containing model weights saved using
             [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
+            The source's non-model files (tokenizer, processor, generation config, remote code)
+            are copied into the export from a local directory only; for a model ID they are
+            skipped with a warning. Get a local copy first with
+            `modelopt.torch.export.ensure_local_checkpoint` to include them.
         export_extra_modules: If True, export extra modules like medusa_heads or
             eagle_module. Otherwise, only export the base model.
         dtype: The weights data type to export the unquantized layers.

From 3abd7a210046103d2650e33893c2b4fd01cee834 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 7 Oct 2026 17:51:15 +0000
Subject: [PATCH 04/11] distill.py: build and export the student from a local
 checkpoint

export_llm_to_hf and save_vlm_to_hf now copy the source's non-model files,
which reads the student from local disk only, but distill.py still passed
--student_hf_path unresolved: a Hub ID would fail the export after
training. Resolve it with ensure_local_checkpoint at the start of main()
and build, probe and export the student from that local copy;
--student_hf_path stays as given.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 examples/megatron_bridge/distill.py           | 19 +++++--
 .../examples/megatron_bridge/test_distill.py  | 53 +++++++++++++++++++
 2 files changed, 67 insertions(+), 5 deletions(-)

diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py
index 333ccbbcb3d..8fbb3ab2fb1 100644
--- a/examples/megatron_bridge/distill.py
+++ b/examples/megatron_bridge/distill.py
@@ -74,6 +74,7 @@
 
 import modelopt.torch.distill as mtd
 import modelopt.torch.utils.distributed as dist
+from modelopt.torch.export import ensure_local_checkpoint
 from modelopt.torch.opt.conversion import ModeloptStateManager
 from modelopt.torch.utils import print_args, print_rank_0, warn_rank_0
 from modelopt.torch.utils.mlflow import Tool, masked_args
@@ -497,6 +498,12 @@ def _tokenizer_prepends_bos(args) -> bool:
 
 
 def main(args: argparse.Namespace, owns_the_run: bool = True):
+    # --student_hf_path stays as given. The student is built and exported from the local copy of
+    # the whole checkpoint in student_local_checkpoint_path, since the HF export reads the source's
+    # files from local disk only; student_hub_model_id is the Hub ID, or None for a local path.
+    args.student_hub_model_id, args.student_local_checkpoint_path = ensure_local_checkpoint(
+        args.student_hf_path
+    )
     student_has_modelopt_state = args.student_megatron_path is not None and has_modelopt_state(
         args.student_megatron_path
     )
@@ -505,7 +512,7 @@ def main(args: argparse.Namespace, owns_the_run: bool = True):
     # reads either layout, so it keeps the faster grouped GEMM.
     moe_grouped_gemm = (
         use_moe_grouped_gemm(
-            args.student_hf_path,
+            args.student_local_checkpoint_path,
             trust_remote_code=args.trust_remote_code,
             force_sequential=args.no_moe_grouped_gemm,
         )
@@ -553,7 +560,7 @@ def _build_model_provider(hf_path, load_weights=True, moe_grouped_gemm=True):
     # the built student inside the patched provide() below).
     # Only the student's layout is pinned -- it must match --student_megatron_path (see quantize.py).
     student_provider = _build_model_provider(
-        args.student_hf_path,
+        args.student_local_checkpoint_path,
         load_weights=args.student_megatron_path is None,
         moe_grouped_gemm=moe_grouped_gemm,
     )
@@ -591,7 +598,9 @@ def _build_model_provider(hf_path, load_weights=True, moe_grouped_gemm=True):
 
     # HF VLM configs expose ``vision_config``; Megatron-Bridge nests the text model under
     # ``language_model`` (used as ``distill_submodule`` below).
-    is_vlm = is_vlm_config(args.student_hf_path, trust_remote_code=args.trust_remote_code)
+    is_vlm = is_vlm_config(
+        args.student_local_checkpoint_path, trust_remote_code=args.trust_remote_code
+    )
 
     if is_vlm:
         warn_rank_0(
@@ -808,7 +817,7 @@ def _restore_student_hook(model_chunks):
         save_vlm_to_hf(
             full_student,
             args.hf_export_path,
-            args.student_hf_path,
+            args.student_local_checkpoint_path,
             trust_remote_code=args.trust_remote_code,
         )
         print_rank_0(f"Saved distilled VLM to {args.hf_export_path} in HF format")
@@ -825,7 +834,7 @@ def _restore_student_hook(model_chunks):
             export_llm_to_hf(
                 megatron_path=f"{checkpoint_dir}/iter_{args.train_iters:07d}",
                 hf_export_path=args.hf_export_path,
-                student_hf_path=args.student_hf_path,
+                student_hf_path=args.student_local_checkpoint_path,
                 template_hf=args.student_hf_model,
                 trust_remote_code=args.trust_remote_code,
             )
diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py
index 3bdcc460c27..61bde312901 100644
--- a/tests/examples/megatron_bridge/test_distill.py
+++ b/tests/examples/megatron_bridge/test_distill.py
@@ -16,6 +16,7 @@
 
 import argparse
 import json
+import shutil
 import sys
 from pathlib import Path
 
@@ -31,6 +32,7 @@
 )
 from transformers import AutoModelForImageTextToText
 
+import modelopt.torch.export
 from modelopt.torch.puzzletron.anymodel import convert_model
 
 sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "examples" / "megatron_bridge"))
@@ -239,6 +241,57 @@ def test_distill_validate_only(tmp_path, num_gpus):
     assert not (hf_export_path / "config.json").exists()
 
 
+@pytest.mark.parametrize("model", ["llm", "vlm"])
+@pytest.mark.timeout(420)
+def test_distill_hf_export_reads_the_resolved_local_checkpoint(tmp_path, monkeypatch, model):
+    """The inline HF export reads the student from the copy ensure_local_checkpoint resolved.
+
+    A Hub ID cannot run end to end here (get_args() loads the tokenizers from --student_hf_path),
+    so the resolver reports a Hub ID and returns a different local copy holding a marker file: the
+    marker reaching the export shows which copy the export read. One rank, so the step runs in
+    this process and sees the patch.
+    """
+    if model == "llm":
+        student_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True)
+    else:
+        student_hf_path = create_tiny_qwen3_5_vl_dir(
+            tmp_path, with_tokenizer=True, num_hidden_layers=2, intermediate_size=128
+        )
+    resolved = tmp_path / "resolved"
+    shutil.copytree(student_hf_path, resolved)
+    (resolved / "resolved_only.md").write_text("from the resolved copy\n")
+    resolved_from = []
+
+    def ensure_local_checkpoint(model_name_or_path, group=None):
+        resolved_from.append(str(model_name_or_path))
+        return "org/tiny-student", str(resolved)
+
+    monkeypatch.setattr(modelopt.torch.export, "ensure_local_checkpoint", ensure_local_checkpoint)
+    distilled_hf_path = tmp_path / "distilled_hf"
+    distill_cmd_parts = extend_cmd_parts(
+        ["torchrun", "--nproc_per_node=1", "distill.py", "--use_mock_data"],
+        student_hf_path=student_hf_path,
+        teacher_hf_path=student_hf_path,
+        output_dir=tmp_path / "distill_output",
+        tp_size=1,
+        pp_size=1,
+        seq_length=16,
+        mbs=1,
+        gbs=4,
+        train_iters=1,
+        lr_warmup_iters=1,
+        eval_interval=1,
+        eval_iters=1,
+        log_interval=1,
+        hf_export_path=distilled_hf_path,
+    )
+    run_example_command(distill_cmd_parts, example_path="megatron_bridge")
+
+    assert resolved_from == [str(student_hf_path)]
+    assert (distilled_hf_path / "config.json").exists()
+    assert (distilled_hf_path / "resolved_only.md").read_text() == "from the resolved copy\n"
+
+
 # NOTE: Qwen3.5-VL-MoE covered by test_qad.py
 @pytest.mark.timeout(420)  # sometimes times out in CI env with default 300s timeout
 def test_distill_vlm(tmp_path, num_gpus):

From fb3be2e02f2c990a6b8ec63edc9ac07d1694b847 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 7 Oct 2026 17:59:09 +0000
Subject: [PATCH 05/11] Patch the runner's cached distill module in the
 resolved-checkpoint test

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 tests/examples/megatron_bridge/test_distill.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py
index 61bde312901..1c84e8b285c 100644
--- a/tests/examples/megatron_bridge/test_distill.py
+++ b/tests/examples/megatron_bridge/test_distill.py
@@ -22,6 +22,7 @@
 
 import pytest
 import torch
+from _test_utils.examples.megatron_example_runner import _load_example_module
 from _test_utils.examples.run_command import extend_cmd_parts, run_example_command
 from _test_utils.torch.puzzletron.utils import create_and_save_small_hf_model
 from _test_utils.torch.transformers_models import (
@@ -32,7 +33,6 @@
 )
 from transformers import AutoModelForImageTextToText
 
-import modelopt.torch.export
 from modelopt.torch.puzzletron.anymodel import convert_model
 
 sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "examples" / "megatron_bridge"))
@@ -249,7 +249,7 @@ def test_distill_hf_export_reads_the_resolved_local_checkpoint(tmp_path, monkeyp
     A Hub ID cannot run end to end here (get_args() loads the tokenizers from --student_hf_path),
     so the resolver reports a Hub ID and returns a different local copy holding a marker file: the
     marker reaching the export shows which copy the export read. One rank, so the step runs in
-    this process and sees the patch.
+    this process, on the module the runner caches -- the one patched here.
     """
     if model == "llm":
         student_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True)
@@ -266,7 +266,8 @@ def ensure_local_checkpoint(model_name_or_path, group=None):
         resolved_from.append(str(model_name_or_path))
         return "org/tiny-student", str(resolved)
 
-    monkeypatch.setattr(modelopt.torch.export, "ensure_local_checkpoint", ensure_local_checkpoint)
+    distill = _load_example_module("distill.py", "megatron_bridge")
+    monkeypatch.setattr(distill, "ensure_local_checkpoint", ensure_local_checkpoint)
     distilled_hf_path = tmp_path / "distilled_hf"
     distill_cmd_parts = extend_cmd_parts(
         ["torchrun", "--nproc_per_node=1", "distill.py", "--use_mock_data"],

From a5d950248bff410f752c3c8c2d27b52ba19e5637 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 7 Oct 2026 18:01:39 +0000
Subject: [PATCH 06/11] Test that quantize, export and prune run a Hub ID from
 its local copy

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 tests/examples/megatron_bridge/conftest.py    | 46 +++++++++++++++++++
 .../megatron_bridge/test_prune_minitron.py    | 36 +++++++++++++++
 .../megatron_bridge/test_quantize_export.py   | 38 +++++++++++++++
 3 files changed, 120 insertions(+)

diff --git a/tests/examples/megatron_bridge/conftest.py b/tests/examples/megatron_bridge/conftest.py
index 1093fb7b281..f058c72c4b5 100644
--- a/tests/examples/megatron_bridge/conftest.py
+++ b/tests/examples/megatron_bridge/conftest.py
@@ -15,9 +15,12 @@
 """Run example steps without a ``torchrun`` launch; see ``_test_utils.examples.megatron_example_runner``."""
 
 import os
+import shutil
 
 import pytest
+import yaml
 from _test_utils.examples.megatron_example_runner import (
+    _load_example_module,
     reset_megatron_global_state,
     run_example_step,
 )
@@ -55,3 +58,46 @@ def _isolate_megatron_global_state():
         reset_megatron_global_state()
         os.environ.clear()
         os.environ.update(env_before)
+
+
+class FakeHubCheckpoint:
+    """Serves a local checkpoint to example scripts as if it were Hub model ``hub_model_id``.
+
+    ``serve(checkpoint, *scripts)`` patches ``ensure_local_checkpoint`` on each script's module --
+    the one the in-process runner caches and runs -- to resolve ``hub_model_id``, and only it, to
+    a copy of ``checkpoint`` holding the file ``marker``. Pass ``hub_model_id`` on the command
+    line: a step that read it without resolving it fails, and ``marker`` reaching an export shows
+    the export read the resolved copy. Single-rank steps only; torchrun's workers would not see
+    the patch.
+    """
+
+    hub_model_id = "org/tiny-model"
+    marker = "from_the_resolved_copy.md"
+
+    def __init__(self, monkeypatch, tmp_path):
+        self._monkeypatch = monkeypatch
+        self._resolved = tmp_path / "resolved_hub_checkpoint"
+
+    def serve(self, checkpoint, *scripts):
+        shutil.copytree(checkpoint, self._resolved)
+        (self._resolved / self.marker).write_text("from the resolved copy\n")
+
+        def ensure_local_checkpoint(model_name_or_path, group=None):
+            assert str(model_name_or_path) == self.hub_model_id
+            return self.hub_model_id, str(self._resolved)
+
+        for script in scripts:
+            module = _load_example_module(script, "megatron_bridge")
+            self._monkeypatch.setattr(module, "ensure_local_checkpoint", ensure_local_checkpoint)
+
+    @staticmethod
+    def saved_tokenizer_model(megatron_path) -> str:
+        """The tokenizer a Megatron-Bridge checkpoint records in its ``run_config.yaml``."""
+        (run_config,) = sorted(megatron_path.rglob("run_config.yaml"))
+        return yaml.safe_load(run_config.read_text())["tokenizer"]["tokenizer_model"]
+
+
+@pytest.fixture
+def fake_hub_checkpoint(monkeypatch, tmp_path):
+    """See :class:`FakeHubCheckpoint`."""
+    return FakeHubCheckpoint(monkeypatch, tmp_path)
diff --git a/tests/examples/megatron_bridge/test_prune_minitron.py b/tests/examples/megatron_bridge/test_prune_minitron.py
index 71230f54dcd..05c7781f3a4 100644
--- a/tests/examples/megatron_bridge/test_prune_minitron.py
+++ b/tests/examples/megatron_bridge/test_prune_minitron.py
@@ -100,6 +100,42 @@ def test_prune_minitron(tmp_path, num_gpus, create_teacher, expected_pruned_conf
         assert getattr(pruned_model.config, field) == expected
 
 
+@pytest.mark.parametrize("output", ["megatron", "hf"])
+def test_prune_minitron_a_hub_model_id(tmp_path, fake_hub_checkpoint, output):
+    """A Hub ID is pruned from its local copy: the HF output carries that copy's non-model files,
+    and the Megatron output records the Hub ID as its tokenizer so it stays valid on other hosts."""
+    teacher_hf_path, teacher_model = create_tiny_qwen3_dir(
+        tmp_path, with_tokenizer=True, return_model=True, num_hidden_layers=1
+    )
+    fake_hub_checkpoint.serve(teacher_hf_path, "prune_minitron.py")
+    pruned_path = tmp_path / "pruned"
+    prune_command_parts = extend_cmd_parts(
+        ["torchrun", "--nproc_per_node=1", "prune_minitron.py"],
+        hf_model_name_or_path=fake_hub_checkpoint.hub_model_id,
+        **{f"output_{output}_path": pruned_path},
+        pp_size=1,
+        calib_dataset_name="cnn_dailymail",
+        calib_num_samples=8,
+        seq_length=16,
+        prune_target_params=int(sum(p.numel() for p in teacher_model.parameters()) * 0.8),
+        prune_score_func="mmlu_1pct_bs32",
+        score_lower_bound=0.0,
+        ss_channel_divisor=4,
+        hparams_to_skip="num_attention_heads",
+        top_k=1,
+    )
+    run_example_command(prune_command_parts, example_path="megatron_bridge")
+
+    if output == "megatron":
+        assert (
+            fake_hub_checkpoint.saved_tokenizer_model(pruned_path)
+            == fake_hub_checkpoint.hub_model_id
+        )
+    else:
+        assert (pruned_path / "config.json").exists()
+        assert (pruned_path / fake_hub_checkpoint.marker).exists()
+
+
 @pytest.mark.parametrize(
     "create_teacher",
     [
diff --git a/tests/examples/megatron_bridge/test_quantize_export.py b/tests/examples/megatron_bridge/test_quantize_export.py
index 9d46482f328..d1bfaae6109 100644
--- a/tests/examples/megatron_bridge/test_quantize_export.py
+++ b/tests/examples/megatron_bridge/test_quantize_export.py
@@ -119,3 +119,41 @@ def test_quantize_and_export(
     # )
     # outputs = llm.generate(["Hello!"], vllm.SamplingParams(max_tokens=4))
     # assert outputs and outputs[0].outputs and outputs[0].outputs[0].text
+
+
+@pytest.mark.timeout(360)
+def test_quantize_and_export_a_hub_model_id(tmp_path: Path, fake_hub_checkpoint):
+    """A Hub ID is made local once: both steps load and export from that copy, while the
+    Megatron checkpoint records the Hub ID as its tokenizer so it stays valid on other hosts."""
+    fake_hub_checkpoint.serve(
+        create_tiny_qwen3_moe_dir(tmp_path, with_tokenizer=True, **_DENSE_KWARGS),
+        "quantize.py",
+        "export_quantized_megatron_to_hf.py",
+    )
+    megatron_path = tmp_path / "quantized_megatron"
+    hf_export_path = tmp_path / "quantized_hf"
+
+    quantize_cmd = extend_cmd_parts(
+        ["torchrun", "--nproc_per_node=1", "quantize.py", "--skip_generate"],
+        hf_model_name_or_path=fake_hub_checkpoint.hub_model_id,
+        recipe="general/ptq/nvfp4_default-kv_fp8",
+        calib_dataset_name="cnn_dailymail",
+        calib_num_samples=4,
+        calib_batch_size=2,
+        seq_length=16,
+        export_megatron_path=megatron_path,
+    )
+    run_example_command(quantize_cmd, example_path="megatron_bridge", setup_free_port=True)
+    assert (
+        fake_hub_checkpoint.saved_tokenizer_model(megatron_path) == fake_hub_checkpoint.hub_model_id
+    )
+
+    export_cmd = extend_cmd_parts(
+        ["torchrun", "--nproc_per_node=1", "export_quantized_megatron_to_hf.py"],
+        hf_model_name_or_path=fake_hub_checkpoint.hub_model_id,
+        megatron_path=megatron_path,
+        export_unified_hf_path=hf_export_path,
+    )
+    run_example_command(export_cmd, example_path="megatron_bridge", setup_free_port=True)
+    assert (hf_export_path / "hf_quant_config.json").exists()
+    assert (hf_export_path / fake_hub_checkpoint.marker).exists()

From c2a27d742dbef47254fbdc941e8ee29cc30d5287 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Mon, 19 Oct 2026 20:04:30 +0000
Subject: [PATCH 07/11] Move the Megatron export changelog entries to 0.49

0.48 has branched (release/0.48.0), so this lands in 0.49: restore the
0.48 entry for the HF exporters to its released wording and file the
Megatron export's feature, breaking-change and deprecation entries under
0.49.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 CHANGELOG.rst | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 133d203e583..9ee5a31b096 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -21,6 +21,7 @@ Changelog
 - Add ``--sft_in_batch_packing``, ``--moe_aux_loss_coeff`` and ``--dsa_kernel_backend`` to ``examples/megatron_bridge/distill.py`` (``--dsa_kernel_backend`` also to ``quantize.py``): pack ``--sft_hf_dataset`` conversations per micro-batch instead of padding each to ``--seq_length``, override the MoE load-balancing loss weight (e.g. 0 for padded SFT data), and select fused DeepSeek sparse-attention kernels for long sequences.
 - Add ``--vllm_fakequant_export`` to ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` to evaluate formats vLLM has no kernels for (e.g. GGML IQ) as BF16 fake-quant checkpoints, and allow ``--ep_size`` > 1 for checkpoints built with ``--no_moe_grouped_gemm``.
 - GGML GPTQ checkpoints from ``examples/megatron_bridge`` keep their calibrated codes when the Megatron checkpoint is reloaded for QAD or export, and layerwise calibration now runs on Megatron vision-language models.
+- The Megatron unified HF export (``export_mcore_gpt_to_hf``) now copies the source checkpoint's non-model files (tokenizer, processor, remote code, chat templates, ...) into the export unchanged, from local disk only. The ``examples/megatron_bridge`` scripts make a Hugging Face Hub model local first with ``modelopt.torch.export.ensure_local_checkpoint``.
 
 *Misc*
 
@@ -28,8 +29,12 @@ Changelog
 
 **Backward Breaking Changes**
 
+- ``export_mcore_gpt_to_hf`` given a Hugging Face Hub model ID no longer writes the tokenizer, processor, generation config or remote-code files. Pass a local directory instead, e.g. from ``modelopt.torch.export.ensure_local_checkpoint``.
+
 **Deprecations**
 
+- ``modelopt.torch.export.copy_non_safetensor_files_from_ckpt`` is deprecated; use ``copy_non_model_files``.
+
 **Bug Fixes**
 
 - Fix ONNX quantization placing a per-tensor Q/DQ pair on a rank-0 (scalar) constant input, which TensorRT's ONNX parser rejects (``convertAxis`` asserts the quantization axis lies in ``[0, nbDims]`` but ``nbDims == 0``). Rank-0 constant operands are now kept unquantized in the int8, fp8, and autotune export paths; the existing single-element Mul/Div/Pad stripping stays gated under ``target_dla``.
@@ -108,7 +113,7 @@ Changelog
 - Add ``--modelopt-*`` options to the ``examples/vllm_serve`` launcher for fakequant calibration and checkpoint reload. Point it at an exported fakequant model directory to reload automatically: Hugging Face exports use ``vllm_fq_modelopt_state.pth``, while Megatron-Core exports now include ``quant_recipe.yaml`` alongside ``quantizer_state.pth`` for reload.
 - Add a ``kd_loss_alpha`` option to HuggingFace ``KDTrainer`` to blend cross-entropy and knowledge distillation losses during training.
 - A tracked ``examples/hf_ptq/hf_ptq.py`` run now writes ``.experiment.json`` into ``--export_path`` and uploads the same file with the run, so a checkpoint on disk names the experiment and MLflow run id that produced it. The pointer is written only once the export completes, and an export that is not tracked removes one it would otherwise inherit from a reused ``--export_path`` or from a quantized source checkpoint.
-- ``export_hf_checkpoint``, the vLLM fake-quant export and the Megatron unified HF export now copy the source checkpoint's non-model files (tokenizer, processor, remote code, chat templates, ...) into the export unchanged. They read the source from local disk only: get a local copy of a Hugging Face Hub model with ``modelopt.torch.export.ensure_local_checkpoint`` and load it from there.
+- ``export_hf_checkpoint`` and the vLLM fake-quant export now copy the source checkpoint's non-model files (tokenizer, processor, remote code, chat templates, ...) into the export unchanged. They read the source from local disk only: get a local copy of a Hugging Face Hub model with ``modelopt.torch.export.ensure_local_checkpoint`` and load it from there.
 
 - Add Parallel Decoding Distillation (PDD) to ``modelopt.torch.fastgen`` with Qwen-Image training, distributed-checkpoint export, and PDD-2/4/8 inference. AutoModel remains an unmodified pinned runtime dependency.
 
@@ -120,7 +125,6 @@ Changelog
 - ``LogitsAndIntermediatesLossBalancer`` (Megatron distillation plugin) no longer rescales the distillation loss to the magnitude of the LM loss when the LM loss is included. The total is now the fixed convex combination ``(1 - alpha) * lm_loss + alpha * kd_loss`` with ``DistillationConfig.kd_loss_alpha`` (in [0, 1]), matching Megatron-LM's offline cached-logits KD. The default ``kd_loss_alpha=1.0`` skips the LM loss, as before. ``DistillationConfig.skip_lm_loss`` and ``kd_loss_scale`` are removed and raise ``ValueError`` if passed; set ``kd_loss_alpha`` instead. In ``examples/megatron_bridge/distill.py``, ``--kd_loss_alpha`` replaces ``--no_skip_lm_loss`` and ``--kd_loss_scale``.
 - The ``examples/vllm_serve`` fakequant launcher no longer supports vLLM 0.9.0. Upgrade to a version listed as tested in the example README.
 - The Megatron-Core DeepSeek-V4 indexer (``CSAIndexer``) is now a quantization module and persists its quantizer state in the checkpoint as ``indexer._extra_state``. A DeepSeek-V4 model quantized with an earlier release resumes from its ``torch_dist`` checkpoint only with a non-strict load (``--dist-ckpt-strictness log_unexpected`` in Megatron-LM) until it is saved again.
-- ``export_mcore_gpt_to_hf`` given a Hugging Face Hub model ID no longer writes the tokenizer, processor, generation config or remote-code files. Pass a local directory instead, e.g. from ``modelopt.torch.export.ensure_local_checkpoint``.
 - ``examples/hf_ptq`` no longer detects MTP layers by name. Weights the loader could not place -- an MTP head, an auxiliary tower -- are identified from Transformers' own accounting: the model is loaded with ``from_pretrained(..., output_loading_info=True)`` and the reported ``unexpected_keys`` (present in the checkpoint, not in the model's architecture) are recorded on the model and carried into the export unchanged. Everything the loader *did* place goes through the normal export path. This removes ``load_mtp_weights``, ``mtp_layer_prefixes_from_checkpoint`` and their support matrix of MTP storage conventions, along with ``_add_mtp_exclusions`` and the pre-quantization ``enable: False`` entries ``hf_ptq`` appended to the recipe's ``quant_cfg``. Two consequences: MTP layers now follow the recipe like any other module instead of being force-excluded by the script -- matching ``examples/megatron_bridge``, which has no MTP-specific code at all -- and ``quantization_config.ignore`` can no longer claim a layer is unquantized that the export in fact quantized. Recipes importing ``configs/ptq/units/default_disabled_quantizers`` still disable ``mtp.*``, so their behaviour is unchanged; a recipe omitting that unit will now quantize an MTP the model actually built.
 - ``examples/hf_ptq --vllm_fakequant_export`` now raises ``NotImplementedError`` when the checkpoint holds weights the model has no parameter for and a shard actually provides them (an MTP head, an auxiliary tower). The fake-quant exporter writes only model-backed state, so it would otherwise drop those weights silently -- and a fake-quant checkpoint is evaluated, where a missing head changes the score rather than failing loudly. Use the unified HF export, which carries them through. Buffers Transformers recomputes are not weights to lose: ``*.inv_freq`` is skipped even when a shard provides it, since older Llama/Mistral-lineage conversions do list it in the index and refusing an export over it would reject checkpoints that export correctly today. The check runs immediately after the model loads, not at export time, so an incompatible run fails before calibration rather than after it.
 - The ``modelopt.onnx.quantization.ort_patching`` module has been removed with no
@@ -174,7 +178,6 @@ Changelog
 - The ``evaluation`` agent skill no longer supports GDPVal: its task recipe, example config and Apptainer SIF helper are removed, and the shared NeMo Gym machinery it carried now lives at ``references/gym.md`` with the launcher wrapper renamed ``scripts/nel-gym.sh``. GDPVal is an AA-suite member, so an "AA" request now generates the ``aa/`` tasks only -- report per-task scores rather than an aggregate compared against a published AA Index.
 - DBRX model support is removed: ModelOpt no longer quantizes or exports DBRX (``DbrxForCausalLM``) checkpoints. Use ModelOpt 0.47 to quantize DBRX, or write your own plugin as described in the customized-model quantization guide.
 - Deprecate ``metadata.recipe_type`` in recipe YAML, for both single-file recipes and a directory recipe's ``metadata.yml``. A recipe now says what kind it is with a ``# modelopt-schema:`` comment naming its schema class, or by delegating to a recipe that does (single-file recipes only -- a directory recipe has no body to delegate through); ``recipe_type`` is still read and still honoured, so a recipe outside this repo keeps working unchanged, but new recipes should leave it out and every recipe shipped here has been converted. Where both are present they must agree, and so must a recipe and the recipe it delegates to -- a disagreement is an error rather than a silent preference.
-- ``modelopt.torch.export.copy_non_safetensor_files_from_ckpt`` is deprecated; use ``copy_non_model_files``.
 
 **Bug Fixes**
 

From d1bd8138ecfced7fa71fcb559736086c6a7f047b Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Mon, 19 Oct 2026 20:22:02 +0000
Subject: [PATCH 08/11] Document that the vLLM fake-quant Megatron export skips
 non-model files for a Hub model ID

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/plugins/vllm_fakequant_megatron.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py
index 1355b833189..de873f6d673 100644
--- a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py
+++ b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py
@@ -509,6 +509,10 @@ def export_mcore_gpt_to_hf_vllm_fq(
             pretrained model hosted inside a model repo on huggingface.co; or
             a *directory* containing model weights saved using
             [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
+            The source's non-model files (tokenizer, processor, generation config, remote code)
+            are copied into the export from a local directory only; for a model ID they are
+            skipped with a warning. Get a local copy first with
+            `modelopt.torch.export.ensure_local_checkpoint` to include them.
         export_extra_modules: If True, export extra modules like medusa_heads or
             eagle_module. Otherwise, only export the base model.
         dtype: The weights data type to export the unquantized layers.

From 7a6037942bd7f86b1ccc37867ccaa951b02047ed Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 20 Oct 2026 00:54:24 +0000
Subject: [PATCH 09/11] Fold the Hub-ID regressions into the existing
 Megatron-Bridge tests

The example tests are slow, so check the Hub-ID handling inside the
runs that already exist instead of adding new ones: quantize/export and
prune take the Hub ID, distillation's inline export and both standalone
distillation exports read the resolved copy, and the quantized checkpoint
records the Hub ID as its tokenizer. fake_hub_checkpoint only patches a
single-rank step, which runs in the pytest process; with more ranks the
tests run on the local path as before.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 tests/examples/megatron_bridge/conftest.py    | 23 +++--
 .../examples/megatron_bridge/test_distill.py  | 84 ++++++-------------
 .../megatron_bridge/test_prune_minitron.py    | 47 ++---------
 .../megatron_bridge/test_quantize_export.py   | 59 ++++---------
 4 files changed, 66 insertions(+), 147 deletions(-)

diff --git a/tests/examples/megatron_bridge/conftest.py b/tests/examples/megatron_bridge/conftest.py
index f058c72c4b5..beb86565eb6 100644
--- a/tests/examples/megatron_bridge/conftest.py
+++ b/tests/examples/megatron_bridge/conftest.py
@@ -63,12 +63,14 @@ def _isolate_megatron_global_state():
 class FakeHubCheckpoint:
     """Serves a local checkpoint to example scripts as if it were Hub model ``hub_model_id``.
 
-    ``serve(checkpoint, *scripts)`` patches ``ensure_local_checkpoint`` on each script's module --
-    the one the in-process runner caches and runs -- to resolve ``hub_model_id``, and only it, to
-    a copy of ``checkpoint`` holding the file ``marker``. Pass ``hub_model_id`` on the command
-    line: a step that read it without resolving it fails, and ``marker`` reaching an export shows
-    the export read the resolved copy. Single-rank steps only; torchrun's workers would not see
-    the patch.
+    ``serve(checkpoint, *scripts, nproc=...)`` patches ``ensure_local_checkpoint`` on each script's
+    module -- the one the in-process runner caches and runs -- to resolve ``hub_model_id`` (or
+    ``checkpoint`` itself, for a script that must read its argument before resolving it) to a copy
+    of ``checkpoint`` holding the file ``marker``, and returns ``hub_model_id`` to pass on the
+    command line: a step that read it without resolving it fails, and ``marker`` reaching an
+    export shows the export read the resolved copy. Only a single-rank step runs in this process
+    and sees the patch, so with ``nproc > 1`` nothing is patched, ``checkpoint`` is returned
+    unchanged and ``active`` is False: check the marker and the Hub ID only when it is True.
     """
 
     hub_model_id = "org/tiny-model"
@@ -77,18 +79,23 @@ class FakeHubCheckpoint:
     def __init__(self, monkeypatch, tmp_path):
         self._monkeypatch = monkeypatch
         self._resolved = tmp_path / "resolved_hub_checkpoint"
+        self.active = False
 
-    def serve(self, checkpoint, *scripts):
+    def serve(self, checkpoint, *scripts, nproc: int) -> str:
+        self.active = nproc == 1
+        if not self.active:
+            return str(checkpoint)
         shutil.copytree(checkpoint, self._resolved)
         (self._resolved / self.marker).write_text("from the resolved copy\n")
 
         def ensure_local_checkpoint(model_name_or_path, group=None):
-            assert str(model_name_or_path) == self.hub_model_id
+            assert str(model_name_or_path) in (self.hub_model_id, str(checkpoint))
             return self.hub_model_id, str(self._resolved)
 
         for script in scripts:
             module = _load_example_module(script, "megatron_bridge")
             self._monkeypatch.setattr(module, "ensure_local_checkpoint", ensure_local_checkpoint)
+        return self.hub_model_id
 
     @staticmethod
     def saved_tokenizer_model(megatron_path) -> str:
diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py
index 1c84e8b285c..6d146957426 100644
--- a/tests/examples/megatron_bridge/test_distill.py
+++ b/tests/examples/megatron_bridge/test_distill.py
@@ -16,13 +16,11 @@
 
 import argparse
 import json
-import shutil
 import sys
 from pathlib import Path
 
 import pytest
 import torch
-from _test_utils.examples.megatron_example_runner import _load_example_module
 from _test_utils.examples.run_command import extend_cmd_parts, run_example_command
 from _test_utils.torch.puzzletron.utils import create_and_save_small_hf_model
 from _test_utils.torch.transformers_models import (
@@ -42,9 +40,17 @@
 
 
 @pytest.mark.timeout(360)  # distill + inline HF export + export of every iteration
-def test_distill_llm_hf_export(tmp_path, num_gpus):
-    """Distill with an inline HF export, then re-export every saved iteration standalone."""
+def test_distill_llm_hf_export(tmp_path, num_gpus, fake_hub_checkpoint):
+    """Distill with an inline HF export, then re-export every saved iteration standalone.
+
+    Both exports read the student from the copy ``ensure_local_checkpoint`` resolved (see
+    ``fake_hub_checkpoint``). distill.py keeps the local path on its command line, since get_args()
+    loads the tokenizer from it; the standalone export is passed the Hub ID.
+    """
     teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True)
+    student_export_arg = fake_hub_checkpoint.serve(
+        teacher_hf_path, "distill.py", "export_distilled_megatron_to_hf.py", nproc=num_gpus
+    )
     train_iters = 2
     distill_output_dir = tmp_path / "distill_output"
     distilled_hf_path = tmp_path / "distilled_hf"
@@ -72,11 +78,13 @@ def test_distill_llm_hf_export(tmp_path, num_gpus):
     checkpoint_root = distill_output_dir / "checkpoints"
     assert (checkpoint_root / f"iter_{train_iters:07d}").exists()
     assert (distilled_hf_path / "config.json").exists()
+    if fake_hub_checkpoint.active:
+        assert (distilled_hf_path / fake_hub_checkpoint.marker).exists()
 
     # export_distilled_megatron_to_hf.py converts any saved iteration, not just the final one.
     all_export_cmd_parts = extend_cmd_parts(
         ["torchrun", "--nproc_per_node=1", "export_distilled_megatron_to_hf.py"],
-        student_hf_path=teacher_hf_path,
+        student_hf_path=student_export_arg,
         megatron_path=checkpoint_root,
         hf_export_path=all_exports,
         export_iterations="all",
@@ -85,6 +93,8 @@ def test_distill_llm_hf_export(tmp_path, num_gpus):
 
     for it in range(1, train_iters + 1):
         assert (all_exports / f"iter_{it:07d}/config.json").exists()
+        if fake_hub_checkpoint.active:
+            assert (all_exports / f"iter_{it:07d}" / fake_hub_checkpoint.marker).exists()
 
 
 def test_distill_llm_sft(tmp_path, num_gpus):
@@ -241,61 +251,9 @@ def test_distill_validate_only(tmp_path, num_gpus):
     assert not (hf_export_path / "config.json").exists()
 
 
-@pytest.mark.parametrize("model", ["llm", "vlm"])
-@pytest.mark.timeout(420)
-def test_distill_hf_export_reads_the_resolved_local_checkpoint(tmp_path, monkeypatch, model):
-    """The inline HF export reads the student from the copy ensure_local_checkpoint resolved.
-
-    A Hub ID cannot run end to end here (get_args() loads the tokenizers from --student_hf_path),
-    so the resolver reports a Hub ID and returns a different local copy holding a marker file: the
-    marker reaching the export shows which copy the export read. One rank, so the step runs in
-    this process, on the module the runner caches -- the one patched here.
-    """
-    if model == "llm":
-        student_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True)
-    else:
-        student_hf_path = create_tiny_qwen3_5_vl_dir(
-            tmp_path, with_tokenizer=True, num_hidden_layers=2, intermediate_size=128
-        )
-    resolved = tmp_path / "resolved"
-    shutil.copytree(student_hf_path, resolved)
-    (resolved / "resolved_only.md").write_text("from the resolved copy\n")
-    resolved_from = []
-
-    def ensure_local_checkpoint(model_name_or_path, group=None):
-        resolved_from.append(str(model_name_or_path))
-        return "org/tiny-student", str(resolved)
-
-    distill = _load_example_module("distill.py", "megatron_bridge")
-    monkeypatch.setattr(distill, "ensure_local_checkpoint", ensure_local_checkpoint)
-    distilled_hf_path = tmp_path / "distilled_hf"
-    distill_cmd_parts = extend_cmd_parts(
-        ["torchrun", "--nproc_per_node=1", "distill.py", "--use_mock_data"],
-        student_hf_path=student_hf_path,
-        teacher_hf_path=student_hf_path,
-        output_dir=tmp_path / "distill_output",
-        tp_size=1,
-        pp_size=1,
-        seq_length=16,
-        mbs=1,
-        gbs=4,
-        train_iters=1,
-        lr_warmup_iters=1,
-        eval_interval=1,
-        eval_iters=1,
-        log_interval=1,
-        hf_export_path=distilled_hf_path,
-    )
-    run_example_command(distill_cmd_parts, example_path="megatron_bridge")
-
-    assert resolved_from == [str(student_hf_path)]
-    assert (distilled_hf_path / "config.json").exists()
-    assert (distilled_hf_path / "resolved_only.md").read_text() == "from the resolved copy\n"
-
-
 # NOTE: Qwen3.5-VL-MoE covered by test_qad.py
 @pytest.mark.timeout(420)  # sometimes times out in CI env with default 300s timeout
-def test_distill_vlm(tmp_path, num_gpus):
+def test_distill_vlm(tmp_path, num_gpus, fake_hub_checkpoint):
     # Self-distillation of a tiny VLM: only the language model is distilled; the vision tower and the
     # vision->language projector must be left byte-for-byte untouched.
     tp_size = num_gpus
@@ -347,10 +305,14 @@ def _is_language_model(name: str) -> bool:
     megatron_ckpt = distill_output_dir / f"checkpoints/iter_{train_iters:07d}"
     assert megatron_ckpt.exists()
 
-    # Separately convert the distilled Megatron checkpoint to HF with the standalone export script
+    # Separately convert the distilled Megatron checkpoint to HF with the standalone export script,
+    # passing the student as a Hub ID (see ``fake_hub_checkpoint``).
+    student_export_arg = fake_hub_checkpoint.serve(
+        vlm_hf_path, "export_distilled_megatron_to_hf.py", nproc=num_gpus
+    )
     export_cmd_parts = extend_cmd_parts(
         ["torchrun", f"--nproc_per_node={num_gpus}", "export_distilled_megatron_to_hf.py"],
-        student_hf_path=vlm_hf_path,
+        student_hf_path=student_export_arg,
         megatron_path=megatron_ckpt,
         hf_export_path=distilled_hf_path,
         tp_size=tp_size,
@@ -359,6 +321,8 @@ def _is_language_model(name: str) -> bool:
     run_example_command(export_cmd_parts, example_path="megatron_bridge")
 
     assert (distilled_hf_path / "config.json").exists()
+    if fake_hub_checkpoint.active:
+        assert (distilled_hf_path / fake_hub_checkpoint.marker).exists()
 
     # from_pretrained (default strict load) verifies the saved weights match the (unchanged) config.
     distilled_model = AutoModelForImageTextToText.from_pretrained(distilled_hf_path)
diff --git a/tests/examples/megatron_bridge/test_prune_minitron.py b/tests/examples/megatron_bridge/test_prune_minitron.py
index 05c7781f3a4..2cdee235d06 100644
--- a/tests/examples/megatron_bridge/test_prune_minitron.py
+++ b/tests/examples/megatron_bridge/test_prune_minitron.py
@@ -70,15 +70,20 @@
         ),
     ],
 )
-def test_prune_minitron(tmp_path, num_gpus, create_teacher, expected_pruned_config):
+def test_prune_minitron(
+    tmp_path, num_gpus, create_teacher, expected_pruned_config, fake_hub_checkpoint
+):
+    """The teacher is passed as a Hub ID (see ``fake_hub_checkpoint``): it is pruned from its local
+    copy, and the HF output carries that copy's non-model files."""
     teacher_hf_path, teacher_model = create_teacher(tmp_path, num_gpus)
+    teacher_arg = fake_hub_checkpoint.serve(teacher_hf_path, "prune_minitron.py", nproc=num_gpus)
     teacher_params = sum(p.numel() for p in teacher_model.parameters())
     prune_target_params = int(teacher_params * 0.8)
 
     pruned_path = tmp_path / "pruned"
     prune_command_parts = extend_cmd_parts(
         ["torchrun", f"--nproc_per_node={num_gpus}", "prune_minitron.py"],
-        hf_model_name_or_path=teacher_hf_path,
+        hf_model_name_or_path=teacher_arg,
         output_hf_path=pruned_path,
         pp_size=num_gpus,
         calib_dataset_name="cnn_dailymail",
@@ -94,48 +99,14 @@ def test_prune_minitron(tmp_path, num_gpus, create_teacher, expected_pruned_conf
     run_example_command(prune_command_parts, example_path="megatron_bridge")
 
     assert (pruned_path / "config.json").exists()
+    if fake_hub_checkpoint.active:
+        assert (pruned_path / fake_hub_checkpoint.marker).exists()
     pruned_model = AutoModelForCausalLM.from_pretrained(pruned_path)
     assert sum(p.numel() for p in pruned_model.parameters()) <= prune_target_params
     for field, expected in expected_pruned_config.items():
         assert getattr(pruned_model.config, field) == expected
 
 
-@pytest.mark.parametrize("output", ["megatron", "hf"])
-def test_prune_minitron_a_hub_model_id(tmp_path, fake_hub_checkpoint, output):
-    """A Hub ID is pruned from its local copy: the HF output carries that copy's non-model files,
-    and the Megatron output records the Hub ID as its tokenizer so it stays valid on other hosts."""
-    teacher_hf_path, teacher_model = create_tiny_qwen3_dir(
-        tmp_path, with_tokenizer=True, return_model=True, num_hidden_layers=1
-    )
-    fake_hub_checkpoint.serve(teacher_hf_path, "prune_minitron.py")
-    pruned_path = tmp_path / "pruned"
-    prune_command_parts = extend_cmd_parts(
-        ["torchrun", "--nproc_per_node=1", "prune_minitron.py"],
-        hf_model_name_or_path=fake_hub_checkpoint.hub_model_id,
-        **{f"output_{output}_path": pruned_path},
-        pp_size=1,
-        calib_dataset_name="cnn_dailymail",
-        calib_num_samples=8,
-        seq_length=16,
-        prune_target_params=int(sum(p.numel() for p in teacher_model.parameters()) * 0.8),
-        prune_score_func="mmlu_1pct_bs32",
-        score_lower_bound=0.0,
-        ss_channel_divisor=4,
-        hparams_to_skip="num_attention_heads",
-        top_k=1,
-    )
-    run_example_command(prune_command_parts, example_path="megatron_bridge")
-
-    if output == "megatron":
-        assert (
-            fake_hub_checkpoint.saved_tokenizer_model(pruned_path)
-            == fake_hub_checkpoint.hub_model_id
-        )
-    else:
-        assert (pruned_path / "config.json").exists()
-        assert (pruned_path / fake_hub_checkpoint.marker).exists()
-
-
 @pytest.mark.parametrize(
     "create_teacher",
     [
diff --git a/tests/examples/megatron_bridge/test_quantize_export.py b/tests/examples/megatron_bridge/test_quantize_export.py
index d1bfaae6109..a60fb417acd 100644
--- a/tests/examples/megatron_bridge/test_quantize_export.py
+++ b/tests/examples/megatron_bridge/test_quantize_export.py
@@ -70,16 +70,24 @@ def test_quantize_and_export(
     recipe,
     quantize_parallelism,
     export_parallelism,
+    fake_hub_checkpoint,
 ):
-    """Quantize a tiny model via a YAML recipe and export it to a unified HF checkpoint."""
+    """Quantize a tiny model via a YAML recipe and export it to a unified HF checkpoint.
+
+    The model is passed as a Hub ID (see ``fake_hub_checkpoint``): both steps load and export from
+    its local copy, and the Megatron checkpoint records the Hub ID as its tokenizer.
+    """
     hf_model_path = create_model(tmp_path, with_tokenizer=True, **model_kwargs)
+    model_arg = fake_hub_checkpoint.serve(
+        hf_model_path, "quantize.py", "export_quantized_megatron_to_hf.py", nproc=num_gpus
+    )
     megatron_path = tmp_path / "quantized_megatron"
     hf_export_path = tmp_path / "quantized_hf"
 
     # Step 1: quantize and save a Megatron checkpoint
     quantize_cmd = extend_cmd_parts(
         ["torchrun", f"--nproc_per_node={num_gpus}", "quantize.py", "--skip_generate"],
-        hf_model_name_or_path=hf_model_path,
+        hf_model_name_or_path=model_arg,
         recipe=recipe,
         **{quantize_parallelism: num_gpus},
         calib_dataset_name="cnn_dailymail",
@@ -91,11 +99,16 @@ def test_quantize_and_export(
     run_example_command(quantize_cmd, example_path="megatron_bridge", setup_free_port=True)
     assert (megatron_path / "latest_checkpointed_iteration.txt").exists()
     assert_has_modelopt_state(megatron_path)
+    if fake_hub_checkpoint.active:
+        assert (
+            fake_hub_checkpoint.saved_tokenizer_model(megatron_path)
+            == fake_hub_checkpoint.hub_model_id
+        )
 
     # Step 2: export to HF
     export_cmd = extend_cmd_parts(
         ["torchrun", f"--nproc_per_node={num_gpus}", "export_quantized_megatron_to_hf.py"],
-        hf_model_name_or_path=hf_model_path,
+        hf_model_name_or_path=model_arg,
         megatron_path=megatron_path,
         export_unified_hf_path=hf_export_path,
         **{export_parallelism: num_gpus},
@@ -104,6 +117,8 @@ def test_quantize_and_export(
     assert (hf_export_path / "config.json").exists()
     assert (hf_export_path / "hf_quant_config.json").exists()
     assert_exported_checkpoint_matches(hf_export_path, hf_model_path)
+    if fake_hub_checkpoint.active:
+        assert (hf_export_path / fake_hub_checkpoint.marker).exists()
 
     # The exported unified checkpoint should be loadable and runnable by vLLM. The deployment check below
     # is disabled because it takes too long in CI (likely because of first run)
@@ -119,41 +134,3 @@ def test_quantize_and_export(
     # )
     # outputs = llm.generate(["Hello!"], vllm.SamplingParams(max_tokens=4))
     # assert outputs and outputs[0].outputs and outputs[0].outputs[0].text
-
-
-@pytest.mark.timeout(360)
-def test_quantize_and_export_a_hub_model_id(tmp_path: Path, fake_hub_checkpoint):
-    """A Hub ID is made local once: both steps load and export from that copy, while the
-    Megatron checkpoint records the Hub ID as its tokenizer so it stays valid on other hosts."""
-    fake_hub_checkpoint.serve(
-        create_tiny_qwen3_moe_dir(tmp_path, with_tokenizer=True, **_DENSE_KWARGS),
-        "quantize.py",
-        "export_quantized_megatron_to_hf.py",
-    )
-    megatron_path = tmp_path / "quantized_megatron"
-    hf_export_path = tmp_path / "quantized_hf"
-
-    quantize_cmd = extend_cmd_parts(
-        ["torchrun", "--nproc_per_node=1", "quantize.py", "--skip_generate"],
-        hf_model_name_or_path=fake_hub_checkpoint.hub_model_id,
-        recipe="general/ptq/nvfp4_default-kv_fp8",
-        calib_dataset_name="cnn_dailymail",
-        calib_num_samples=4,
-        calib_batch_size=2,
-        seq_length=16,
-        export_megatron_path=megatron_path,
-    )
-    run_example_command(quantize_cmd, example_path="megatron_bridge", setup_free_port=True)
-    assert (
-        fake_hub_checkpoint.saved_tokenizer_model(megatron_path) == fake_hub_checkpoint.hub_model_id
-    )
-
-    export_cmd = extend_cmd_parts(
-        ["torchrun", "--nproc_per_node=1", "export_quantized_megatron_to_hf.py"],
-        hf_model_name_or_path=fake_hub_checkpoint.hub_model_id,
-        megatron_path=megatron_path,
-        export_unified_hf_path=hf_export_path,
-    )
-    run_example_command(export_cmd, example_path="megatron_bridge", setup_free_port=True)
-    assert (hf_export_path / "hf_quant_config.json").exists()
-    assert (hf_export_path / fake_hub_checkpoint.marker).exists()

From 4f05440509f779df171cac5c1b4b915deb54d6a8 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 20 Oct 2026 00:58:58 +0000
Subject: [PATCH 10/11] Serve the fake Hub ID's tokenizer in the
 Megatron-Bridge tests

Saving a Megatron checkpoint loads the tokenizer it records, here the
Hub ID, which failed for the fake org/tiny-model in CI. Map
AutoTokenizer.from_pretrained(hub_model_id) to the resolved copy, as
the Hugging Face cache does for a real Hub ID.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 tests/examples/megatron_bridge/conftest.py | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/tests/examples/megatron_bridge/conftest.py b/tests/examples/megatron_bridge/conftest.py
index beb86565eb6..0a7157bdc7e 100644
--- a/tests/examples/megatron_bridge/conftest.py
+++ b/tests/examples/megatron_bridge/conftest.py
@@ -18,6 +18,7 @@
 import shutil
 
 import pytest
+import transformers
 import yaml
 from _test_utils.examples.megatron_example_runner import (
     _load_example_module,
@@ -68,7 +69,9 @@ class FakeHubCheckpoint:
     ``checkpoint`` itself, for a script that must read its argument before resolving it) to a copy
     of ``checkpoint`` holding the file ``marker``, and returns ``hub_model_id`` to pass on the
     command line: a step that read it without resolving it fails, and ``marker`` reaching an
-    export shows the export read the resolved copy. Only a single-rank step runs in this process
+    export shows the export read the resolved copy. Like the Hugging Face cache would for a real Hub
+    ID, ``AutoTokenizer.from_pretrained(hub_model_id)`` loads the copy's tokenizer: Megatron-Bridge
+    loads the tokenizer a checkpoint records. Only a single-rank step runs in this process
     and sees the patch, so with ``nproc > 1`` nothing is patched, ``checkpoint`` is returned
     unchanged and ``active`` is False: check the marker and the Hub ID only when it is True.
     """
@@ -95,6 +98,15 @@ def ensure_local_checkpoint(model_name_or_path, group=None):
         for script in scripts:
             module = _load_example_module(script, "megatron_bridge")
             self._monkeypatch.setattr(module, "ensure_local_checkpoint", ensure_local_checkpoint)
+
+        load_tokenizer = transformers.AutoTokenizer.from_pretrained
+
+        def from_pretrained(name, *args, **kwargs):
+            if str(name) == self.hub_model_id:
+                name = str(self._resolved)
+            return load_tokenizer(name, *args, **kwargs)
+
+        self._monkeypatch.setattr(transformers.AutoTokenizer, "from_pretrained", from_pretrained)
         return self.hub_model_id
 
     @staticmethod

From e1fc2c7c8a8a1067d2721745ea4153a668880626 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 20 Oct 2026 23:45:36 +0000
Subject: [PATCH 11/11] Accept the keyword form in the fake Hub ID's tokenizer
 loader

Megatron-LM calls AutoTokenizer.from_pretrained with
pretrained_model_name_or_path= as a keyword, which the wrapper only took
positionally.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 tests/examples/megatron_bridge/conftest.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/tests/examples/megatron_bridge/conftest.py b/tests/examples/megatron_bridge/conftest.py
index 0a7157bdc7e..d8da852f11a 100644
--- a/tests/examples/megatron_bridge/conftest.py
+++ b/tests/examples/megatron_bridge/conftest.py
@@ -101,10 +101,13 @@ def ensure_local_checkpoint(model_name_or_path, group=None):
 
         load_tokenizer = transformers.AutoTokenizer.from_pretrained
 
-        def from_pretrained(name, *args, **kwargs):
-            if str(name) == self.hub_model_id:
-                name = str(self._resolved)
-            return load_tokenizer(name, *args, **kwargs)
+        def from_pretrained(*args, **kwargs):
+            # Megatron-LM passes the name as pretrained_model_name_or_path=.
+            if args and str(args[0]) == self.hub_model_id:
+                args = (str(self._resolved), *args[1:])
+            if str(kwargs.get("pretrained_model_name_or_path")) == self.hub_model_id:
+                kwargs["pretrained_model_name_or_path"] = str(self._resolved)
+            return load_tokenizer(*args, **kwargs)
 
         self._monkeypatch.setattr(transformers.AutoTokenizer, "from_pretrained", from_pretrained)
         return self.hub_model_id