Skip to content

Commit bfc486f

Browse files
committed
Merge remote-tracking branch 'origin/yuya/muse-glimmer-support' into codex/pr5479-hybrid-model
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com> # Conflicts: # src/megatron/bridge/models/gpt/model_config.py # tests/unit_tests/models/test_gpt_model_config.py
2 parents 9498633 + 743d6ad commit bfc486f

9 files changed

Lines changed: 123 additions & 37 deletions

File tree

‎examples/models/muse_glimmer/README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ bridge = AutoBridge.from_hf_pretrained(
3838
model_config = bridge.get_model_config()
3939

4040
# Apply deployment-specific parallelism before construction.
41-
model_config.transformer.tensor_model_parallel_size = 8
41+
model_config.tensor_model_parallel_size = 8
4242
models = bridge.get_model(model_config)
4343
```
4444

‎src/megatron/bridge/models/common/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
BuildConfigT,
1717
ModelBuilder,
1818
ModelConfig,
19+
ModelConfigOverrideMixin,
1920
ModelT,
2021
Serializable,
2122
compose_hooks,
@@ -28,6 +29,7 @@
2829
"BuildConfigT",
2930
"ModelBuilder",
3031
"ModelConfig",
32+
"ModelConfigOverrideMixin",
3133
"ModelT",
3234
"Serializable",
3335
"compose_hooks",

‎src/megatron/bridge/models/common/base.py‎

Lines changed: 38 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515
from dataclasses import fields as dataclass_fields
1616
from dataclasses import is_dataclass
17-
from typing import Any
17+
from typing import Any, ClassVar
1818

1919
from megatron.training.models.base import (
2020
BuildConfigT, # noqa: F401
@@ -27,6 +27,7 @@
2727
ModelConfig as _MegatronModelConfig,
2828
)
2929

30+
from megatron.bridge.models.transformer_config import TransformerConfig
3031
from megatron.bridge.utils.instantiate_utils import _resolve_target, _validate_target_prefix, instantiate
3132

3233

@@ -73,6 +74,42 @@ def _from_dict(subdata: dict[str, Any], full_key: str) -> Any:
7374
return result
7475

7576

77+
class ModelConfigOverrideMixin:
78+
"""Route flat overrides to their declared owner and reject unknown fields.
79+
80+
Builder-backed model configs store Megatron-Core transformer settings in a
81+
nested ``transformer`` dataclass while exposing flat assignment for recipe
82+
compatibility. This mixin keeps that convenience without allowing typos to
83+
create phantom configuration attributes. Subclasses declare the nested
84+
config type through ``transformer_config_class``.
85+
"""
86+
87+
transformer_config_class: ClassVar[type[TransformerConfig]]
88+
89+
def __setattr__(self, name: str, value: Any, /) -> None:
90+
"""Assign a declared outer or nested field and reject phantom fields."""
91+
try:
92+
transformer = object.__getattribute__(self, "transformer")
93+
except AttributeError:
94+
object.__setattr__(self, name, value)
95+
return
96+
97+
model_fields = getattr(type(self), "__dataclass_fields__", {})
98+
transformer_fields = getattr(type(transformer), "__dataclass_fields__", {})
99+
descriptor = getattr(type(self), name, None)
100+
101+
if name == "transformer" or name in model_fields:
102+
object.__setattr__(self, name, value)
103+
elif name in transformer_fields:
104+
setattr(transformer, name, value)
105+
elif name == "builder" or name.startswith("_") or hasattr(descriptor, "__set__"):
106+
object.__setattr__(self, name, value)
107+
else:
108+
raise AttributeError(
109+
f"Neither {type(self).__name__} nor {type(transformer).__name__} declares a field named {name!r}."
110+
)
111+
112+
76113
class ModelConfig(_MegatronModelConfig):
77114
"""Bridge compatibility wrapper for Megatron-LM model configs."""
78115

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

Lines changed: 19 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
Type,
3838
TypeVar,
3939
Union,
40+
cast,
4041
)
4142

4243
import torch
@@ -51,6 +52,7 @@
5152
from transformers.configuration_utils import PretrainedConfig
5253
from transformers.modeling_utils import PreTrainedModel
5354

55+
from megatron.bridge.models.common import ModelConfigOverrideMixin
5456
from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
5557
from megatron.bridge.models.conversion.param_mapping import (
5658
MegatronParamMapping,
@@ -433,11 +435,11 @@ def mapping_registry(self) -> MegatronMappingRegistry:
433435
# provider, such as one implementing MLA.
434436
PROVIDER_CLASS = None
435437

436-
# Builder-backed construction is rolling out incrementally by model family. Override these
437-
# when the standard GPT or Transformer config cannot represent the model, and set
438-
# MODEL_CONFIG_CLASS to None for families that have not yet migrated.
438+
# Builder-backed construction is rolling out incrementally by model family. Override this
439+
# when the standard GPT model config cannot represent the model, and set it to None for
440+
# families that have not yet migrated. The selected model config owns its nested
441+
# transformer config class.
439442
MODEL_CONFIG_CLASS: ClassVar[type[ModelConfig] | None] = BridgeGPTModelConfig
440-
TRANSFORMER_CONFIG_CLASS: ClassVar[type[BridgeTransformerConfig]] = BridgeTransformerConfig
441443

442444
# Leave unset unless HF export must copy nonstandard files in addition to the usual artifacts,
443445
# for example ``["*reasoning_parser.py"]``.
@@ -723,12 +725,23 @@ def hf_config_to_model_config(self, hf_config: PretrainedConfig) -> ModelConfig:
723725
)
724726

725727
config_kwargs = self.hf_config_to_model_config_kwargs(hf_config)
728+
if not issubclass(model_config_class, ModelConfigOverrideMixin):
729+
raise TypeError(f"{model_config_class.__name__} must inherit {ModelConfigOverrideMixin.__name__}.")
730+
model_config_with_overrides = cast(type[ModelConfigOverrideMixin], model_config_class)
731+
transformer_config_class = model_config_with_overrides.transformer_config_class
732+
if not isinstance(transformer_config_class, type) or not issubclass(
733+
transformer_config_class, BridgeTransformerConfig
734+
):
735+
raise TypeError(
736+
f"{model_config_class.__name__}.transformer_config_class must be a "
737+
f"{BridgeTransformerConfig.__name__} subclass."
738+
)
726739
model_kwargs, transformer_kwargs = self._partition_model_config_kwargs(
727740
config_kwargs,
728741
model_config_class,
729-
self.TRANSFORMER_CONFIG_CLASS,
742+
transformer_config_class,
730743
)
731-
transformer_config = self.TRANSFORMER_CONFIG_CLASS(**transformer_kwargs)
744+
transformer_config = transformer_config_class(**transformer_kwargs)
732745
return model_config_class(transformer=transformer_config, **model_kwargs)
733746

734747
# Set by @register_bridge decorator

‎src/megatron/bridge/models/gpt/model_config.py‎

Lines changed: 5 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -15,16 +15,17 @@
1515
"""Serializable Bridge extension of Megatron-LM's GPT model config."""
1616

1717
from dataclasses import dataclass
18-
from typing import Any
18+
from typing import Any, ClassVar
1919

20-
from megatron.bridge.models.common.base import ModelConfig
20+
from megatron.bridge.models.common.base import ModelConfig, ModelConfigOverrideMixin
2121
from megatron.bridge.models.gpt.gpt_builder import GPTModelConfig
22+
from megatron.bridge.models.transformer_config import TransformerConfig
2223
from megatron.bridge.utils.activation_map import callable_to_str, str_to_callable
2324
from megatron.bridge.utils.instantiate_utils import _resolve_target
2425

2526

2627
@dataclass(kw_only=True)
27-
class BridgeGPTModelConfig(GPTModelConfig):
28+
class BridgeGPTModelConfig(ModelConfigOverrideMixin, GPTModelConfig, ModelConfig):
2829
"""GPT model config with strict overrides and callable serialization.
2930
3031
Outer GPT build fields and nested transformer fields keep one declared
@@ -37,29 +38,7 @@ class BridgeGPTModelConfig(GPTModelConfig):
3738
``run_config.yaml`` without adding model-specific config fields.
3839
"""
3940

40-
def __setattr__(self, name: str, value: Any, /) -> None:
41-
"""Assign a declared outer or nested field and reject phantom fields."""
42-
try:
43-
transformer = object.__getattribute__(self, "transformer")
44-
except AttributeError:
45-
object.__setattr__(self, name, value)
46-
return
47-
48-
model_fields = getattr(type(self), "__dataclass_fields__", {})
49-
transformer_fields = getattr(type(transformer), "__dataclass_fields__", {})
50-
51-
if name == "transformer":
52-
object.__setattr__(self, name, value)
53-
elif name in model_fields:
54-
object.__setattr__(self, name, value)
55-
elif name in transformer_fields:
56-
setattr(transformer, name, value)
57-
elif name == "builder" or name.startswith("_"):
58-
object.__setattr__(self, name, value)
59-
else:
60-
raise AttributeError(
61-
f"Neither {type(self).__name__} nor {type(transformer).__name__} declares a field named {name!r}."
62-
)
41+
transformer_config_class: ClassVar[type[TransformerConfig]] = TransformerConfig
6342

6443
def get_builder_cls(self) -> type:
6544
"""Resolve the configured builder through Bridge's target allowlist."""

‎src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,6 @@ class MuseGlimmerBridge(MegatronModelBridge):
138138
"""Builder-backed bridge for the complete Muse Glimmer multimodal model."""
139139

140140
MODEL_CONFIG_CLASS = MuseGlimmerModelConfig
141-
TRANSFORMER_CONFIG_CLASS = MuseGlimmerTransformerConfig
142141

143142
@staticmethod
144143
def _validate_architecture(hf_config: PretrainedConfig) -> tuple[Any, Any]:

‎src/megatron/bridge/models/muse_glimmer/muse_glimmer_config.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121

2222
from megatron.training.models.hybrid import HybridModelConfig
2323

24-
from megatron.bridge.models.common import deserialize_model_config
24+
from megatron.bridge.models.common import ModelConfigOverrideMixin, deserialize_model_config
2525
from megatron.bridge.models.transformer_config import TransformerConfig
2626
from megatron.bridge.utils.activation_map import callable_to_str, str_to_callable
2727
from megatron.bridge.utils.instantiate_utils import _resolve_target
@@ -67,10 +67,11 @@ def spatial_merge_size(self) -> int:
6767

6868

6969
@dataclass(kw_only=True)
70-
class MuseGlimmerModelConfig(HybridModelConfig):
70+
class MuseGlimmerModelConfig(ModelConfigOverrideMixin, HybridModelConfig):
7171
"""Complete builder configuration for Muse Glimmer."""
7272

7373
builder: ClassVar[str] = "megatron.bridge.models.muse_glimmer.MuseGlimmerModelBuilder"
74+
transformer_config_class: ClassVar[type[TransformerConfig]] = MuseGlimmerTransformerConfig
7475
hybrid_attention_layers_include_mlp: ClassVar[bool] = True
7576

7677
vision: MuseGlimmerVisionModelConfig = field(default_factory=MuseGlimmerVisionModelConfig)

‎tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
MuseGlimmerModelBuilder,
2929
MuseGlimmerModelConfig,
3030
MuseGlimmerTextConfig,
31+
MuseGlimmerTransformerConfig,
3132
MuseGlimmerVisionConfig,
3233
)
3334
from megatron.bridge.models.muse_glimmer.modeling_muse_glimmer import (
@@ -131,6 +132,24 @@ def test_config_conversion_uses_model_config_path_only() -> None:
131132
assert model_config.vision_config.spatial_merge_size == 2
132133

133134

135+
def test_model_config_owns_transformer_type_and_routes_flat_overrides() -> None:
136+
model_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config())
137+
138+
assert MuseGlimmerModelConfig.transformer_config_class is MuseGlimmerTransformerConfig
139+
assert "TRANSFORMER_CONFIG_CLASS" not in MuseGlimmerBridge.__dict__
140+
141+
model_config.tensor_model_parallel_size = 8
142+
model_config.seq_length = 512
143+
144+
assert model_config.transformer.tensor_model_parallel_size == 8
145+
assert "tensor_model_parallel_size" not in model_config.__dict__
146+
assert model_config.seq_length == 512
147+
assert "seq_length" in model_config.__dict__
148+
149+
with pytest.raises(AttributeError, match="declares a field"):
150+
model_config.tesnor_model_parallel_size = 4
151+
152+
134153
def test_config_conversion_uses_top_level_embedding_sharing_contract() -> None:
135154
"""The VLM config, not its nested text config, owns HF weight tying."""
136155
hf_config = _tiny_hf_config()

‎tests/unit_tests/models/test_gpt_model_config.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
# limitations under the License.
1414

1515
import inspect
16+
from dataclasses import dataclass
1617
from unittest.mock import Mock, patch
1718

1819
import pytest
@@ -103,10 +104,45 @@ def test_model_config_mapping_rejects_unknown_fields() -> None:
103104
)
104105

105106

107+
def test_model_config_class_selects_nested_transformer_type(monkeypatch: pytest.MonkeyPatch) -> None:
108+
@dataclass
109+
class SpecializedTransformerConfig(TransformerConfig):
110+
specialized_field: bool = False
111+
112+
class SpecializedModelConfig(BridgeGPTModelConfig):
113+
transformer_config_class = SpecializedTransformerConfig
114+
115+
class SpecializedBridge(LlamaBridge):
116+
MODEL_CONFIG_CLASS = SpecializedModelConfig
117+
118+
bridge = SpecializedBridge()
119+
monkeypatch.setattr(
120+
bridge,
121+
"hf_config_to_model_config_kwargs",
122+
lambda _hf_config: {
123+
"num_layers": 2,
124+
"hidden_size": 128,
125+
"num_attention_heads": 4,
126+
"ffn_hidden_size": 256,
127+
"activation_func": F.silu,
128+
"specialized_field": True,
129+
"vocab_size": 256,
130+
},
131+
)
132+
133+
model_config = bridge.hf_config_to_model_config(object())
134+
135+
assert isinstance(model_config, SpecializedModelConfig)
136+
assert isinstance(model_config.transformer, SpecializedTransformerConfig)
137+
assert model_config.specialized_field is True
138+
139+
106140
def test_gpt_model_config_is_the_bridge_default() -> None:
107141
assert MegatronModelBridge.MODEL_CONFIG_CLASS is BridgeGPTModelConfig
108142
assert LlamaBridge.MODEL_CONFIG_CLASS is BridgeGPTModelConfig
109143
assert "MODEL_CONFIG_CLASS" not in LlamaBridge.__dict__
144+
assert BridgeGPTModelConfig.transformer_config_class is TransformerConfig
145+
assert not hasattr(MegatronModelBridge, "TRANSFORMER_CONFIG_CLASS")
110146

111147

112148
def test_bridge_can_explicitly_disable_model_config() -> None:

0 commit comments

Comments
 (0)