|
14 | 14 |
|
15 | 15 | from dataclasses import fields as dataclass_fields |
16 | 16 | from dataclasses import is_dataclass |
17 | | -from typing import Any |
| 17 | +from typing import Any, ClassVar |
18 | 18 |
|
19 | 19 | from megatron.training.models.base import ( |
20 | 20 | BuildConfigT, # noqa: F401 |
|
27 | 27 | ModelConfig as _MegatronModelConfig, |
28 | 28 | ) |
29 | 29 |
|
| 30 | +from megatron.bridge.models.transformer_config import TransformerConfig |
30 | 31 | from megatron.bridge.utils.instantiate_utils import _resolve_target, _validate_target_prefix, instantiate |
31 | 32 |
|
32 | 33 |
|
@@ -73,6 +74,42 @@ def _from_dict(subdata: dict[str, Any], full_key: str) -> Any: |
73 | 74 | return result |
74 | 75 |
|
75 | 76 |
|
| 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 | + |
76 | 113 | class ModelConfig(_MegatronModelConfig): |
77 | 114 | """Bridge compatibility wrapper for Megatron-LM model configs.""" |
78 | 115 |
|
|
0 commit comments