transformers - 💡(How to fix) Fix Gemma-4 training with FSDP2 raises `KeyError` in `Gemma4TextAttention.forward` because `shared_kv_states` is rebuilt per-layer [1 comments, 2 participants]

"""`torchrun --nproc_per_node=1 repro.py` → KeyError(0)."""

import os
import torch
import torch.distributed as dist
from torch.distributed._composable.fsdp import MixedPrecisionPolicy, fully_shard
from torch.distributed.device_mesh import init_device_mesh
from transformers import Gemma4TextConfig, Gemma4TextModel

dist.init_process_group(backend="nccl")
torch.cuda.set_device(dist.get_rank())
cast = os.environ.get("CAST_FORWARD_INPUTS", "1") == "1"
print(f"cast_forward_inputs={cast}", flush=True)

# 2-layer Gemma-4: layer 0 writes shared_kv_states[0], layer 1 reads it.
#   * num_kv_shared_layers=1 — default 0 means no sharing layer, bug can't fire.
#   * layer_types must be set (default None breaks model init); both must be
#     full_attention because Gemma-4 force-overrides the last layer to full,
#     so the first must match for the sharing layer's same-type source lookup.
config = Gemma4TextConfig(
    num_hidden_layers=2,
    num_kv_shared_layers=1,
    layer_types=["full_attention", "full_attention"],
)
model = Gemma4TextModel(config).cuda().train()
model.gradient_checkpointing_enable(
    gradient_checkpointing_kwargs={"use_reentrant": True}
)

mesh = init_device_mesh("cuda", (dist.get_world_size(),))
mp = MixedPrecisionPolicy(param_dtype=torch.bfloat16, cast_forward_inputs=cast)
for layer in model.layers:  # per-layer fully_shard triggers the per-call rebuild
    fully_shard(layer, mesh=mesh, mp_policy=mp)
fully_shard(model, mesh=mesh, mp_policy=mp)

out = model(input_ids=torch.randint(0, config.vocab_size, (1, 8), device="cuda"))
print(f"PASS — last_hidden_state.shape={tuple(out.last_hidden_state.shape)}", flush=True)
dist.destroy_process_group()

---

cast_forward_inputs=True
File ".../transformers/models/gemma4/modeling_gemma4.py", line 1218, in forward
    key_states, value_states = shared_kv_states[self.kv_shared_layer_index]
KeyError: 0

---

File "transformers/modeling_layers.py", line 92, in __call__
    return self._gradient_checkpointing_func(partial(super().__call__, **kwargs), *args)
File "torch/utils/checkpoint.py", line 268, in CheckpointFunction.forward
    outputs = run_function(*args)
File "transformers/models/gemma4/modeling_gemma4.py", line 1219, in Gemma4TextAttention.forward
    key_states, value_states = shared_kv_states[self.kv_shared_layer_index]
KeyError: 22

---

GEMMA4_KV_DEBUG layer_idx=22 store_full_length_kv=True               dict_id=15592302014016 dict_len=0
GEMMA4_KV_DEBUG layer_idx=24 is_kv_shared=True kv_shared_layer_index=22 dict_id=15592302338496 dict_len=0
GEMMA4_KV_DEBUG layer_idx=23 store_full_length_kv=True               dict_id=11041820011008 dict_len=0
GEMMA4_KV_DEBUG layer_idx=24 is_kv_shared=True kv_shared_layer_index=22 dict_id=11041820015872 dict_len=0

---

args, kwargs = (
    _apply_to_tensors(cast_fn, args),
    _apply_to_tensors(cast_fn, kwargs),
)