support modelopt W4A16 & W4A4 NVFP4 export for grouped MoE weights

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

@@ -32,7 +32,6 @@
 
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.transformer_config import MLATransformerConfig, TransformerConfig
-from modelopt.torch.quantization.utils import is_quantized
 from safetensors.torch import save_file
 from transformers.configuration_utils import PretrainedConfig
 from typing_extensions import Unpack
@@ -646,15 +645,16 @@ def export_hf_weights_modelopt(
         quant_mode: str = "nvfp4",
         cpu: bool = False,
         show_progress: bool = True,
-        conversion_tasks: Optional[List[WeightConversionTask | None]] = None,
+        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.
 
         Args:
             model: Megatron model instance or list of instances.
-            quant_mode: ModelOpt quantization mode to export. Currently supports ``"nvfp4"``.
+            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
@@ -673,57 +673,27 @@ def export_hf_weights_modelopt(
             RuntimeError: If a matched quantized Megatron parameter uses a qformat unsupported by
                 ``quant_mode``.
         """
-        from megatron.bridge.models.conversion.modelopt_utils import (
-            build_hf_to_megatron_name_map,
-            collect_modelopt_quant_metadata,
-            get_modelopt_quant_exporter,
-            matches_quant_ignore_pattern,
-            sync_modelopt_quant_metadata,
-        )
-
-        expected_qformat, export_weight = get_modelopt_quant_exporter(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)
-
-        hf_to_megatron_name = build_hf_to_megatron_name_map(conversion_tasks)
-        metadata = collect_modelopt_quant_metadata(conversion_tasks)
-
-        pp_group = model_bridge._get_pp_group(model)
-        if pp_group is not None and dist.is_initialized() and dist.get_world_size(group=pp_group) > 1:
-            sync_modelopt_quant_metadata(metadata, pp_group)
-
+        export_tasks = build_modelopt_export_plan(
+            conversion_tasks,
+            model=model,
+            bridge=self._model_bridge,
+            quant_mode=quant_mode,
+            ignore_patterns=ignore_patterns or [],
+        )
         hf_weights = self.export_hf_weights(
             model,
             cpu=cpu,
             show_progress=show_progress,
-            conversion_tasks=conversion_tasks,
+            conversion_tasks=export_tasks,
             merge_adapter_weights=merge_adapter_weights,
         )
-
-        ignore_patterns = ignore_patterns or []
-        for hf_name, tensor in hf_weights:
-            if "_quantizer." in hf_name:
-                continue
-
-            meta = None
-            if hf_name.endswith(".weight") and not matches_quant_ignore_pattern(hf_name, ignore_patterns):
-                megatron_name = hf_to_megatron_name.get(hf_name)
-                if megatron_name is not None:
-                    meta = metadata.get(megatron_name)
-
-            if meta is None:
-                tensor = tensor.detach()
-                yield HFWeightTuple(hf_name, tensor.cpu() if cpu else tensor)
-                continue
-
-            if meta.qformat != expected_qformat:
-                raise RuntimeError(f"Unsupported qformat for ModelOpt {quant_mode} export: {meta.qformat}")
-
-            for quant_name, quant_tensor in export_weight(hf_name, tensor, meta):
-                yield HFWeightTuple(quant_name, quant_tensor.cpu() if cpu else quant_tensor)
+        yield from hf_weights
 
     def export_hf_weights_quant(
         self,
@@ -1089,6 +1059,10 @@ def save_hf_weights(
         )
         model_instance = self._get_model_instance(model)
         quant_tensors = None
+        # Import lazily so Bridge conversion modules can load before ModelOpt
+        # registers its Megatron-Bridge plugin hooks.
+        from modelopt.torch.quantization.utils import is_quantized
+
         if is_quantized(model_instance):
             quant_tensors = {}
 

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

@@ -14,7 +14,7 @@
 
 import os
 import re
-from typing import List, Optional
+from typing import Any, List, Optional
 
 from megatron.bridge.models.conversion.param_mapping import AutoMapping, MegatronParamMapping
 from megatron.bridge.models.conversion.quant_mapping import convert_to_amax_map, derive_kv_bmm_amax_map
@@ -156,6 +156,7 @@ def __init__(self, *mappings: MegatronParamMapping):
             *mappings: MegatronParamMapping objects
         """
         self.mappings = list(mappings)
+        self._pg_collection = None
         self._add_separate_layernorm_mappings()
         if int(os.environ.get("ENABLE_BRIDGE_QUANT_MAPPING", "0")):
             self._add_quantization_mappings()
@@ -191,6 +192,24 @@ def __init__(self, *mappings: MegatronParamMapping):
                         reverse_dict_patterns[key] = None
                 self._reverse_patterns.append((reverse_dict_patterns, mapping))
 
+    def set_process_groups_from_pg_collection(self, pg_collection: Any) -> None:
+        """Install a process-group collection on all mappings returned by this registry."""
+        self._pg_collection = pg_collection
+        for mapping in self.mappings:
+            mapping.set_process_groups_from_pg_collection(pg_collection)
+
+    def _prepare_mapping(self, mapping: MegatronParamMapping) -> MegatronParamMapping:
+        mapping.set_process_groups_from_pg_collection(self._pg_collection)
+        return mapping
+
+    def resolve_mapping(
+        self,
+        mapping: MegatronParamMapping,
+        captures: tuple[str, ...],
+    ) -> MegatronParamMapping:
+        """Resolve a mapping while preserving this registry's process groups."""
+        return self._prepare_mapping(mapping.resolve(captures))
+
     def megatron_to_hf_lookup(self, megatron_param_name: str) -> Optional[MegatronParamMapping]:
         """
         Get mapping for a Megatron parameter name.
@@ -218,13 +237,13 @@ def megatron_to_hf_lookup(self, megatron_param_name: str) -> Optional[MegatronPa
             if pattern is None:
                 # Direct match
                 if mapping.megatron_param == megatron_param_name:
-                    return mapping
+                    return self._prepare_mapping(mapping)
             else:
                 # Pattern match
                 match = pattern.match(megatron_param_name)
                 if match:
                     # Return resolved mapping with wildcards replaced
-                    return mapping.resolve(match.groups())
+                    return self.resolve_mapping(mapping, match.groups())
         return None
 
     def hf_to_megatron_lookup(self, hf_param_name: str) -> Optional[MegatronParamMapping]:
@@ -247,12 +266,12 @@ def hf_to_megatron_lookup(self, hf_param_name: str) -> Optional[MegatronParamMap
                 if pattern is None:
                     # Direct match
                     if mapping.hf_param == hf_param_name:
-                        return mapping
+                        return self._prepare_mapping(mapping)
                 else:
                     # Pattern match
                     match = pattern.match(hf_param_name)
                     if match:
-                        return mapping.resolve(match.groups())
+                        return self.resolve_mapping(mapping, match.groups())
             else:
                 # Dict destination - check each pattern
                 patterns_dict = pattern_info
@@ -261,12 +280,12 @@ def hf_to_megatron_lookup(self, hf_param_name: str) -> Optional[MegatronParamMap
                         # Direct match
                         if mapping.hf_param[key] == hf_param_name:
                             # Create a simplified mapping for this specific key
-                            return mapping.resolve(())
+                            return self.resolve_mapping(mapping, ())
                     else:
                         # Pattern match
                         match = pattern.match(hf_param_name)
                         if match:
-                            return mapping.resolve(match.groups())
+                            return self.resolve_mapping(mapping, match.groups())
         return None
 
     def get_all_mappings(self) -> List[MegatronParamMapping]: