refactor(modelopt): delegate real-quant export to ModelOpt

docs/modelopt/quantization.md

@@ -140,22 +140,21 @@ Use `AutoBridge.export_hf_weights_modelopt()` when you need to stream ModelOpt d
 already-loaded Megatron model instead of writing a full checkpoint through the export script. This is useful for
 integrations that consume Hugging Face weight names directly, such as inference-engine refit paths.
 
-The API currently only supports `quant_mode="nvfp4"`. Quantized parameters are yielded as the original Hugging Face
-`*.weight` name plus the ModelOpt NVFP4 scale tensors:
-
-- `*.weight`
-- `*.weight_scale`
-- `*.weight_scale_2`
-
-Unquantized parameters are yielded under their regular Hugging Face names. Quantizer-internal tensors are skipped.
+Build one export plan after ModelOpt calibration and pass its `quantization_config` unchanged to the consumer. The
+plan retains only stable topology and format metadata, so it can be reused for every weight stream. Mutable quantizer
+state is captured as each weight is streamed. ModelOpt owns the emitted tensor names and formats; unquantized
+parameters retain their regular Hugging Face names.
 
 ```python
 from safetensors.torch import save_file
 
+plan = bridge.build_hf_modelopt_export_plan(model)
+quantization_config = plan.quantization_config
+
 state_dict = {}
 for name, weight in bridge.export_hf_weights_modelopt(
     model,
-    quant_mode="nvfp4",
+    export_plan=plan,
     cpu=True,
     show_progress=False,
 ):
@@ -170,16 +169,20 @@ full `state_dict`.
 ```python
 for name, weight in bridge.export_hf_weights_modelopt(
     model,
-    quant_mode="nvfp4",
-    ignore_patterns=["lm_head", "*self_attn.o_proj*"],
+    export_plan=plan,
     show_progress=False,
 ):
     refit_engine.replace_weight(name, weight)
 ```
 
-`ignore_patterns` are matched against Hugging Face parameter names. The matcher handles the optional `model.` prefix
-and ModelOpt scale suffixes, so a pattern can target the logical parameter name without separately listing
-`*.weight_scale` and `*.weight_scale_2`.
+Plan construction and export-stream iteration perform WORLD and model-parallel collectives. Every distributed rank
+must therefore build the plan and fully consume each export stream in the same task order, without rank-local early
+termination. The current streaming API supports canonical per-expert Hugging Face MoE layouts; canonical grouped-
+expert Hugging Face tensors are rejected until ModelOpt provides a state-stacking operation.
+
+Quantized adapter-wrapped weights are not supported. Fold adapters into the base weights before ModelOpt calibration
+or QAT. Dimension-permuting mappings are also rejected by the streaming API. This API requires a ModelOpt release
+that provides the functional quantized-weight export interface.
 
 ### Supported Models For PTQ
 

pyproject.toml

@@ -134,6 +134,8 @@ override-dependencies = [
     "mlflow>=3.15.1",                                       # To address CVE-2025-15031
     "cachetools>=5.0.0",
     "cryptography>=43.0.0,<47",
+    # TODO: Before merging the functional export integration, pin the first
+    # ModelOpt release containing NVIDIA/Model-Optimizer#2251.
     "nvidia-modelopt==0.46.0rc1",
     "urllib3>=2.6.3",                                       # To address CVE-2026-21441
     "langchain>=0.3.28",                                    # To address CVE-2025-65106

src/megatron/bridge/models/conversion/auto_bridge.py

@@ -31,6 +31,7 @@
 
 
 if TYPE_CHECKING:
+    from megatron.bridge.models.conversion.modelopt_utils import ModelOptExportPlan
     from megatron.bridge.peft.base import PEFT
 
 from megatron.core.transformer.module import MegatronModule
@@ -775,58 +776,53 @@ def export_hf_weights(
             weight_dtype=weight_dtype,
         )
 
-    def export_hf_weights_modelopt(
+    def build_hf_modelopt_export_plan(
         self,
         model: MegatronModelT | list[MegatronModelT],
-        quant_mode: str = "nvfp4",
-        cpu: bool = False,
-        show_progress: bool = True,
         conversion_tasks: Optional[List[WeightConversionTask]] = None,
-        ignore_patterns: Optional[List[str]] = None,
-        merge_adapter_weights: bool = True,
-    ) -> Iterable["HFWeightTuple"]:
-        """Export Megatron weights to HuggingFace ModelOpt deployment format.
+    ) -> "ModelOptExportPlan":
+        """Prepare canonical real-quant tensors and configuration for streaming export.
 
         Args:
             model: Megatron model instance or list of instances.
-            quant_mode: ModelOpt quantization mode to export. Currently supports
-                ``"nvfp4"`` and ``"w4a16_nvfp4"``.
-            cpu: Whether to move exported tensors to CPU before yielding.
-            show_progress: Display progress bar during base Hugging Face weight export.
             conversion_tasks: Pre-built conversion tasks. If not provided, tasks will be built
                 automatically from the models.
-            ignore_patterns: Hugging Face parameter name patterns that should remain unquantized.
-                Scale tensor suffixes and the optional ``model.`` prefix are ignored when matching.
-            merge_adapter_weights: Whether to gather and merge LoRA adapter weights into the base
-                tensors during export.
-
-        Yields:
-            HFWeightTuple: Named tuples containing Hugging Face parameter names and tensors. Quantized
-            weights yield the packed weight under the original ``*.weight`` name followed by the
-            corresponding ``*.weight_scale`` and ``*.weight_scale_2`` tensors.
-
-        Raises:
-            RuntimeError: If a matched quantized Megatron parameter uses a qformat unsupported by
-                ``quant_mode``.
         """
         from megatron.bridge.models.conversion.modelopt_utils import build_modelopt_export_plan
 
         if not isinstance(model, list):
             model = [model]
         if conversion_tasks is None:
             conversion_tasks = self._model_bridge.build_conversion_tasks(self.hf_pretrained, model)
-        export_tasks = build_modelopt_export_plan(
+        return build_modelopt_export_plan(
             conversion_tasks,
             model=model,
-            bridge=self._model_bridge,
-            quant_mode=quant_mode,
-            ignore_patterns=ignore_patterns or [],
         )
+
+    def export_hf_weights_modelopt(
+        self,
+        model: MegatronModelT | list[MegatronModelT],
+        cpu: bool = False,
+        show_progress: bool = True,
+        export_plan: Optional["ModelOptExportPlan"] = None,
+        merge_adapter_weights: bool = True,
+    ) -> Iterable["HFWeightTuple"]:
+        """Export canonical ModelOpt deployment tensors from a prepared plan."""
+        if not merge_adapter_weights:
+            raise NotImplementedError("ModelOpt export does not support unmerged adapter weights")
+        if not isinstance(model, list):
+            model = [model]
+        if export_plan is None:
+            export_plan = self.build_hf_modelopt_export_plan(model)
+        from megatron.bridge.models.conversion.modelopt_utils import (
+            prepare_modelopt_export_tasks,
+        )
+
         hf_weights = self.export_hf_weights(
             model,
             cpu=cpu,
             show_progress=show_progress,
-            conversion_tasks=export_tasks,
+            conversion_tasks=prepare_modelopt_export_tasks(export_plan),
             merge_adapter_weights=merge_adapter_weights,
         )
         yield from hf_weights