[1/3] Export expert-parallel Megatron-Bridge MoE checkpoints and fix exporter / QAD gaps

CHANGELOG.rst

@@ -36,6 +36,7 @@ Changelog
 
 - Add an end-to-end W4A4 NVFP4 PTQ and QAD tutorial for Qwen3.6-35B-A3B also covering evaluation and vLLM throughput benchmarking. See `examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/>`_ for details.
 - Add ``--mlflow <tracking-uri>`` to the ``examples/megatron_bridge`` scripts that write a checkpoint -- ``prune_minitron.py``, ``quantize.py``, ``distill.py``, ``export_quantized_megatron_to_hf.py`` and ``export_distilled_megatron_to_hf.py`` (MLflow's own ``MLFLOW_TRACKING_URI`` is honoured too). Each run records the invocation, its arguments as searchable params and its log -- a distillation records its training metrics instead of the rank-0 log -- and writes ``.experiment.json`` into the checkpoint it produced, so a pruning, a quantization, the distillation that refines its checkpoint and the export that deploys it can be traced to one another; uploading the checkpoints themselves stays off unless ``--mlflow_log_checkpoints`` is passed.
+- Add ``--ep_size`` to ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` so large MoE models with grouped-GEMM experts can be exported with their experts sharded across GPUs. Checkpoints built with ``--no_moe_grouped_gemm`` must still be exported at ``--ep_size 1``.
 
 *Misc*
 
@@ -90,6 +91,7 @@ Changelog
 
 **Bug Fixes**
 
+- Fix Megatron unified HF export of MoE models with grouped-GEMM experts when only the experts are quantized (e.g. ``nvfp4_experts_only-*`` recipes): ``hf_quant_config.json`` and the ``quantization_config`` in ``config.json`` were not written, so the quantized experts were served as unquantized weights. Re-export such checkpoints.
 - Fix Megatron-Core checkpoint saving for quantized grouped MoE experts when tensor and expert parallelism are both enabled.
 - Fix shared ONNX export metadata and Diffusers attention policy: every ``NVFP4QuantExporter`` post-process now upgrades the default-domain opset to at least 23, all FP8 custom-op exports re-run ONNX shape/type inference after setting output metadata, and quantized SDPA derives FP8 MHA enablement from the live Q/K/V quantizers instead of honoring a caller-set ``_disable_fp8_mha`` attribute.
 - Fix ONNX FP16 conversion failing to preserve public output types when type inference changes a graph output declaration before output casts are inserted.

examples/megatron_bridge/export_quantized_megatron_to_hf.py

@@ -22,7 +22,8 @@
 
 The HuggingFace unified exporter does not gather tensor-parallel-sharded weights, so this script
 always loads the checkpoint at tensor_model_parallel_size=1 (re-sharding from whatever TP was used
-during quantization). Use --pp_size to shard a large model across GPUs for export.
+during quantization). Use --pp_size and/or --ep_size to shard a large model across GPUs for
+export; --ep_size needs grouped-GEMM experts (not --no_moe_grouped_gemm).
 
 Example usage to export an FP8 checkpoint produced by quantize.py:
 
@@ -110,8 +111,14 @@ def get_args() -> argparse.Namespace:
         help="Export extra modules such as Medusa heads, EAGLE, or MTP.",
     )
 
-    # Only Pipeline parallelism is supported for export
+    # Only pipeline and expert parallelism are supported for export
     parser.add_argument("--pp_size", type=int, default=1, help="Pipeline parallel size")
+    parser.add_argument(
+        "--ep_size",
+        type=int,
+        default=1,
+        help="Expert parallel size (grouped-GEMM experts only); shards the experts across ranks",
+    )
     parser.add_argument(
         "--num_layers_in_first_pipeline_stage",
         type=int,
@@ -156,7 +163,7 @@ def main(args: argparse.Namespace):
         provider_overrides={
             "tensor_model_parallel_size": 1,  # Tensor parallelism is not supported
             "pipeline_model_parallel_size": args.pp_size,
-            "expert_model_parallel_size": 1,  # Expert parallelism is not supported
+            "expert_model_parallel_size": args.ep_size,
             "expert_tensor_parallel_size": 1,  # Expert tensor parallelism is not supported
             "num_layers_in_first_pipeline_stage": args.num_layers_in_first_pipeline_stage,
             "num_layers_in_last_pipeline_stage": args.num_layers_in_last_pipeline_stage,

modelopt/torch/distill/plugins/megatron.py

@@ -590,6 +590,11 @@ def _sharded_state_dict(self, *args, **kwargs) -> "ShardedStateDict":
     # An MTP head left out of quantization is exempt from that: there is no quantization
     # error to recover there, and its CE materialises an fp32 [seq, vocab] tensor.
     skip_mtp_loss = _mtp_excluded_from_quantization(model)
+    if skip_mtp_loss:
+        # Freeze the untrained MTP head: DDP's overlapped grad reduce asserts on params with no grad.
+        for name, param in model.named_parameters():
+            if "mtp" in name.split("."):
+                param.requires_grad_(False)
 
     def _compute_student_lm_loss(self, labels, logits) -> Tensor:
         self._lm_loss_call_count += 1

modelopt/torch/export/quant_utils.py

@@ -452,12 +452,6 @@ def uses_iq_quantization(module) -> bool:
 
     This reads ``num_bits`` directly rather than resolving each layer's full format, so an
     unrelated unsupported quantizer elsewhere in the model cannot turn the check into an error.
-
-    Known gap, shared with ``get_quantization_format``: ``weight_attr_names`` yields nothing for
-    a TEGroupedLinear, whose parameters are ``weight0..N`` while its quantizer is a single
-    ``GroupedQuantizer`` under ``weight_quantizer``. Neither function sees such a module, so an
-    experts-only IQ model reports no format at all -- not just here. Closing it belongs in
-    ``weight_attr_names``, where it affects every format, rather than in this helper.
     """
     for weight_name in weight_attr_names(module):
         weight_quantizer = representative_weight_quantizer(module, weight_name)

modelopt/torch/export/unified_export_megatron.py

@@ -177,7 +177,9 @@ def __init__(
         # Update hf_config
         self._hf_text_config.num_hidden_layers = language_model.config.num_layers
         self._hf_text_config.hidden_size = language_model.config.hidden_size
-        self._hf_text_config.head_dim = language_model.config.kv_channels
+        # MLA's kv_channels is the V head dim, not HF's head_dim (e.g. glm5_next derives it from RoPE).
+        if not getattr(language_model.config, "multi_latent_attention", False):
+            self._hf_text_config.head_dim = language_model.config.kv_channels
         self._hf_text_config.num_attention_heads = language_model.config.num_attention_heads
         self._hf_text_config.num_key_value_heads = language_model.config.num_query_groups
         self.is_multimodal = isinstance(model, LLaVAModel)
@@ -201,7 +203,10 @@ def __init__(
                 del self._hf_config.quantization_config
         self.all_rules = self._populate_rule_book()
         self.rules = self.all_rules[self.arch]
-        self.exclude_modules = []
+        # The vision tower is copied through unquantized, so deployments must not treat it as such.
+        self.exclude_modules = [
+            prefix.removesuffix(".") + "*" for prefix in self.vision_passthrough_prefixes or ()
+        ]
         self.layer_config_dict = {}
 
         if not hasattr(model, "_modelopt_state"):
@@ -316,6 +321,11 @@ def save_pretrained(
         # medusa_heads and eagle_module only exist in the last stage.
         is_last_stage_main_rank = pp_rank == pp_size - 1 and tp_rank == 0
         is_writer_rank = self._is_sidecar_writer_rank(is_last_stage_main_rank)
+        # One writer per pipeline stage: other TP / DP / EP ranks hold the same layers (EP>1 ranks
+        # hold no gathered experts at all), and writing them too would race on the same files.
+        writes_layers = (
+            tp_rank == 0 and get_data_parallel_rank() == 0 and get_expert_model_parallel_rank() == 0
+        )
 
         quantization_format = self._get_quantization_format(self.model)
         if self._any_rank_uses_iq_quantization():
@@ -369,7 +379,11 @@ def save_pretrained(
                         self._hf_pretrained_model_name,
                         trust_remote_code=self.trust_remote_code,
                     )
-                    generation_config.save_pretrained(save_directory)
+                    # 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.
@@ -394,8 +408,9 @@ def save_pretrained(
                 except (OSError, ValueError, ImportError):
                     pass
 
-            # MTP load mutates per-rank layer_state_dicts, so it runs on every last-stage main rank.
-            mtp_state_dict = self._get_mtp_state_dict()
+            # 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.
+            mtp_state_dict = self._get_mtp_state_dict(copy_from_pretrained=writes_layers)
             if len(mtp_state_dict) > 0:
                 layer_state_dicts[self.model.config.num_layers].update(mtp_state_dict)
                 print(f"Successfully loaded {len(mtp_state_dict)} MTP tensors")
@@ -433,7 +448,7 @@ def save_pretrained(
         # Add multimodal components to state_dict. Since only support decoder model quantization,
         # no changes will be made to the multimodal components. We copy the multimodal components
         # from the pretrained model directly to the state_dict to avoid implementing the export logic.
-        if is_first_stage_main_rank:
+        if is_first_stage_main_rank and writes_layers:
             # layer_state_dicts is keyed by layer_number (1-indexed), so the first
             # decoder layer on this (first) PP stage is the smallest key, not 0.
             # Merge the multimodal components into that shard so they land in a file
@@ -461,9 +476,8 @@ def save_pretrained(
                 json.dump(config_dict, f, indent=4)
         torch.distributed.barrier()
 
-        # save_safetensors(state_dict, save_directory)
         save_safetensors_by_layer_index(
-            layer_state_dicts=layer_state_dicts,
+            layer_state_dicts=layer_state_dicts if writes_layers else {},
             total_layers=self.model.config.num_layers,
             save_directory=save_directory,
             name_template="model-{:05d}-of-{:05d}",
@@ -733,6 +747,15 @@ def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False):
                             layer.mlp.shared_experts.gate_weight, layer_id, is_mtp=is_mtp
                         )
                 if hasattr(layer.mlp.experts, "local_experts"):
+                    # SequentialMLP rules index experts by local position, so EP>1 would collide.
+                    if (
+                        torch.distributed.is_initialized()
+                        and get_expert_model_parallel_world_size() > 1
+                    ):
+                        raise NotImplementedError(
+                            "Export at expert parallel size > 1 needs grouped-GEMM experts; "
+                            "export SequentialMLP (--no_moe_grouped_gemm) checkpoints at EP=1."
+                        )
                     if not self.rules.get("use_packed_local_experts", False):
                         for expert_id, expert in enumerate(layer.mlp.experts.local_experts):
                             self.rules["local_experts.linear_fc1"](
@@ -772,12 +795,12 @@ def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False):
                 self.rules["linear_fc1"](layer.mlp.linear_fc1, layer_id, is_mtp=is_mtp)
                 self.rules["linear_fc2"](layer.mlp.linear_fc2, layer_id, is_mtp=is_mtp)
 
-    def _get_mtp_state_dict(self) -> dict[str, torch.Tensor]:
-        """Export the live MTP module, or copy it from the pretrained model if absent."""
+    def _get_mtp_state_dict(self, copy_from_pretrained: bool = True) -> dict[str, torch.Tensor]:
+        """Export the live MTP module, or copy it from the pretrained model if absent (and allowed)."""
         model = getattr(self, "model", None)
         mtp = getattr(model, "mtp", None)
         if mtp is None or not hasattr(mtp, "layers") or len(mtp.layers) == 0:
-            return self._copy_mtp_state_dict_from_pretrained()
+            return self._copy_mtp_state_dict_from_pretrained() if copy_from_pretrained else {}
 
         # Inner layers reuse the base walker with is_mtp=True (retargets backbone -> mtp).
         saved_state_dict = self._state_dict
@@ -1089,7 +1112,8 @@ def _get_weight_bias(
             and module.expert_bias is not None
             and module.expert_bias.numel() > 0
         ):
-            name_to_value["expert_bias"] = module.expert_bias.to(dtype).cpu()
+            # FP32 like Megatron's buffer and HF's e_score_correction_bias: it decides expert routing.
+            name_to_value["expert_bias"] = module.expert_bias.float().cpu()
 
         return name_to_value
 
@@ -1585,12 +1609,18 @@ def _grouped_mlp_slicing(
             torch.save(local_expert_state, _buf)
             local_bytes = _buf.getvalue()
             del _buf
-            gathered_bytes: list = [None] * ep_size
-            torch.distributed.all_gather_object(
-                gathered_bytes, local_bytes, group=get_expert_model_parallel_group()
+            # Gather to EP rank 0 only, which writes the shards: holding every expert on every EP
+            # rank multiplies host memory by EP (a full GLM-5.3-Flash export OOMs at EP4).
+            ep_group = get_expert_model_parallel_group()
+            gathered_bytes: list | None = [None] * ep_size if ep_rank == 0 else None
+            torch.distributed.gather_object(
+                local_bytes,
+                gathered_bytes,
+                dst=torch.distributed.get_global_rank(ep_group, 0),
+                group=ep_group,
             )
             del local_bytes
-            for b in gathered_bytes:
+            for b in gathered_bytes or ():
                 # weights_only=False: our own torch.save output from a sibling EP rank
                 # in this job's collective, not user-supplied.
                 s_loaded = torch.load(io.BytesIO(b), map_location="cpu", weights_only=False)

modelopt/torch/quantization/utils/core_utils.py

@@ -247,18 +247,24 @@ def representative_weight_quantizer(module: nn.Module, weight_name: str = "weigh
 def weight_attr_names(module: nn.Module) -> "Generator[str, None, None]":
     """Get the weight param attribute names in a converted module, non-recursive.
 
-    Covers three layouts:
+    Covers four layouts:
 
     - standard ``nn.Linear``: ``weight`` + ``weight_quantizer``.
     - custom per-weight quantizer (e.g. ``Llama4TextExperts`` with ``gate_up_proj`` +
       ``gate_up_proj_weight_quantizer``).
     - fused-experts ``nn.ModuleList`` quantizers (``_QuantFusedExperts`` with
       ``<first_proj>`` + ``<first_proj>_weight_quantizers`` plural list).
+    - TEGroupedLinear: ``weight0..N`` behind one ``GroupedQuantizer``, reported as ``weight``.
     """
+    from ..nn import GroupedQuantizer  # local: ..nn imports this utils package at module scope
+
     # standard: "weight" + "weight_quantizer" (singular) or "weight_quantizers" (plural)
     if getattr(module, "weight", None) is not None:
         if representative_weight_quantizer(module, "weight") is not None:
             yield "weight"
+    elif isinstance(getattr(module, "weight_quantizer", None), GroupedQuantizer):
+        # TEGroupedLinear: ``weight0..N`` share one GroupedQuantizer and there is no ``weight``.
+        yield "weight"
 
     # per-parameter custom attr names
     for name, _ in module.named_parameters(recurse=False):

tests/examples/megatron_bridge/test_quantize_export.py

@@ -41,30 +41,47 @@
 
 
 @pytest.mark.parametrize(
-    ("create_model", "model_kwargs"),
+    ("create_model", "model_kwargs", "recipe", "quantize_parallelism", "export_parallelism"),
     [
-        # MoE: routed experts used to be dropped silently from the export.
-        (create_tiny_qwen3_moe_dir, _DENSE_KWARGS),
-        # Dense VLM: only the language model is quantized, vision is copied through.
-        (create_tiny_qwen3vl_dir, {}),
-        # Mamba hybrid + MoE, and the one architecture that keeps grouped-GEMM experts.
-        (create_tiny_nemotron_h_dir, {}),
+        # MoE: routed experts used to be dropped silently from the export. Calibrates with context
+        # parallelism (not on the VLM: M-RoPE CP needs nemo:26.10).
+        (
+            create_tiny_qwen3_moe_dir,
+            _DENSE_KWARGS,
+            "general/ptq/nvfp4_default-kv_fp8",
+            "cp_size",
+            "pp_size",
+        ),
+        # Dense VLM: only the language model is quantized, vision is copied through. Keeps the
+        # FP8 script-to-checkpoint coverage.
+        (create_tiny_qwen3vl_dir, {}, "general/ptq/fp8_default-kv_fp8", "tp_size", "pp_size"),
+        # Mamba hybrid + MoE with grouped-GEMM experts: exported with its experts sharded
+        # across ranks (EP), resharded from the TP-quantized checkpoint.
+        (create_tiny_nemotron_h_dir, {}, "general/ptq/nvfp4_default-kv_fp8", "tp_size", "ep_size"),
     ],
     ids=["qwen3_moe", "qwen3vl", "nemotron_h"],
 )
 @pytest.mark.timeout(360)  # quantize + export in one test; 1-gpu CI exceeds the default 300s
-def test_quantize_and_export(tmp_path: Path, num_gpus, create_model, model_kwargs):
+def test_quantize_and_export(
+    tmp_path: Path,
+    num_gpus,
+    create_model,
+    model_kwargs,
+    recipe,
+    quantize_parallelism,
+    export_parallelism,
+):
     """Quantize a tiny model via a YAML recipe and export it to a unified HF checkpoint."""
     hf_model_path = create_model(tmp_path, with_tokenizer=True, **model_kwargs)
-    megatron_path = tmp_path / "fp8_megatron"
-    hf_export_path = tmp_path / "fp8_hf"
+    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,
-        recipe="general/ptq/fp8_default-kv_fp8",
-        tp_size=num_gpus,
+        recipe=recipe,
+        **{quantize_parallelism: num_gpus},
         calib_dataset_name="cnn_dailymail",
         calib_num_samples=4,
         calib_batch_size=2,
@@ -81,7 +98,7 @@ def test_quantize_and_export(tmp_path: Path, num_gpus, create_model, model_kwarg
         hf_model_name_or_path=hf_model_path,
         megatron_path=megatron_path,
         export_unified_hf_path=hf_export_path,
-        pp_size=num_gpus,
+        **{export_parallelism: num_gpus},
     )
     run_example_command(export_cmd, example_path="megatron_bridge", setup_free_port=True)
     assert (hf_export_path / "config.json").exists()