Address review: distilled-export layout, hub-id self-check, NVFP4 scale merge

export_distilled_megatron_to_hf.py chose its MoE layout with `use_moe_grouped_gemm`,
but it now rejects quantized checkpoints, so every checkpoint it handles was written by
distill.py's unquantized branch using grouped GEMM. For an architecture with no
`experts.linear_fc1` rule the two disagreed, and the script would have built
SequentialMLP over a grouped-GEMM checkpoint. It mirrors distill.py's choice now.

`_verify_exported_keys` bailed out whenever the source was not a local directory, which
is the documented invocation -- every README and PR snippet passes a hub repo id. The
guard meant to protect user runs on architectures CI never sees was therefore inactive
for exactly those runs. It now fetches just `*.safetensors.index.json`, and still skips
when the source cannot be reached.

The NVFP4 merge kept each expert's block scales but replaced `weight_scale_2` with the
cross-expert maximum, inflating every quieter expert's effective scale by
`scale_2_max / scale_2_i`. Values round-tripped, which is why the reference check passed,
but the FP4 range was wasted. Block scales are now rescaled onto the merged global scale;
measured on the 256-expert half-depth Qwen3.5, worst-case relative error against the
reference drops from 0.16667 to 0.15331. Both packed paths share one helper so they
cannot drift.

Grouped-expert packing only forwarded the four suffixes it names, silently dropping
anything else (`bias` on a future `add_bias_linear` MoE); leftovers now assert.
`load_modelopt_megatron_checkpoint` reads the checkpoint keys already, so it now checks
the stored expert layout against the built model instead of trusting the config.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>

Author: kevalmorabia97

examples/megatron_bridge/export_distilled_megatron_to_hf.py

@@ -69,7 +69,6 @@
     is_vlm_config,
     load_mbridge_model_from_hf,
     load_modelopt_megatron_checkpoint,
-    use_moe_grouped_gemm,
 )
 
 # Megatron-Bridge checkpoint iteration directories use names like ``iter_0000100``.
@@ -256,11 +255,8 @@ def main(args: argparse.Namespace):
         _bridge, _provider, _model, full_model, _tokenizer = load_mbridge_model_from_hf(
             hf_model_name_or_path=args.student_hf_path,
             trust_remote_code=args.trust_remote_code,
-            moe_grouped_gemm=use_moe_grouped_gemm(
-                args.student_hf_path,
-                trust_remote_code=args.trust_remote_code,
-                force_sequential=args.no_moe_grouped_gemm,
-            ),
+            # Mirrors distill.py's unquantized branch
+            moe_grouped_gemm=not args.no_moe_grouped_gemm,
             provider_overrides={
                 "tensor_model_parallel_size": args.tp_size,
                 "pipeline_model_parallel_size": args.pp_size,

modelopt/torch/export/unified_export_megatron.py

@@ -29,7 +29,7 @@
 
 import torch
 import torch.distributed
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, snapshot_download
 from huggingface_hub.errors import EntryNotFoundError
 from safetensors import safe_open
 from safetensors.torch import save_file
@@ -50,6 +50,7 @@
     QUANTIZATION_W4A16_NVFP4,
 )
 from .plugins.hf_checkpoint_utils import (
+    _is_hf_hub_offline,
     copy_hf_ckpt_remote_code,
     copy_non_safetensor_files_from_ckpt,
     load_multimodal_components,
@@ -445,16 +446,25 @@ def save_pretrained(
 
     def _verify_exported_keys(self, save_directory, pretrained_model_name_or_path) -> None:
         """Raise if the export dropped tensors the source has: a missing rule emits nothing."""
-        if pretrained_model_name_or_path is None or not os.path.isdir(
-            str(pretrained_model_name_or_path)
-        ):
-            return  # hub id: not worth a download inside export
+        if pretrained_model_name_or_path is None:
+            return
+        source_dir = str(pretrained_model_name_or_path)
+        if not os.path.isdir(source_dir):
+            # A repo id is the documented invocation, so fetch just the index rather than skip.
+            try:
+                source_dir = snapshot_download(
+                    repo_id=source_dir,
+                    allow_patterns=["*.safetensors.index.json"],
+                    local_files_only=_is_hf_hub_offline(),
+                )
+            except Exception:
+                return  # source unreachable: skip rather than fail an otherwise good export
         index_file = Path(save_directory) / "model.safetensors.index.json"
         if not index_file.exists():
             return
         with open(index_file) as f:
             exported = set(json.load(f)["weight_map"])
-        source = _read_checkpoint_keys(pretrained_model_name_or_path)
+        source = _read_checkpoint_keys(source_dir)
         if not source:
             return
 
@@ -1267,6 +1277,11 @@ def collect(suffix):
         weights = collect(".weight")
         if not weights:
             return
+        handled = (".weight", ".weight_scale", ".weight_scale_2", ".input_scale", ".output_scale")
+        unhandled = {k.split(".", 1)[1] for k in per_expert if not k.endswith(handled)}
+        assert not unhandled, (
+            f"{prefix}: grouped-expert packing has no rule for {sorted(unhandled)}"
+        )
         # Record against the packed prefix, as _pack_name_remapping does for the other packed path.
         if qformat in (None, QUANTIZATION_NONE):
             self._record_excluded_module(prefix)
@@ -1282,9 +1297,8 @@ def collect(suffix):
             self._state_dict[prefix] = merged_weight
         else:
             if scales_2:
-                # NVFP4 keeps each expert's block scales; only the global scale is merged.
-                merged_scale = torch.stack(scales, dim=0)
-                merged_scale_2 = torch.max(torch.stack(scales_2, dim=0), dim=0)[0]
+                # NVFP4 keeps each expert's block scales, rescaled onto the merged global scale.
+                merged_scale, merged_scale_2 = self._merge_nvfp4_expert_scales(scales, scales_2)
             else:
                 merged_scale, merged_scale_2 = torch.max(torch.stack(scales, dim=0), dim=0)[0], None
             self._state_dict[prefix] = to_quantized_weight(
@@ -1771,6 +1785,20 @@ def _self_attention_scaling(
                 # FP8 KV Cache is supported in VLLM; NVFP4 supported in TRTLLM
                 self.kv_cache_dtype = kv_cache_dtype
 
+    @staticmethod
+    def _merge_nvfp4_expert_scales(scales: list, scales_2: list):
+        """Merge per-expert NVFP4 scales onto one global scale, preserving each expert's FP4 range.
+
+        Each expert's block scales were derived against its own ``scale_2``; rescaling them by
+        ``scale_2_i / scale_2_max`` keeps the quieter experts from losing a mantissa bit.
+        """
+        merged_scale_2 = torch.max(torch.stack(scales_2, dim=0), dim=0)[0]
+        stacked_2 = torch.stack(scales_2, dim=0).reshape(-1, *([1] * scales[0].dim()))
+        merged_scale = (
+            torch.stack(scales, dim=0).to(torch.float32) * (stacked_2 / merged_scale_2)
+        ).to(scales[0].dtype)
+        return merged_scale, merged_scale_2
+
     def _pack_name_remapping(self, module, prefix, layer_type=None, is_mtp=False, transpose=True):
         """Pack per-expert weights into one tensor; ``transpose`` for HF [E, in, out] layouts."""
         if is_mtp:
@@ -1812,8 +1840,9 @@ def _pack_name_remapping(self, module, prefix, layer_type=None, is_mtp=False, tr
                 merged_weight_scale = None
         else:
             # NVFP4
-            merged_weight_scale_2 = torch.max(torch.stack(weight_scale_2_list, dim=0), dim=0)[0]
-            merged_weight_scale = torch.stack(weight_scale_list, dim=0)
+            merged_weight_scale, merged_weight_scale_2 = self._merge_nvfp4_expert_scales(
+                weight_scale_list, weight_scale_2_list
+            )
             if transpose:
                 merged_weight_scale = merged_weight_scale.transpose(-2, -1).contiguous()
 

modelopt/torch/utils/plugins/mbridge.py

@@ -230,6 +230,22 @@ def load_modelopt_megatron_checkpoint(
         print_rank_0("Language-model-only checkpoint: loading into the VLM's `.language_model`.")
         model = [get_language_model(m)[0] for m in unwrapped_model]
 
+    # The layout is baked into the checkpoint, so a model built with the other one silently
+    # mismatches. The keys are already read above, so check rather than trust the config.
+    ckpt_grouped = any(".experts.weight0" in key for key in checkpoint_keys)
+    ckpt_sequential = any(".local_experts." in key for key in checkpoint_keys)
+    if ckpt_grouped or ckpt_sequential:
+        built_sequential = any(
+            ".local_experts." in name for m in unwrap_model(model) for name, _ in m.named_modules()
+        )
+        if ckpt_sequential != built_sequential:
+            raise ValueError(
+                f"{megatron_path} stores MoE experts as "
+                f"{'SequentialMLP' if ckpt_sequential else 'grouped GEMM (TEGroupedMLP)'} but the "
+                f"model was built as {'SequentialMLP' if built_sequential else 'grouped GEMM'}. "
+                "Pass the same --no_moe_grouped_gemm setting used to write the checkpoint."
+            )
+
     # Restore the ModelOpt state before loading weights.
     # has_modelopt_state / load_modelopt_state resolves the latest iter_* directory
     if restore_modelopt_state: