fix(model): default MODEL_CONFIG_CLASS to None and opt in per family (#6305)

Signed-off-by: Eugenio "Jay" Zuccarelli <11176606+jayzuccarelli@users.noreply.github.com>

Author: jayzuccarelli

docs/bridge-guide.md

@@ -115,6 +115,12 @@ family's Megatron Core `ModelBuilder`, such as `GPTModelBuilder` or
 while other model families migrate. Calling the legacy provider API for a
 builder-backed family emits a deprecation warning.
 
+The builder path is opt-in per family: a bridge opts in by setting
+`MODEL_CONFIG_CLASS`, which defaults to `None`. For a family that has not
+migrated, `get_model_config()` raises `ModelConfigNotSupportedError` and
+`to_megatron_provider()` remains the supported path, without a deprecation
+warning.
+
 ```python
 from megatron.bridge import AutoBridge
 

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

@@ -60,7 +60,6 @@ class BagelBridge(MegatronModelBridge):
     """Import official BAGEL tensors with the strict native checkpoint mapper."""
 
     SUPPORTS_HF_PRETRAINED_EXPORT = False
-    MODEL_CONFIG_CLASS = None
 
     def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> BagelModelProvider:
         """Build the BAGEL provider from the nested official configuration."""

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

@@ -1929,8 +1929,8 @@ def get_model_config(self) -> ModelConfig:
             Serializable model config linked to its model builder.
 
         Raises:
-            ModelConfigNotSupportedError: If the selected model family explicitly
-                disables the builder path.
+            ModelConfigNotSupportedError: If the selected model family has not
+                migrated to builder-backed configs.
         """
         hf_config = (
             self.hf_pretrained.config

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

@@ -89,7 +89,6 @@
     unwrap_model,
 )
 from megatron.bridge.models.decorators.dispatch import dispatch
-from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig
 from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM
 from megatron.bridge.models.model_provider import ModelProviderMixin
 from megatron.bridge.models.transformer_config import TransformerConfig as BridgeTransformerConfig
@@ -529,16 +528,18 @@ def mapping_registry(self) -> MegatronMappingRegistry:
     # provider, such as one implementing MLA.
     PROVIDER_CLASS = None
 
-    # Builder-backed construction is rolling out incrementally by model family. Override this
-    # when the standard GPT model config cannot represent the model, and set it to None for
-    # families that have not yet migrated. The selected model config owns its nested
-    # transformer config class.
-    MODEL_CONFIG_CLASS: ClassVar[type[ModelConfig] | None] = BridgeGPTModelConfig
+    # Builder-backed construction is rolling out incrementally by model family, so it is opt-in.
+    # None means the family has not migrated: get_model_config() raises and the provider path
+    # stays the supported one. A family that has mapped its Hugging Face config for the builder
+    # path sets this to BridgeGPTModelConfig, or to a specialized config when the standard GPT
+    # model config cannot represent the model. Subclasses of a migrated bridge inherit the
+    # opt-in. The selected model config owns its nested transformer config class.
+    MODEL_CONFIG_CLASS: ClassVar[type[ModelConfig] | None] = None
 
     # Conversion still uses the provider-backed construction path unless a model family has
     # explicitly migrated its conversion lifecycle to ModelBuilder. Keeping this separate
-    # from MODEL_CONFIG_CLASS preserves get_model_config() for legacy families without
-    # silently changing how their checkpoints are constructed.
+    # from MODEL_CONFIG_CLASS lets a family expose get_model_config() without silently
+    # changing how its checkpoints are constructed. Requires MODEL_CONFIG_CLASS to be set.
     USE_MODEL_CONFIG_FOR_CONVERSION: ClassVar[bool] = False
 
     # Leave unset unless HF export must copy nonstandard files in addition to the usual artifacts,
@@ -817,9 +818,9 @@ def _partition_model_config_kwargs(
     def hf_config_to_model_config(self, hf_config: PretrainedConfig) -> ModelConfig:
         """Convert a Hugging Face config directly to a builder-backed config.
 
-        GPT-style model families use ``BridgeGPTModelConfig`` by default.
-        A bridge may select a specialized config class or explicitly disable
-        this path by setting ``MODEL_CONFIG_CLASS`` to ``None``.
+        This path is opt-in per model family. ``MODEL_CONFIG_CLASS`` defaults to
+        ``None``; a migrated family sets it to ``BridgeGPTModelConfig`` or to a
+        specialized config class.
 
         Args:
             hf_config: Hugging Face architecture configuration.
@@ -828,14 +829,16 @@ def hf_config_to_model_config(self, hf_config: PretrainedConfig) -> ModelConfig:
             Serializable model configuration linked to its builder.
 
         Raises:
-            ModelConfigNotSupportedError: If the model family explicitly disables this path.
+            ModelConfigNotSupportedError: If the model family has not migrated to
+                builder-backed configs (``MODEL_CONFIG_CLASS`` is ``None``).
             ValueError: If mapped fields do not belong to either config dataclass.
         """
         model_config_class = self.MODEL_CONFIG_CLASS
         if model_config_class is None:
             raise ModelConfigNotSupportedError(
                 f"ModelConfig conversion is not implemented for {type(self).__name__}. "
-                "This model family sets MODEL_CONFIG_CLASS to None."
+                "This model family has not migrated to builder-backed configs "
+                "(MODEL_CONFIG_CLASS is None); use AutoBridge.to_megatron_provider() instead."
             )
 
         config_kwargs = self.hf_config_to_model_config_kwargs(hf_config)

src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py

@@ -60,10 +60,6 @@ class GLM5Bridge(MegatronModelBridge):
         >>> provider = bridge.to_megatron_provider()
     """
 
-    # Until the declarative Hybrid builder supports the GLM adapters, fail clearly
-    # rather than returning the inherited GPT builder configuration.
-    MODEL_CONFIG_CLASS = None
-
     @classmethod
     def megatron_to_hf_config(cls, provider: GLM5ModelProvider) -> dict:
         """Export HF block counts and preserve the reference model's context length."""

src/megatron/bridge/models/llama/llama_bridge.py

@@ -27,6 +27,7 @@
     GatedMLPMapping,
     QKVMapping,
 )
+from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig
 from megatron.bridge.models.gpt_provider import GPTModelProvider
 
 
@@ -46,6 +47,8 @@ class LlamaBridge(MegatronModelBridge):
         >>> model_config = bridge.get_model_config()
     """
 
+    MODEL_CONFIG_CLASS = BridgeGPTModelConfig
+
     def hf_config_to_model_config_kwargs(self, hf_config: Any) -> dict[str, Any]:
         """Convert a Hugging Face Llama config to builder config kwargs."""
         config_kwargs = super().hf_config_to_model_config_kwargs(hf_config)

src/megatron/bridge/models/qwen/qwen3_bridge.py

@@ -25,6 +25,7 @@
     GatedMLPMapping,
     QKVMapping,
 )
+from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig
 
 
 @MegatronModelBridge.register_bridge(source=Qwen3ForCausalLM, target=GPTModel, model_type="qwen3")
@@ -42,6 +43,7 @@ class Qwen3Bridge(MegatronModelBridge):
         >>> model_config = bridge.get_model_config()
     """
 
+    MODEL_CONFIG_CLASS = BridgeGPTModelConfig
     USE_MODEL_CONFIG_FOR_CONVERSION = True
 
     def hf_config_to_model_config_kwargs(self, hf_config: PretrainedConfig) -> dict[str, Any]:

tests/unit_tests/models/test_gpt_model_config.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import inspect
+import warnings
 from dataclasses import dataclass
 from unittest.mock import Mock, patch
 
@@ -21,12 +22,17 @@
 import torch.nn.functional as F
 from megatron.core.models.gpt import GPTModel
 from megatron.core.transformer import ModuleSpec
+from transformers import Qwen2Config
 
+import megatron.bridge.models  # noqa: F401  (registers every model-family bridge)
 from megatron.bridge.models.common.base import ModelConfig
 from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge, ModelConfigNotSupportedError
 from megatron.bridge.models.gpt.gpt_builder import GPTModelBuilder
 from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig
 from megatron.bridge.models.llama.llama_bridge import LlamaBridge
+from megatron.bridge.models.muse_glimmer.muse_glimmer_bridge import MuseGlimmerBridge
+from megatron.bridge.models.qwen.qwen2_bridge import Qwen2Bridge
+from megatron.bridge.models.qwen.qwen3_bridge import Qwen3Bridge
 from megatron.bridge.models.transformer_config import TransformerConfig
 
 
@@ -137,22 +143,88 @@ class SpecializedBridge(LlamaBridge):
     assert model_config.specialized_field is True
 
 
-def test_gpt_model_config_is_the_bridge_default() -> None:
-    assert MegatronModelBridge.MODEL_CONFIG_CLASS is BridgeGPTModelConfig
-    assert LlamaBridge.MODEL_CONFIG_CLASS is BridgeGPTModelConfig
-    assert "MODEL_CONFIG_CLASS" not in LlamaBridge.__dict__
+def _bridge_subclasses(bridge_class: type[MegatronModelBridge]) -> set[type[MegatronModelBridge]]:
+    subclasses = set(bridge_class.__subclasses__())
+    return subclasses.union(*(_bridge_subclasses(subclass) for subclass in subclasses))
+
+
+def _all_bridge_classes() -> list[type[MegatronModelBridge]]:
+    """Return every bridge class registered by importing ``megatron.bridge.models``."""
+    return sorted(
+        (
+            bridge_class
+            for bridge_class in _bridge_subclasses(MegatronModelBridge)
+            if bridge_class.__module__.startswith("megatron.bridge.")
+        ),
+        key=lambda bridge_class: bridge_class.__name__,
+    )
+
+
+def _has_model_config_mapping(bridge_class: type[MegatronModelBridge]) -> bool:
+    return (
+        bridge_class.hf_config_to_model_config_kwargs is not MegatronModelBridge.hf_config_to_model_config_kwargs
+        or bridge_class.hf_config_to_model_config is not MegatronModelBridge.hf_config_to_model_config
+    )
+
+
+def test_model_config_is_opt_in_per_family() -> None:
+    assert MegatronModelBridge.MODEL_CONFIG_CLASS is None
     assert MegatronModelBridge.USE_MODEL_CONFIG_FOR_CONVERSION is False
+    for bridge_class in (LlamaBridge, Qwen3Bridge):
+        assert bridge_class.__dict__["MODEL_CONFIG_CLASS"] is BridgeGPTModelConfig
+    assert MuseGlimmerBridge.__dict__["MODEL_CONFIG_CLASS"] is not None
     assert LlamaBridge.USE_MODEL_CONFIG_FOR_CONVERSION is False
     assert BridgeGPTModelConfig.transformer_config_class is TransformerConfig
     assert not hasattr(MegatronModelBridge, "TRANSFORMER_CONFIG_CLASS")
 
 
-def test_bridge_can_explicitly_disable_model_config() -> None:
-    class UnsupportedBridge(LlamaBridge):
-        MODEL_CONFIG_CLASS = None
+def test_bridge_subclass_inherits_model_config_opt_in() -> None:
+    class DerivedBridge(LlamaBridge):
+        pass
+
+    assert DerivedBridge.MODEL_CONFIG_CLASS is BridgeGPTModelConfig
+
+
+@pytest.mark.parametrize("bridge_class", _all_bridge_classes(), ids=lambda bridge_class: bridge_class.__name__)
+def test_bridge_can_explicitly_disable_model_config(bridge_class: type[MegatronModelBridge]) -> None:
+    if bridge_class.MODEL_CONFIG_CLASS is None:
+        # Unmigrated family: the builder path is rejected instead of returning generic defaults.
+        assert not _has_model_config_mapping(bridge_class)
+        assert bridge_class.USE_MODEL_CONFIG_FOR_CONVERSION is False
+        with pytest.raises(ModelConfigNotSupportedError, match=f"{bridge_class.__name__}.*to_megatron_provider"):
+            bridge_class().hf_config_to_model_config(object())
+    else:
+        # Opted-in family: it must map its Hugging Face config for the builder path.
+        assert _has_model_config_mapping(bridge_class)
+
+    # A bridge that replaces hf_config_to_model_config() wholesale does not read MODEL_CONFIG_CLASS.
+    if bridge_class.hf_config_to_model_config is MegatronModelBridge.hf_config_to_model_config:
+
+        class UnsupportedBridge(bridge_class):
+            MODEL_CONFIG_CLASS = None
+
+        with pytest.raises(ModelConfigNotSupportedError, match="UnsupportedBridge.*to_megatron_provider"):
+            UnsupportedBridge().hf_config_to_model_config(object())
+
+
+def test_unmigrated_bridge_keeps_provider_without_warning() -> None:
+    hf_config = Qwen2Config(
+        num_hidden_layers=2,
+        hidden_size=64,
+        intermediate_size=128,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        vocab_size=256,
+    )
+
+    assert Qwen2Bridge.MODEL_CONFIG_CLASS is None
+    with warnings.catch_warnings():
+        warnings.simplefilter("error", FutureWarning)
+        provider = Qwen2Bridge().provider_bridge(Mock(config=hf_config))
 
-    with pytest.raises(ModelConfigNotSupportedError, match="sets MODEL_CONFIG_CLASS to None"):
-        UnsupportedBridge().hf_config_to_model_config(object())
+    assert provider.position_embedding_type == "rope"
+    assert provider.add_qkv_bias is True
+    assert provider.add_bias_linear is False
 
 
 @pytest.mark.skipif(