fix(models): support Transformers 5.15 conversions

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

src/megatron/bridge/training/utils/config_utils.py

@@ -102,8 +102,12 @@ def _convert_pretrained_config_to_dict(
 
         if include_target and is_dataclass(value):
             field_names = {field.name for field in dataclass_fields(value) if not field.name.startswith("_")}
+            raw_values = vars(value)
             config_items = [
-                (field.name, getattr(value, field.name))
+                (
+                    field.name,
+                    raw_values[field.name] if field.name in raw_values else object.__getattribute__(value, field.name),
+                )
                 for field in dataclass_fields(value)
                 if not field.name.startswith("_")
             ]

tests/functional_tests/test_groups/models/deepseek/test_deepseek_conversion.py

@@ -42,7 +42,9 @@
     "num_attention_heads": 32,
     "num_experts_per_tok": 4,
     "num_hidden_layers": 2,
-    "num_key_value_heads": 4,
+    # DeepSeek MLA expands compressed KV latents to every attention head before
+    # dispatching to the HF attention backend, so it uses MHA rather than GQA.
+    "num_key_value_heads": 32,
     "num_nextn_predict_layers": 0,
     "q_lora_rank": 512,
     "topk_group": 2,

tests/functional_tests/test_groups/models/nemotron_omni/test_nemotron_omni_conversion.py

@@ -40,15 +40,10 @@
 # 52-layer pattern, which keeps a representative mamba (M), MoE-MLP (E) and
 # attention (*) layer mix.
 _LLM_LAYER_TYPES = ["mamba", "moe", "mamba", "moe", "mamba", "attention"]
-# Transformers 5.15 serializes the same hybrid blocks under their canonical names.
-_SERIALIZED_LLM_LAYER_TYPES = [
-    "linear_attention",
-    "moe",
-    "linear_attention",
-    "moe",
-    "linear_attention",
-    "full_attention",
-]
+_LLM_LAYER_TYPE_ALIASES = {
+    "linear_attention": "mamba",
+    "full_attention": "attention",
+}
 _LLM_OVERRIDES = {
     "hybrid_override_pattern": "MEMEM*",
     "layer_types": _LLM_LAYER_TYPES,
@@ -72,6 +67,11 @@ def _apply_overrides(sub_config, overrides: dict) -> None:
             setattr(sub_config, key, value)
 
 
+def _normalize_llm_layer_types(layer_types: list[str]) -> list[str]:
+    """Normalize equivalent remote-config names to the toy model's layer vocabulary."""
+    return [_LLM_LAYER_TYPE_ALIASES.get(layer_type, layer_type) for layer_type in layer_types]
+
+
 def _fix_tied_weights_keys(model: nn.Module) -> None:
     """Convert _tied_weights_keys from list to dict for transformers 5.x compatibility."""
     for module in model.modules():
@@ -223,11 +223,23 @@ def test_nemotron_omni_conversion_roundtrip(self, nemotron_omni_toy_model_path,
         with open(config_file) as f:
             saved_config = json.load(f)
 
+        source_config_file = Path(hf_model_id) / "config.json"
+        if source_config_file.exists():
+            with open(source_config_file) as f:
+                source_config = json.load(f)
+        else:
+            source_config = AutoConfig.from_pretrained(hf_model_id, trust_remote_code=True).to_dict()
+        source_llm_config = source_config["llm_config"]
+
         assert saved_config["architectures"][0] == "NemotronH_Nano_Omni_Reasoning_V3"
         assert saved_config["model_type"] == "NemotronH_Nano_Omni_Reasoning_V3"
         assert "llm_config" in saved_config
         assert "vision_config" in saved_config
         assert "sound_config" in saved_config
         assert saved_config["llm_config"]["num_hidden_layers"] == 6
-        assert saved_config["llm_config"]["layer_types"] == _SERIALIZED_LLM_LAYER_TYPES
-        assert saved_config["llm_config"]["layers_block_type"] == _SERIALIZED_LLM_LAYER_TYPES
+        expected_layer_types = _normalize_llm_layer_types(source_llm_config["layer_types"])
+        expected_block_types = _normalize_llm_layer_types(source_llm_config["layers_block_type"])
+        assert expected_layer_types == _LLM_LAYER_TYPES
+        assert expected_block_types == _LLM_LAYER_TYPES
+        assert _normalize_llm_layer_types(saved_config["llm_config"]["layer_types"]) == expected_layer_types
+        assert _normalize_llm_layer_types(saved_config["llm_config"]["layers_block_type"]) == expected_block_types

tests/unit_tests/training/utils/test_config_utils.py

@@ -25,6 +25,7 @@
 import pytest
 import torch
 from megatron.core.msc_utils import MultiStorageClientFeature
+from transformers import PreTrainedConfig
 
 from megatron.bridge.models.common import Serializable
 from megatron.bridge.training.utils.config_utils import _ConfigContainerBase, create_ddp_config
@@ -774,6 +775,24 @@ class ConfigWithNone(_ConfigContainerBase):
         assert result["optional_field"] is None
         assert result["required_field"] == "required"
 
+    def test_dataclass_pretrained_config_reads_raw_fields(self):
+        """Dataclass config serialization must not invoke guarded dynamic access."""
+
+        @dataclass
+        class HeterogeneousConfig(PreTrainedConfig):
+            num_key_value_heads: int = 8
+
+            def __getattribute__(self, name):
+                if name == "num_key_value_heads":
+                    raise RuntimeError("read this value from the per-layer config")
+                return super().__getattribute__(name)
+
+        config = HeterogeneousConfig()
+
+        serialized = _ConfigContainerBase._convert_pretrained_config_to_dict(config, include_target=True)
+
+        assert serialized["num_key_value_heads"] == 8
+
     def test_config_with_complex_nested_types(self):
         """Test ConfigContainer with complex nested types."""