fix(models): repair Transformers 5.15 compatibility

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

Author: yaoyu-33

src/megatron/bridge/models/gemma/gemma4_bridge.py

@@ -98,7 +98,14 @@ def _infer_attn_pattern(layer_types: list[str]) -> tuple[int, int] | list[str]:
     return (len(layer_types), 0)
 
 
-def _attention_config_value(hf_config: Any, layer_type: str, field_name: str, default: Any) -> Any:
+def _attention_config_value(
+    hf_config: Any,
+    layer_type: str,
+    field_name: str,
+    default: Any,
+    *,
+    legacy_field_name: str | None = None,
+) -> Any:
     """Read an attention field from a concrete HF layer config when available."""
     layer_types = getattr(hf_config, "layer_types", None)
     per_layer_config = getattr(hf_config, "per_layer_config", None)
@@ -110,9 +117,10 @@ def _attention_config_value(hf_config: Any, layer_type: str, field_name: str, de
         else:
             return getattr(per_layer_config[layer_index], field_name, default)
 
-    # Transformers 5.15 rejects direct reads of fields that vary by layer. The
-    # global value is still the correct fallback for the local-attention config.
-    return getattr(hf_config, "__dict__", {}).get(field_name, default)
+    # Transformers 5.15 rejects direct reads of fields that vary by layer.
+    # Older configs instead serialize separate flat local/global field names.
+    fallback_name = legacy_field_name or field_name
+    return getattr(hf_config, "__dict__", {}).get(fallback_name, default)
 
 
 def _layer_types_from_provider(provider: Gemma4ModelProvider | Gemma4DenseProvider) -> list[str]:
@@ -245,7 +253,8 @@ def _build_dense_provider(self, hf_config) -> Gemma4DenseProvider:
             hf_config,
             "full_attention",
             "num_key_value_heads",
-            getattr(hf_config, "num_global_key_value_heads", num_query_groups),
+            num_query_groups,
+            legacy_field_name="num_global_key_value_heads",
         )
 
         self._dense_num_attention_heads = num_attention_heads
@@ -265,7 +274,13 @@ def _build_dense_provider(self, hf_config) -> Gemma4DenseProvider:
             num_attention_heads=num_attention_heads,
             num_query_groups=num_query_groups,
             kv_channels=_attention_config_value(hf_config, "sliding_attention", "head_dim", 256),
-            global_kv_channels=_attention_config_value(hf_config, "full_attention", "head_dim", 512),
+            global_kv_channels=_attention_config_value(
+                hf_config,
+                "full_attention",
+                "head_dim",
+                512,
+                legacy_field_name="global_head_dim",
+            ),
             num_global_query_groups=num_global_query_groups,
             seq_length=hf_config.max_position_embeddings,
             vocab_size=hf_config.vocab_size,
@@ -307,12 +322,19 @@ def _build_moe_provider(self, hf_config) -> Gemma4ModelProvider:
         provider.kv_channels = head_dim
         provider.qk_layernorm = True
 
-        provider.global_head_dim = _attention_config_value(hf_config, "full_attention", "head_dim", 512)
+        provider.global_head_dim = _attention_config_value(
+            hf_config,
+            "full_attention",
+            "head_dim",
+            512,
+            legacy_field_name="global_head_dim",
+        )
         provider.num_global_key_value_heads = _attention_config_value(
             hf_config,
             "full_attention",
             "num_key_value_heads",
-            getattr(hf_config, "num_global_key_value_heads", 2),
+            2,
+            legacy_field_name="num_global_key_value_heads",
         )
         provider.attention_k_eq_v = getattr(hf_config, "attention_k_eq_v", False)
 
@@ -447,16 +469,23 @@ def maybe_modify_loaded_hf_weight(
                     text_config, "num_attention_heads", getattr(self, "_dense_num_attention_heads", 8)
                 )
                 kv_head_dim = q_weight.shape[0] // num_q_heads
-                num_kv_heads = getattr(text_config, "num_key_value_heads", getattr(self, "_dense_num_query_groups", 2))
+                num_kv_heads = _attention_config_value(
+                    text_config,
+                    "sliding_attention",
+                    "num_key_value_heads",
+                    getattr(self, "_dense_num_query_groups", 2),
+                )
                 layer_match = re.search(r"layers\.(\d+)\.", q_name)
                 layer_types = getattr(text_config, "layer_types", None)
                 if layer_match and layer_types:
                     layer_idx = int(layer_match.group(1))
                     if layer_idx < len(layer_types) and layer_types[layer_idx] == "full_attention":
-                        num_global_kv_heads = getattr(
+                        num_global_kv_heads = _attention_config_value(
                             text_config,
-                            "num_global_key_value_heads",
+                            "full_attention",
+                            "num_key_value_heads",
                             getattr(self, "_dense_num_global_query_groups", None),
+                            legacy_field_name="num_global_key_value_heads",
                         )
                         if num_global_kv_heads is not None:
                             num_kv_heads = num_global_kv_heads

src/megatron/bridge/models/gemma_vl/gemma4_vl_bridge.py

@@ -49,6 +49,7 @@
 )
 from megatron.bridge.models.gemma.gemma4_bridge import (
     Gemma4Bridge,
+    _attention_config_value,
     _Gemma4QKVMapping,
     _infer_attn_pattern,
 )
@@ -99,6 +100,12 @@ def provider_bridge(
             return self._build_moe_provider(text_config)
 
         provider_kwargs = self.hf_config_to_provider_kwargs(text_config)
+        provider_kwargs["num_query_groups"] = _attention_config_value(
+            text_config,
+            "sliding_attention",
+            "num_key_value_heads",
+            4,
+        )
         provider = Gemma4VLModelProvider(**provider_kwargs)
 
         provider.window_size = getattr(text_config, "sliding_window", 1024)
@@ -107,13 +114,25 @@ def provider_bridge(
             rope_theta_from_hf(text_config),
         )
 
-        head_dim = getattr(text_config, "head_dim", 256)
+        head_dim = _attention_config_value(text_config, "sliding_attention", "head_dim", 256)
         provider.softmax_scale = 1.0
         provider.kv_channels = head_dim
         provider.qk_layernorm = True
 
-        provider.global_head_dim = getattr(text_config, "global_head_dim", 512)
-        provider.num_global_key_value_heads = getattr(text_config, "num_global_key_value_heads", 2)
+        provider.global_head_dim = _attention_config_value(
+            text_config,
+            "full_attention",
+            "head_dim",
+            512,
+            legacy_field_name="global_head_dim",
+        )
+        provider.num_global_key_value_heads = _attention_config_value(
+            text_config,
+            "full_attention",
+            "num_key_value_heads",
+            2,
+            legacy_field_name="num_global_key_value_heads",
+        )
         provider.attention_k_eq_v = getattr(text_config, "attention_k_eq_v", False)
 
         rope_params = getattr(text_config, "rope_parameters", {})

tests/unit_tests/models/gemma/test_gemma4_bridge.py

@@ -221,6 +221,8 @@ def test_runtime_preserves_nonperiodic_attention_schedule(self, bridge):
         provider = bridge.provider_bridge(pretrained)
         assert provider.kv_channels == 256
         assert provider.global_head_dim == 512
+        assert provider.num_query_groups == hf_config.per_layer_config[1].num_key_value_heads
+        assert provider.num_global_key_value_heads == hf_config.per_layer_config[0].num_key_value_heads
         runtime_layer_types = [
             "sliding_attention"
             if _is_local_attn_layer(layer_number, provider.interleaved_attn_pattern)

tests/unit_tests/recipes/test_muse_glimmer_recipes.py

@@ -149,6 +149,8 @@ def test_muse_glimmer_performance_recipe_is_dense_decoder_only() -> None:
         "CUDA_DEVICE_MAX_CONNECTIONS": 1,
         "NCCL_GRAPH_REGISTER": 0,
         "NCCL_NVLS_ENABLE": 0,
+        "NVTE_BWD_LAYERNORM_SM_MARGIN": 0,
+        "NVTE_FWD_LAYERNORM_SM_MARGIN": 0,
         "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
         "TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
         "TORCH_NCCL_HIGH_PRIORITY": 1,