[Bug Report] `TransformerBridge` cannot load heterogeneous Gemma 4 configs with Transformers 5.15

[hijohnnylin]

Summary

With transformers==5.15.0, TransformerBridge.boot_transformers cannot load Gemma 4 because map_default_transformer_lens_config reads per-layer attention attributes from the whole-model config.

Gemma 4 now represents heterogeneous attention geometry through per_layer_config. Accessing a heterogeneous field such as head_dim or num_key_value_heads on the global config raises AmbiguousGlobalPerLayerAttributeError. TransformerLens performs that access inside hasattr(...) guards, but hasattr suppresses only AttributeError, so the guard itself is the crash site.

This happens before any model weights are loaded.

Versions

Runtime reproduction:

• transformer-lens==3.7.0
• transformers==5.15.0
• PyTorch 2.9
• Python 3.12

The same checkpoints load through this stage with transformers==5.14.1, which used the older global_head_dim / num_global_key_value_heads representation.

I also checked the transformer-lens==3.7.1 source. The raising global attribute probes are unchanged.

Minimal reproduction

pip install "transformer-lens==3.7.1" "transformers==5.15.0"

from transformer_lens.model_bridge import TransformerBridge

TransformerBridge.boot_transformers("google/gemma-4-E2B", device="cpu")

No weights are required to reproduce the underlying failure directly:

from transformers import AutoConfig
from transformer_lens.model_bridge.sources.transformers import (
    map_default_transformer_lens_config,
)

config = AutoConfig.from_pretrained("google/gemma-4-E2B")
map_default_transformer_lens_config(config)

Observed exception:

  File ".../transformer_lens/model_bridge/sources/transformers.py", line 613, in boot
    tl_config = map_default_transformer_lens_config(hf_config)
  File ".../transformer_lens/model_bridge/sources/transformers.py", line 194, in map_default_transformer_lens_config
    if hasattr(source_config, "head_dim") and source_config.head_dim is not None:
  File ".../transformers/integrations/heterogeneity/configuration_utils.py", line 298, in __getattribute__
    raise AmbiguousGlobalPerLayerAttributeError(
transformers.integrations.heterogeneity.configuration_utils.AmbiguousGlobalPerLayerAttributeError:
'head_dim' is a per-layer attribute and may vary across layers. Access it via the individual layer
configs instead (e.g. config.per_layer_config[i].head_dim).

On google/gemma-4-31B, the same config-mapping call raises earlier on num_key_value_heads, which is also heterogeneous for that checkpoint.

Expected behavior

TransformerBridge.boot_transformers should load Gemma 4 without reading an ambiguous global value. Per-layer attention geometry should either be preserved explicitly in TransformerLens metadata or left to the Hugging Face modules that the Gemma 4 bridge already uses for the actual attention computation.

Actual behavior

The bridge raises during config translation, before model construction or weight loading, so every Gemma 4 checkpoint using the new heterogeneous config representation is unreachable.

Root cause

TransformerLens currently probes num_key_value_heads and head_dim as global attributes:

• https://github.com/TransformerLensOrg/TransformerLens/blob/v3.7.1/transformer_lens/model_bridge/sources/transformers.py#L86-L105
• https://github.com/TransformerLensOrg/TransformerLens/blob/v3.7.1/transformer_lens/model_bridge/sources/transformers.py#L170-L188

Transformers 5.15 intentionally raises when a caller reads a per-layer attribute from a heterogeneous global config:

• https://github.com/huggingface/transformers/blob/v5.15.0/src/transformers/integrations/heterogeneity/configuration_utils.py#L257-L275

The Gemma 4 config migration made head_dim and, where applicable, num_key_value_heads explicit per-layer values:

• Use new per_layer_config for Gemma 4 so that heterogeneous attention config is explicit huggingface/transformers#47384

For E2B, sliding-attention layers use head_dim=256 while the seven full-attention layers use head_dim=512. For 31B, the sliding and full-attention layers also use different key/value-head counts. Enabling allow_global_per_layer_attribute_access merely suppresses the exception and returns one global fallback value; it does not preserve the model's per-layer geometry.

Impact

This is a complete load failure for Gemma 4 on Transformers 5.15. It also blocks testing or using other Gemma 4 bridge features because the architecture adapter is never reached.

Suggested fix

Before probing any field that may be heterogeneous, detect the config's heterogeneity metadata and read the concrete layer configs instead of the global object. In particular:

• do not use hasattr(config, field) as a safety check for heterogeneous fields;
• do not set allow_global_per_layer_attribute_access=True as a library-wide workaround;
• preserve per-layer head_dim and num_key_value_heads values, or explicitly avoid consuming them when the Gemma 4 bridge delegates the corresponding math to Hugging Face;
• retain the scalar head_dim / num_key_value_heads path and support the legacy global_head_dim / num_global_key_value_heads fields where needed for pre-5.15 configs.

A possible first step is to branch before the global probes:

if getattr(source_config, "is_heterogeneous", False):
    layer_configs = [
        source_config.per_layer_config[i]
        for i in range(source_config.num_hidden_layers)
    ]
    per_layer_head_dim = [cfg.head_dim for cfg in layer_configs]
    per_layer_num_kv_heads = [cfg.num_key_value_heads for cfg in layer_configs]
else:
    # Existing scalar/global fallback path.
    ...

The exact TransformerLens config representation may require a broader change than this sketch, but the important point is that the global read must not occur and a single fallback value must not be treated as correct for every layer.

Suggested regression tests:

1. Config-only translation of google/gemma-4-E2B on Transformers 5.15 does not raise.
2. Config-only translation of google/gemma-4-31B does not raise on num_key_value_heads.
3. The resulting metadata retains both attention geometries rather than silently broadcasting one value.
4. A pre-5.15-style Gemma 4 config still follows the legacy global-field path.