Gemma 4 training fails as `flash-attn` rejects `head_dim=512` on global attention layers

[jamesbraza]

Training any google/gemma-4-* model through skyrl-train crashes on the first training step with:

File ".../transformers/models/gemma4/modeling_gemma4.py", line 1243, in forward
    attn_output, attn_weights = attention_interface(...)
File ".../transformers/integrations/flash_attention.py", line 69, in flash_attention_forward
    attn_output = _flash_attention_forward(...)
File ".../flash_attn/flash_attn_interface.py", line 91, in _flash_attn_forward
    out, softmax_lse, S_dmask, rng_state = flash_attn_gpu.fwd(...)
RuntimeError: FlashAttention forward only supports head dimension at most 256

Gemma4TextConfig's global-attention layers use global_head_dim=512 (link 1, link2). FlashAttention 2's CUDA kernel is hard-capped at head_dim ≤ 256, so the first global layer raises the above RuntimeError.

SkyRL currently:

1. Sets attn_implementation="flash_attention_2" model-wide whenever trainer.flash_attn=True (link)
2. Requires FA2 whenever trainer.use_sample_packing=True (link)

This makes Gemma 4 is unusable on the FSDP path.

Notably, FA3 on Hopper also caps both forward and backward at head_dim=256. FA4 beta currently has this 256 dim cap as well. So upgrading the flash-attn pin won't help.

---

A proposed fix is to support per-layer-type attn_implementation for models whose config.layer_types alternates between attention variants. transformers v5 already accepts a dict form:

model_config._attn_implementation = {
    "sliding_attention": "flash_attention_2",
    "full_attention": "sdpa",  # or "flex_attention"
}

[tamoghnokandar]

This issue has already been opened in Flash attention's repository - Link . I think it's better to wait for their fix.