megatron_bridge: load the HF model from local disk; Megatron export copies its non-model files

CHANGELOG.rst

@@ -21,15 +21,20 @@ 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*
 
 - The ``evaluation`` agent skill can stop looping ("doom loop") reasoning early: ask it to avoid doom loops and it adds vLLM's ``repetition_detection`` request field (vLLM 0.17.0 or later) to every task, ending a request once its output repeats. Use the same setting for baseline and candidate runs.
 
 **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``.

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,
             )

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_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)
     # 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,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_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_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,
@@ -309,7 +320,7 @@ def main(args: argparse.Namespace):
             save_vlm_to_hf(
                 full_model,
                 str(hf_export_path),
-                args.student_hf_path,
+                args.student_local_checkpoint_path,
                 trust_remote_code=args.trust_remote_code,
             )
             record_exported_checkpoint(args, hf_export_path, dist.is_master())
@@ -327,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_path,
+                    student_hf_path=args.student_local_checkpoint_path,
                     template_hf=args.student_hf_model,
                     trust_remote_code=args.trust_remote_code,
                 )

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,22 @@ 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 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_name_or_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_name_or_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_name_or_path,
+            args.local_checkpoint_path,
             trust_remote_code=trust_remote_code,
             force_sequential=args.no_moe_grouped_gemm,
         ),
@@ -214,7 +220,7 @@ def main(args: argparse.Namespace):
     )
     export_fn(
         unwrapped_model,
-        args.hf_model_name_or_path,
+        args.local_checkpoint_path,
         export_extra_modules=export_extra_modules,
         dtype=torch.bfloat16,
         export_dir=args.export_unified_hf_path,

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,15 @@ 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 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_name_or_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
@@ -510,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_name_or_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)
@@ -643,7 +650,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.hub_model_id or args.local_checkpoint_path,
             hf_tokenizer_kwargs=tokenizer_kwargs,
         )
         args.checkpoint_exported = True
@@ -752,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_name_or_path
+                model, args.output_hf_path, source_path=args.local_checkpoint_path
             )
         else:
             dummy_model_cls = AutoModelForImageTextToText if is_vlm else AutoModelForCausalLM
@@ -764,7 +772,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.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")
 

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,23 @@ 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 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_name_or_path
+        trust_remote_code=args.trust_remote_code, hf_path=args.local_checkpoint_path
     )
 
     moe_grouped_gemm = use_moe_grouped_gemm(
-        args.hf_model_name_or_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_name_or_path,
+        hf_model_name_or_path=args.local_checkpoint_path,
         trust_remote_code=trust_remote_code,
         moe_grouped_gemm=moe_grouped_gemm,
         provider_overrides={
@@ -447,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_name_or_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)
@@ -491,7 +497,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.hub_model_id or args.local_checkpoint_path,
         hf_tokenizer_kwargs={"trust_remote_code": trust_remote_code},
     )
     args.checkpoint_exported = True

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.