fix(modelopt): exclude quantizer state from real-quant task discovery

Signed-off-by: Meng Xin <164952785+mxinO@users.noreply.github.com>

Author: mxinO

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

@@ -888,7 +888,12 @@ def build_hf_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)
+            from modelopt.torch.quantization.nn import AnyQuantizer
+
+            # ModelOpt reads quantizer state from each owning weight, not separate HF mappings.
+            conversion_tasks = self._model_bridge.build_conversion_tasks(
+                self.hf_pretrained, model, exclude_module_types=AnyQuantizer
+            )
         return build_modelopt_export_plan(
             conversion_tasks,
             model=model,

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

@@ -107,6 +107,25 @@
 _BridgeImplClass = TypeVar("_BridgeImplClass", bound="MegatronModelBridge")
 
 
+def _named_conversion_parameters(
+    model: torch.nn.Module,
+    exclude_module_types: tuple[type[torch.nn.Module], ...] = (),
+) -> Iterable[tuple[str, torch.Tensor]]:
+    """Iterate parameters and persistent buffers outside excluded module subtrees."""
+    parameters = itertools.chain(model.named_parameters(), persistent_buffers(model))
+    if not exclude_module_types:
+        yield from parameters
+        return
+
+    excluded_names = {name for name, module in model.named_modules() if isinstance(module, exclude_module_types)}
+    for name, tensor in parameters:
+        parent_name = name.rpartition(".")[0]
+        while parent_name and parent_name not in excluded_names:
+            parent_name = parent_name.rpartition(".")[0]
+        if parent_name not in excluded_names:
+            yield name, tensor
+
+
 class ModelConfigNotSupportedError(NotImplementedError):
     """Raised when a bridge has no builder-backed model config path."""
 
@@ -1045,11 +1064,14 @@ def mapping_registry(self):
         raise NotImplementedError("Subclass must implement mapping_registry method")
 
     def _megatron_global_param_names_all_pp_ranks(
-        self, megatron_model: Union[MegatronModel, List[MegatronModel]]
+        self,
+        megatron_model: Union[MegatronModel, List[MegatronModel]],
+        *,
+        exclude_module_types: tuple[type[torch.nn.Module], ...] = (),
     ) -> List[str]:
         """Get all parameter names across all pipeline parallel ranks."""
-        # Cache the result after first call
-        if hasattr(self, "_cached_param_names"):
+        # Filtered discovery must neither reuse nor replace ordinary conversion's cache.
+        if not exclude_module_types and hasattr(self, "_cached_param_names"):
             return self._cached_param_names
 
         # Compute the result
@@ -1061,8 +1083,7 @@ def _megatron_global_param_names_all_pp_ranks(
         models_list = megatron_model if isinstance(megatron_model, list) else [megatron_model]
 
         for vp_stage, model in enumerate(models_list):
-            # persistent buffers are part of the model's state_dict, but not the named_parameters, so we must include them here separately
-            for local_param_name, _ in itertools.chain(model.named_parameters(), persistent_buffers(model)):
+            for local_param_name, _ in _named_conversion_parameters(model, exclude_module_types):
                 if "_extra_state" in local_param_name:
                     continue
                 local_param_name = self._unwrap_name(local_param_name)
@@ -1085,10 +1106,10 @@ def _megatron_global_param_names_all_pp_ranks(
         # change this might cause a hang
         gathered_global_param_names = sorted(flattened_names, key=extract_sort_key)
 
-        # Cache the result
-        self._cached_param_names = gathered_global_param_names
+        if not exclude_module_types:
+            self._cached_param_names = gathered_global_param_names
 
-        return self._cached_param_names
+        return gathered_global_param_names
 
     def _with_progress_tracking(self, tasks, description: str, show_progress: bool = True):
         """Helper method to wrap an iterable with progress tracking.
@@ -2233,11 +2254,15 @@ def build_conversion_tasks(
         hf_pretrained: HFPreTrained,
         megatron_model: List[MegatronModel],
         weight_dtype: Optional[torch.dtype] = None,
+        *,
+        exclude_module_types: tuple[type[torch.nn.Module], ...] = (),
     ) -> List[WeightConversionTask]:
         """Construct the conversion tasks between HF and megatron.
 
         Args:
             weight_dtype: Export dtype recorded on each task. Overrides must forward it.
+            exclude_module_types: Module types whose parameters and persistent buffers
+                are excluded from discovery. Remaining parameters still require mappings.
 
         The algorithm walks over every parameter of every destination model,
         asks the :class:`MegatronMappingRegistry` whether it has a mapping for that
@@ -2263,7 +2288,10 @@ def build_conversion_tasks(
         model_config = unwrapped_model.config
         embeddings_are_tied = self._share_embeddings_and_output_weights(model_config)
         pp_rank = _get_pp_rank(megatron_model)
-        sorted_global_param_names_all_pp_ranks = self._megatron_global_param_names_all_pp_ranks(megatron_model)
+        sorted_global_param_names_all_pp_ranks = self._megatron_global_param_names_all_pp_ranks(
+            megatron_model,
+            **({"exclude_module_types": exclude_module_types} if exclude_module_types else {}),
+        )
 
         # Filter out the output_layer weight if embeddings are tied to it -- it doesn't exist
         # as a separate parameter in that case. Other `output_layer.*` parameters (e.g. a
@@ -2284,8 +2312,7 @@ def build_conversion_tasks(
 
         pending_tasks: list[WeightConversionTask | None] = [None] * len(sorted_global_param_names_all_pp_ranks)
         for vp_stage, model in enumerate(megatron_model):
-            # persistent buffers are part of the model's state_dict, but not the named_parameters, so we must include them here separately
-            for local_name, _ in itertools.chain(model.named_parameters(), persistent_buffers(model)):
+            for local_name, _ in _named_conversion_parameters(model, exclude_module_types):
                 if "_extra_state" in local_name or self._is_adapter_param_name(local_name):
                     continue
 

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

@@ -47,7 +47,7 @@ class ModelOptExportPlan(NamedTuple):
     conversion_tasks: list[WeightConversionTask]
     quantization_config: dict[str, Any]
     quantized_params: frozenset[str]
-    local_quantization_specs: dict[str, tuple[object | None, bool]]
+    local_quantization_specs: dict[str, object | None]
 
 
 @dataclass(frozen=True)
@@ -57,7 +57,6 @@ class _SourceState:
     parallelism: str | None = None
     qkv_layout: tuple[int, int, int, int, bool] | None = None
     weight_dtype: torch.dtype | None = None
-    input_quantizer_enabled: bool = False
 
 
 def _same_storage(left: object, right: object) -> bool:
@@ -163,22 +162,16 @@ def _capture_source_spec(task: WeightConversionTask) -> _SourceState | None:
         return _SourceState(None, tuple(task.param_weight.shape), weight_dtype=task.param_weight.dtype)
 
     from modelopt.torch.export.quantized_weight_export import get_quantized_weight_export_spec
-    from modelopt.torch.quantization.utils import quantizer_attr_names
 
     spec = get_quantized_weight_export_spec(task.megatron_module, weight_name)
-    input_quantizer = None
     if spec is not None:
         _validate_quantized_task(task)
-        input_quantizer = getattr(task.megatron_module, quantizer_attr_names(weight_name).input_quantizer, None)
-        if input_quantizer is None:
-            input_quantizer = getattr(task.megatron_module, "input_quantizer", None)
     return _SourceState(
         spec,
         tuple(task.param_weight.shape),
         _mapping_parallelism(task.mapping, task.megatron_module) if spec is not None else None,
         _qkv_layout(task.mapping, task.megatron_module) if spec is not None else None,
         task.param_weight.dtype,
-        input_quantizer is not None and input_quantizer.is_enabled,
     )
 
 
@@ -786,7 +779,7 @@ def build_modelopt_export_plan(
         export_tasks,
         quantization_config,
         quantized_params,
-        {name: (source.state, source.input_quantizer_enabled) for name, source in local_specs.items()},
+        {name: source.state for name, source in local_specs.items()},
     )
 
 
@@ -798,11 +791,7 @@ def prepare_modelopt_export_tasks(plan: ModelOptExportPlan) -> list[WeightConver
             if task.global_param_name not in plan.local_quantization_specs:
                 continue
             source = _capture_source_spec(task)
-            if (
-                source is None
-                or (source.state, source.input_quantizer_enabled)
-                != plan.local_quantization_specs[task.global_param_name]
-            ):
+            if source is None or source.state != plan.local_quantization_specs[task.global_param_name]:
                 raise ValueError(f"ModelOpt quantization configuration changed for {task.global_param_name}")
     except Exception as error:
         validation_error = f"{type(error).__name__}: {error}"

src/megatron/bridge/models/kimi/kimi_k3_bridge.py

@@ -345,6 +345,8 @@ def build_conversion_tasks(
         hf_pretrained,
         megatron_model,
         weight_dtype=None,
+        *,
+        exclude_module_types: tuple[type[torch.nn.Module], ...] = (),
     ) -> list[WeightConversionTask]:
         """Expose virtual BF16 expert weights for MXFP4 packed/scale pairs."""
         original_get_all_keys = hf_pretrained.state.source.get_all_keys
@@ -362,7 +364,12 @@ def _get_all_keys_with_virtual() -> list[str]:
 
         hf_pretrained.state.source.get_all_keys = _get_all_keys_with_virtual
         try:
-            return super().build_conversion_tasks(hf_pretrained, megatron_model, weight_dtype=weight_dtype)
+            return super().build_conversion_tasks(
+                hf_pretrained,
+                megatron_model,
+                weight_dtype=weight_dtype,
+                **({"exclude_module_types": exclude_module_types} if exclude_module_types else {}),
+            )
         finally:
             hf_pretrained.state.source.get_all_keys = original_get_all_keys
 

src/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py

@@ -212,6 +212,8 @@ def build_conversion_tasks(
         hf_pretrained,
         megatron_model,
         weight_dtype=None,
+        *,
+        exclude_module_types: tuple[type[torch.nn.Module], ...] = (),
     ) -> List:
         """Override to synthesize virtual weight keys from INT4 quantized triplets.
 
@@ -235,7 +237,12 @@ def _get_all_keys_with_virtual():
 
         hf_pretrained.state.source.get_all_keys = _get_all_keys_with_virtual
         try:
-            return super().build_conversion_tasks(hf_pretrained, megatron_model, weight_dtype=weight_dtype)
+            return super().build_conversion_tasks(
+                hf_pretrained,
+                megatron_model,
+                weight_dtype=weight_dtype,
+                **({"exclude_module_types": exclude_module_types} if exclude_module_types else {}),
+            )
         finally:
             hf_pretrained.state.source.get_all_keys = original_get_all_keys
 

tests/unit_tests/models/test_model_bridge.py

@@ -1650,3 +1650,56 @@ def __init__(self, megatron_param, hf_param):
     assert len(tasks) == 1
     assert tasks[0].global_param_name == "decoder.weight"
     assert tasks[0].megatron_module is None
+
+
+def test_conversion_discovery_excludes_complete_module_subtrees():
+    class Excluded(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.ones(1))
+            self.register_buffer("amax", torch.ones(1))
+            self.child = torch.nn.Linear(1, 1, bias=False)
+
+    model = torch.nn.Module()
+    model.quantizer = Excluded()
+    model.quantizer_neighbor = torch.nn.Linear(1, 1, bias=False)
+    ordinary = dict(model_bridge_module._named_conversion_parameters(model))
+
+    filtered = dict(model_bridge_module._named_conversion_parameters(model, (Excluded,)))
+
+    assert set(filtered) == {"quantizer_neighbor.weight"}
+    assert "quantizer.amax" in ordinary
+    assert "quantizer.child.weight" in ordinary
+    assert dict(model_bridge_module._named_conversion_parameters(model)).keys() == ordinary.keys()
+    assert list(model_bridge_module._named_conversion_parameters(model.quantizer, (Excluded,))) == []
+
+
+@pytest.mark.parametrize("model_family", ["kimi_k3", "kimi_k25"])
+def test_kimi_conversion_tasks_forward_module_exclusions(monkeypatch, model_family):
+    from megatron.bridge.models.kimi.kimi_k3_bridge import KimiK3Bridge
+    from megatron.bridge.models.kimi_vl.kimi_k25_vl_bridge import KimiK25VLBridge
+
+    bridge_cls = KimiK3Bridge if model_family == "kimi_k3" else KimiK25VLBridge
+    keys = ["expert.weight_packed", "expert.weight_scale", "expert.weight_shape"]
+    get_all_keys = lambda: keys
+    source = SimpleNamespace(get_all_keys=get_all_keys)
+    pretrained = SimpleNamespace(state=SimpleNamespace(source=source))
+    tasks = []
+
+    def build_tasks(_self, hf_pretrained, model, *, weight_dtype, exclude_module_types):
+        assert model == []
+        assert weight_dtype == torch.bfloat16
+        assert exclude_module_types == (torch.nn.ReLU,)
+        assert "expert.weight" in hf_pretrained.state.source.get_all_keys()
+        return tasks
+
+    monkeypatch.setattr(MegatronModelBridge, "build_conversion_tasks", build_tasks)
+    bridge = bridge_cls()
+
+    assert (
+        bridge.build_conversion_tasks(
+            pretrained, [], weight_dtype=torch.bfloat16, exclude_module_types=(torch.nn.ReLU,)
+        )
+        is tasks
+    )
+    assert source.get_all_keys is get_all_keys

tests/unit_tests/models/test_modelopt_utils.py

@@ -18,13 +18,15 @@
 
 import pytest
 import torch
+from transformers import PretrainedConfig
 
 
 mtq = pytest.importorskip("modelopt.torch.quantization")
 weight_export = pytest.importorskip("modelopt.torch.export.quantized_weight_export")
 
 from megatron.bridge.models.conversion import modelopt_utils
 from megatron.bridge.models.conversion.auto_bridge import AutoBridge
+from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
 from megatron.bridge.models.conversion.model_bridge import HFWeightTuple, MegatronModelBridge, WeightConversionTask
 from megatron.bridge.models.conversion.param_mapping import (
     AutoMapping,
@@ -119,6 +121,95 @@ def test_build_plan_delegates_fp8_packing_and_config_to_modelopt():
     assert plan.quantization_config["config_groups"]["group_0"]["targets"] == ["Linear"]
 
 
+@pytest.mark.parametrize("cached_names", [False, True])
+@pytest.mark.parametrize("quantization", ["fp8", "static_nvfp4"])
+def test_public_modelopt_plan_excludes_quantizer_state_before_mapping_validation(
+    monkeypatch, cached_names, quantization
+):
+    monkeypatch.setenv("ENABLE_BRIDGE_QUANT_MAPPING", "1")
+
+    class ProjectionBridge(MegatronModelBridge):
+        def provider_bridge(self, hf_pretrained):
+            raise NotImplementedError
+
+        def mapping_registry(self):
+            return MegatronMappingRegistry(ColumnParallelMapping("projection.weight", "model.projection.weight"))
+
+    implementation = ProjectionBridge()
+    monkeypatch.setattr(implementation, "_megatron_global_adapters_info_all_pp_ranks", lambda _models: [])
+
+    class ProjectionAutoBridge(AutoBridge):
+        @property
+        def _model_bridge(self):
+            return implementation
+
+    model = torch.nn.Module()
+    model.config = SimpleNamespace(share_embeddings_and_output_weights=False)
+    if quantization == "static_nvfp4":
+        model.projection = mtq.quantize(
+            torch.nn.Linear(16, 4, bias=False),
+            {
+                "quant_cfg": [
+                    {"quantizer_name": "*", "enable": False},
+                    {
+                        "quantizer_name": "*weight_quantizer",
+                        "enable": True,
+                        "cfg": {
+                            "num_bits": (2, 1),
+                            "block_sizes": {-1: 16, "type": "static", "scale_bits": (4, 3)},
+                        },
+                    },
+                ],
+                "algorithm": "max",
+            },
+            lambda candidate: candidate(torch.ones(2, 16)),
+        )
+        assert isinstance(model.projection.weight_quantizer, mtq.nn.NVFP4StaticQuantizer)
+        assert model.projection.weight_quantizer.global_amax is not None
+    else:
+        model.projection = _fp8_linear()
+    original_state = dict(model.named_buffers())
+    model_bridge = modelopt_utils.model_bridge_utils
+    monkeypatch.setattr(model_bridge, "_get_pg_collection_from_model", lambda _model: None)
+    monkeypatch.setattr(model_bridge, "_get_pp_group", lambda _model: SimpleNamespace(size=lambda: 1))
+    monkeypatch.setattr(model_bridge, "_get_pp_rank", lambda _model: 0)
+    monkeypatch.setattr(model_bridge, "unwrap_model", lambda candidate: candidate)
+    monkeypatch.setattr(model_bridge, "_megatron_local_name_to_global", lambda _models, _config, name, _vp: name)
+    monkeypatch.setattr(
+        torch.distributed, "all_gather_object", lambda output, value, group: output.__setitem__(0, value)
+    )
+    monkeypatch.setattr(model_bridge.parallel_state, "get_tensor_model_parallel_group", lambda: None)
+    monkeypatch.setattr(model_bridge.parallel_state, "get_expert_tensor_parallel_group", lambda: None)
+    monkeypatch.setattr(model_bridge.parallel_state, "get_expert_model_parallel_group", lambda: None)
+    bridge = ProjectionAutoBridge(PretrainedConfig())
+    if cached_names:
+        implementation._megatron_global_param_names_all_pp_ranks([model])
+
+    plan = bridge.build_hf_modelopt_export_plan(model)
+
+    assert [task.global_param_name for task in plan.conversion_tasks] == ["projection.weight"]
+    assert hasattr(implementation, "_cached_param_names") == cached_names
+    exported = dict(bridge.export_hf_weights_modelopt(model, cpu=True, show_progress=False, export_plan=plan))
+    if quantization == "static_nvfp4":
+        assert exported["model.projection.weight"].dtype == torch.uint8
+        expected_scale = model.projection.weight_quantizer.global_amax / (6 * 448)
+        torch.testing.assert_close(exported["model.projection.weight_scale_2"], expected_scale)
+        assert "model.projection.weight_scale" in exported
+    else:
+        assert exported["model.projection.weight"].dtype == torch.float8_e4m3fn
+    assert dict(model.named_buffers()).keys() == original_state.keys()
+    assert all(dict(model.named_buffers())[name] is value for name, value in original_state.items())
+    if quantization == "static_nvfp4":
+        with pytest.raises(ValueError, match="_global_amax"):
+            implementation.build_conversion_tasks(bridge.hf_pretrained, [model])
+    else:
+        ordinary_tasks = implementation.build_conversion_tasks(bridge.hf_pretrained, [model])
+        assert any(isinstance(task.mapping, AmaxMapping) for task in ordinary_tasks)
+    model.register_buffer("unmapped_buffer", torch.tensor(1.0))
+    with pytest.raises(ValueError, match="unmapped_buffer"):
+        bridge.build_hf_modelopt_export_plan(model)
+
+
 def test_numbered_grouped_weight_uses_exact_modelopt_quantizer():
     module = _GroupedWeights()
     hf_name = "model.layers.0.mlp.experts.0.down_proj.weight"