Skip to content

Commit d039149

Browse files
committed
fix(model): default MODEL_CONFIG_CLASS to None and opt in per family
Make the builder-backed config path opt-in. MegatronModelBridge.MODEL_CONFIG_CLASS now defaults to None, and the migrated families (Llama, Qwen3, Muse Glimmer) set it explicitly. Unmigrated families reject get_model_config() through the existing MODEL_CONFIG_CLASS is None branch, and provider_bridge() only emits the deprecation FutureWarning for opted-in families. This replaces the _has_model_config_mapping() hook-override guard. The explicit MODEL_CONFIG_CLASS = None on GLM5Bridge and BagelBridge is now the default and is removed. test_bridge_can_explicitly_disable_model_config is parametrized over every registered bridge. Signed-off-by: Eugenio "Jay" Zuccarelli <11176606+jayzuccarelli@users.noreply.github.com>
1 parent 3815dc6 commit d039149

8 files changed

Lines changed: 87 additions & 61 deletions

File tree

‎docs/bridge-guide.md‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,12 @@ family's Megatron Core `ModelBuilder`, such as `GPTModelBuilder` or
115115
while other model families migrate. Calling the legacy provider API for a
116116
builder-backed family emits a deprecation warning.
117117

118+
The builder path is opt-in per family: a bridge opts in by setting
119+
`MODEL_CONFIG_CLASS`, which defaults to `None`. For a family that has not
120+
migrated, `get_model_config()` raises `ModelConfigNotSupportedError` and
121+
`to_megatron_provider()` remains the supported path, without a deprecation
122+
warning.
123+
118124
```python
119125
from megatron.bridge import AutoBridge
120126

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,6 @@ class BagelBridge(MegatronModelBridge):
6060
"""Import official BAGEL tensors with the strict native checkpoint mapper."""
6161

6262
SUPPORTS_HF_PRETRAINED_EXPORT = False
63-
MODEL_CONFIG_CLASS = None
6463

6564
def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> BagelModelProvider:
6665
"""Build the BAGEL provider from the nested official configuration."""

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1934,8 +1934,8 @@ def get_model_config(self) -> ModelConfig:
19341934
Serializable model config linked to its model builder.
19351935
19361936
Raises:
1937-
ModelConfigNotSupportedError: If the selected model family explicitly
1938-
disables the builder path.
1937+
ModelConfigNotSupportedError: If the selected model family has not
1938+
migrated to builder-backed configs.
19391939
"""
19401940
hf_config = (
19411941
self.hf_pretrained.config

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

Lines changed: 17 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,6 @@
8989
unwrap_model,
9090
)
9191
from megatron.bridge.models.decorators.dispatch import dispatch
92-
from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig
9392
from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM
9493
from megatron.bridge.models.model_provider import ModelProviderMixin
9594
from megatron.bridge.models.transformer_config import TransformerConfig as BridgeTransformerConfig
@@ -529,16 +528,18 @@ def mapping_registry(self) -> MegatronMappingRegistry:
529528
# provider, such as one implementing MLA.
530529
PROVIDER_CLASS = None
531530

532-
# Builder-backed construction is rolling out incrementally by model family. Override this
533-
# when the standard GPT model config cannot represent the model, and set it to None for
534-
# families that have not yet migrated. The selected model config owns its nested
535-
# transformer config class.
536-
MODEL_CONFIG_CLASS: ClassVar[type[ModelConfig] | None] = BridgeGPTModelConfig
531+
# Builder-backed construction is rolling out incrementally by model family, so it is opt-in.
532+
# None means the family has not migrated: get_model_config() raises and the provider path
533+
# stays the supported one. A family that has mapped its Hugging Face config for the builder
534+
# path sets this to BridgeGPTModelConfig, or to a specialized config when the standard GPT
535+
# model config cannot represent the model. Subclasses of a migrated bridge inherit the
536+
# opt-in. The selected model config owns its nested transformer config class.
537+
MODEL_CONFIG_CLASS: ClassVar[type[ModelConfig] | None] = None
537538

538539
# Conversion still uses the provider-backed construction path unless a model family has
539540
# explicitly migrated its conversion lifecycle to ModelBuilder. Keeping this separate
540-
# from MODEL_CONFIG_CLASS preserves get_model_config() for legacy families without
541-
# silently changing how their checkpoints are constructed.
541+
# from MODEL_CONFIG_CLASS lets a family expose get_model_config() without silently
542+
# changing how its checkpoints are constructed. Requires MODEL_CONFIG_CLASS to be set.
542543
USE_MODEL_CONFIG_FOR_CONVERSION: ClassVar[bool] = False
543544

544545
# Leave unset unless HF export must copy nonstandard files in addition to the usual artifacts,
@@ -814,26 +815,12 @@ def _partition_model_config_kwargs(
814815
}
815816
return model_kwargs, transformer_kwargs
816817

817-
def _has_model_config_mapping(self) -> bool:
818-
"""Return whether this bridge maps its Hugging Face config to a builder-backed config.
819-
820-
The base ``hf_config_to_model_config_kwargs()`` maps only the generic
821-
``CONFIG_MAPPING`` fields. Family-specific settings such as position
822-
embeddings, normalization, or bias flags stay in ``provider_bridge()``
823-
until a family overrides one of the builder-backed hooks.
824-
"""
825-
bridge_class = type(self)
826-
return (
827-
bridge_class.hf_config_to_model_config_kwargs is not MegatronModelBridge.hf_config_to_model_config_kwargs
828-
or bridge_class.hf_config_to_model_config is not MegatronModelBridge.hf_config_to_model_config
829-
)
830-
831818
def hf_config_to_model_config(self, hf_config: PretrainedConfig) -> ModelConfig:
832819
"""Convert a Hugging Face config directly to a builder-backed config.
833820
834-
GPT-style model families use ``BridgeGPTModelConfig`` by default.
835-
A bridge may select a specialized config class or explicitly disable
836-
this path by setting ``MODEL_CONFIG_CLASS`` to ``None``.
821+
This path is opt-in per model family. ``MODEL_CONFIG_CLASS`` defaults to
822+
``None``; a migrated family sets it to ``BridgeGPTModelConfig`` or to a
823+
specialized config class.
837824
838825
Args:
839826
hf_config: Hugging Face architecture configuration.
@@ -842,21 +829,16 @@ def hf_config_to_model_config(self, hf_config: PretrainedConfig) -> ModelConfig:
842829
Serializable model configuration linked to its builder.
843830
844831
Raises:
845-
ModelConfigNotSupportedError: If the model family explicitly disables this path
846-
or has not migrated its config mapping.
832+
ModelConfigNotSupportedError: If the model family has not migrated to
833+
builder-backed configs (``MODEL_CONFIG_CLASS`` is ``None``).
847834
ValueError: If mapped fields do not belong to either config dataclass.
848835
"""
849836
model_config_class = self.MODEL_CONFIG_CLASS
850837
if model_config_class is None:
851838
raise ModelConfigNotSupportedError(
852839
f"ModelConfig conversion is not implemented for {type(self).__name__}. "
853-
"This model family sets MODEL_CONFIG_CLASS to None."
854-
)
855-
if not self._has_model_config_mapping():
856-
raise ModelConfigNotSupportedError(
857-
f"ModelConfig conversion is not implemented for {type(self).__name__}. "
858-
"This model family has not migrated its Hugging Face config mapping; "
859-
"use AutoBridge.to_megatron_provider() instead."
840+
"This model family has not migrated to builder-backed configs "
841+
"(MODEL_CONFIG_CLASS is None); use AutoBridge.to_megatron_provider() instead."
860842
)
861843

862844
config_kwargs = self.hf_config_to_model_config_kwargs(hf_config)
@@ -912,7 +894,7 @@ def provider_bridge(self, hf_pretrained: HFPreTrained) -> ModelProviderTarget:
912894
from megatron.bridge.models.gpt_provider import GPTModelProvider
913895
from megatron.bridge.models.mla_provider import MLAModelProvider
914896

915-
if self.MODEL_CONFIG_CLASS is not None and self._has_model_config_mapping():
897+
if self.MODEL_CONFIG_CLASS is not None:
916898
warnings.warn(
917899
f"{type(self).__name__}.provider_bridge() and AutoBridge.to_megatron_provider() "
918900
"are deprecated for builder-backed models. Use AutoBridge.get_model_config() "

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

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -60,10 +60,6 @@ class GLM5Bridge(MegatronModelBridge):
6060
>>> provider = bridge.to_megatron_provider()
6161
"""
6262

63-
# Until the declarative Hybrid builder supports the GLM adapters, fail clearly
64-
# rather than returning the inherited GPT builder configuration.
65-
MODEL_CONFIG_CLASS = None
66-
6763
@classmethod
6864
def megatron_to_hf_config(cls, provider: GLM5ModelProvider) -> dict:
6965
"""Export HF block counts and preserve the reference model's context length."""

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
GatedMLPMapping,
2828
QKVMapping,
2929
)
30+
from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig
3031
from megatron.bridge.models.gpt_provider import GPTModelProvider
3132

3233

@@ -46,6 +47,8 @@ class LlamaBridge(MegatronModelBridge):
4647
>>> model_config = bridge.get_model_config()
4748
"""
4849

50+
MODEL_CONFIG_CLASS = BridgeGPTModelConfig
51+
4952
def hf_config_to_model_config_kwargs(self, hf_config: Any) -> dict[str, Any]:
5053
"""Convert a Hugging Face Llama config to builder config kwargs."""
5154
config_kwargs = super().hf_config_to_model_config_kwargs(hf_config)

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
GatedMLPMapping,
2626
QKVMapping,
2727
)
28+
from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig
2829

2930

3031
@MegatronModelBridge.register_bridge(source=Qwen3ForCausalLM, target=GPTModel, model_type="qwen3")
@@ -42,6 +43,7 @@ class Qwen3Bridge(MegatronModelBridge):
4243
>>> model_config = bridge.get_model_config()
4344
"""
4445

46+
MODEL_CONFIG_CLASS = BridgeGPTModelConfig
4547
USE_MODEL_CONFIG_FOR_CONVERSION = True
4648

4749
def hf_config_to_model_config_kwargs(self, hf_config: PretrainedConfig) -> dict[str, Any]:

‎tests/unit_tests/models/test_gpt_model_config.py‎

Lines changed: 57 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -22,13 +22,16 @@
2222
import torch.nn.functional as F
2323
from megatron.core.models.gpt import GPTModel
2424
from megatron.core.transformer import ModuleSpec
25-
from transformers import Qwen3Config
25+
from transformers import Qwen2Config
2626

27+
import megatron.bridge.models # noqa: F401 (registers every model-family bridge)
2728
from megatron.bridge.models.common.base import ModelConfig
2829
from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge, ModelConfigNotSupportedError
2930
from megatron.bridge.models.gpt.gpt_builder import GPTModelBuilder
3031
from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig
3132
from megatron.bridge.models.llama.llama_bridge import LlamaBridge
33+
from megatron.bridge.models.muse_glimmer.muse_glimmer_bridge import MuseGlimmerBridge
34+
from megatron.bridge.models.qwen.qwen2_bridge import Qwen2Bridge
3235
from megatron.bridge.models.qwen.qwen3_bridge import Qwen3Bridge
3336
from megatron.bridge.models.transformer_config import TransformerConfig
3437

@@ -140,48 +143,83 @@ class SpecializedBridge(LlamaBridge):
140143
assert model_config.specialized_field is True
141144

142145

143-
def test_gpt_model_config_is_the_bridge_default() -> None:
144-
assert MegatronModelBridge.MODEL_CONFIG_CLASS is BridgeGPTModelConfig
145-
assert LlamaBridge.MODEL_CONFIG_CLASS is BridgeGPTModelConfig
146-
assert "MODEL_CONFIG_CLASS" not in LlamaBridge.__dict__
146+
def _bridge_subclasses(bridge_class: type[MegatronModelBridge]) -> set[type[MegatronModelBridge]]:
147+
subclasses = set(bridge_class.__subclasses__())
148+
return subclasses.union(*(_bridge_subclasses(subclass) for subclass in subclasses))
149+
150+
151+
def _all_bridge_classes() -> list[type[MegatronModelBridge]]:
152+
"""Return every bridge class registered by importing ``megatron.bridge.models``."""
153+
return sorted(
154+
(
155+
bridge_class
156+
for bridge_class in _bridge_subclasses(MegatronModelBridge)
157+
if bridge_class.__module__.startswith("megatron.bridge.")
158+
),
159+
key=lambda bridge_class: bridge_class.__name__,
160+
)
161+
162+
163+
def _has_model_config_mapping(bridge_class: type[MegatronModelBridge]) -> bool:
164+
return (
165+
bridge_class.hf_config_to_model_config_kwargs is not MegatronModelBridge.hf_config_to_model_config_kwargs
166+
or bridge_class.hf_config_to_model_config is not MegatronModelBridge.hf_config_to_model_config
167+
)
168+
169+
170+
def test_model_config_is_opt_in_per_family() -> None:
171+
assert MegatronModelBridge.MODEL_CONFIG_CLASS is None
147172
assert MegatronModelBridge.USE_MODEL_CONFIG_FOR_CONVERSION is False
173+
for bridge_class in (LlamaBridge, Qwen3Bridge):
174+
assert bridge_class.__dict__["MODEL_CONFIG_CLASS"] is BridgeGPTModelConfig
175+
assert MuseGlimmerBridge.__dict__["MODEL_CONFIG_CLASS"] is not None
148176
assert LlamaBridge.USE_MODEL_CONFIG_FOR_CONVERSION is False
149177
assert BridgeGPTModelConfig.transformer_config_class is TransformerConfig
150178
assert not hasattr(MegatronModelBridge, "TRANSFORMER_CONFIG_CLASS")
151179

152180

153-
def test_bridge_can_explicitly_disable_model_config() -> None:
181+
def test_bridge_subclass_inherits_or_disables_model_config_opt_in() -> None:
182+
class DerivedBridge(LlamaBridge):
183+
pass
184+
154185
class UnsupportedBridge(LlamaBridge):
155186
MODEL_CONFIG_CLASS = None
156187

157-
with pytest.raises(ModelConfigNotSupportedError, match="sets MODEL_CONFIG_CLASS to None"):
188+
assert DerivedBridge.MODEL_CONFIG_CLASS is BridgeGPTModelConfig
189+
with pytest.raises(ModelConfigNotSupportedError, match="UnsupportedBridge.*to_megatron_provider"):
158190
UnsupportedBridge().hf_config_to_model_config(object())
159191

160192

161-
def _make_qwen3_config() -> Qwen3Config:
162-
return Qwen3Config(
193+
@pytest.mark.parametrize("bridge_class", _all_bridge_classes(), ids=lambda bridge_class: bridge_class.__name__)
194+
def test_bridge_can_explicitly_disable_model_config(bridge_class: type[MegatronModelBridge]) -> None:
195+
if bridge_class.MODEL_CONFIG_CLASS is None:
196+
# Unmigrated family: the builder path is rejected instead of returning generic defaults.
197+
assert not _has_model_config_mapping(bridge_class)
198+
assert bridge_class.USE_MODEL_CONFIG_FOR_CONVERSION is False
199+
with pytest.raises(ModelConfigNotSupportedError, match=f"{bridge_class.__name__}.*to_megatron_provider"):
200+
bridge_class().hf_config_to_model_config(object())
201+
else:
202+
# Opted-in family: it must map its Hugging Face config for the builder path.
203+
assert _has_model_config_mapping(bridge_class)
204+
205+
206+
def test_unmigrated_bridge_keeps_provider_without_warning() -> None:
207+
hf_config = Qwen2Config(
163208
num_hidden_layers=2,
164209
hidden_size=64,
165210
intermediate_size=128,
166211
num_attention_heads=4,
167212
num_key_value_heads=2,
168-
head_dim=16,
169213
vocab_size=256,
170214
)
171215

172-
173-
def test_bridge_without_model_config_mapping_rejects_model_config() -> None:
174-
with pytest.raises(ModelConfigNotSupportedError, match="Qwen3Bridge.*to_megatron_provider"):
175-
Qwen3Bridge().hf_config_to_model_config(_make_qwen3_config())
176-
177-
178-
def test_bridge_without_model_config_mapping_keeps_provider_without_warning() -> None:
216+
assert Qwen2Bridge.MODEL_CONFIG_CLASS is None
179217
with warnings.catch_warnings():
180218
warnings.simplefilter("error", FutureWarning)
181-
provider = Qwen3Bridge().provider_bridge(Mock(config=_make_qwen3_config()))
219+
provider = Qwen2Bridge().provider_bridge(Mock(config=hf_config))
182220

183221
assert provider.position_embedding_type == "rope"
184-
assert provider.qk_layernorm is True
222+
assert provider.add_qkv_bias is True
185223
assert provider.add_bias_linear is False
186224

187225

0 commit comments

Comments
 (0)