PrefixTuningConfig fails on google/gemma-4-e2b-it: tensor expand size mismatch in attention forward (peft 0.19.1, transformers 5.6.2)

[stharrold]

Summary

Training with PrefixTuningConfig on google/gemma-4-e2b-it crashes in the
first forward pass with a tensor expand-size mismatch. The error is consistent
with #2881 (Qwen3 GQA
prefix tuning), which was fixed for Qwen3 in
#2883 (peft 0.18.0). We're on
peft 0.19.1, so the Qwen3 fix is in — but Gemma 4 has a different attention
implementation that the same fix doesn't cover.

Reproducer

import torch
from peft import PrefixTuningConfig, TaskType, get_peft_model
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "google/gemma-4-e2b-it"

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    dtype=torch.bfloat16,
    device_map="auto",
    attn_implementation="sdpa",
)
cfg = PrefixTuningConfig(
    task_type=TaskType.CAUSAL_LM,
    num_virtual_tokens=20,
    prefix_projection=False,
)
model = get_peft_model(model, cfg)

tok = AutoTokenizer.from_pretrained(model_id)
text = "Hello, world." * 100  # any non-trivial prompt
inputs = tok(text, return_tensors="pt", truncation=True, max_length=1632).to(model.device)
inputs["labels"] = inputs["input_ids"]

with torch.no_grad():
    out = model(**inputs)
print("loss:", out.loss.item())

Error

RuntimeError: The expanded size of the tensor (1632) must match the existing
size (1652) at non-singleton dimension 3.
Target sizes: [1, 8, 1632, 1632].  Tensor sizes: [1, 1, 1632, 1652]

File "/opt/venv/lib/python3.12/site-packages/peft/peft_model.py", line 1964, in forward
... (wrapped calls into Gemma-4 attention)

The 20-token gap between target K-dim (1632 = sequence length) and tensor
K-dim (1652 = sequence length + num_virtual_tokens) confirms peft has
correctly prefix-extended the K/V, but Gemma-4's attention implementation
is computing a square [B, H, Q, Q] target shape — i.e., not honoring the
prefix-extended K.

Environment

• peft==0.19.1
• transformers==5.6.2
• torch==2.11.0
• Python 3.12
• Hardware: NVIDIA A100 80GB; attn_implementation="sdpa" (also expected to
  fail on other CUDA devices with sdpa)

Related issues

• #2881 — same shape-mismatch
  pattern for Qwen3 GQA prefix tuning. Fixed in
  #2883, shipped in peft 0.18.0.
  Gemma-4 needs the same class of adaptation.
• #1901 — original GQA
  support for prefix tuning (Qwen2-era).
• #3129 — open Gemma-4 +
  peft compatibility tracking. Currently scoped to LoRA's Gemma4ClippableLinear,
  but Gemma-4 + PrefixTuning is the same ecosystem question.

Workaround attempt (in progress)

We suspect attn_implementation="eager" may bypass the bug (sdpa's fused
kernel doesn't honor manually-prepended K/V the way the eager implementation
does). Will update with results once verified. If eager works, that points the
fix toward the sdpa attention path in the Gemma-4 modeling code.

Suggested fix

Mirror the Qwen3 fix from #2883 for Gemma-4's attention implementation. The
fix likely lives near peft/peft_model.py:1964 where prefix-extended K/V is
fed to the base model's forward — the attention-mask construction needs to be
aware of the additional num_virtual_tokens for Gemma-4's specific attention
layout.

Happy to test patches.

[stharrold]

Update — eager-attn workaround tested

attn_implementation="eager" does NOT fix the underlying issue. It DOES dodge the original SDPA expand-size mismatch, but exposes a second bug further along the forward path:

File "/opt/venv/lib/python3.12/site-packages/transformers/models/gemma4/modeling_gemma4.py", line 2262, in forward
    position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
RuntimeError: CUDA error: an illegal memory access was encountered

The math at that line is:

position_ids = torch.arange(seq_len) + past_seen_tokens

When peft PrefixTuningConfig is active, past_seen_tokens reflects the prefix-extended cache (i.e., includes num_virtual_tokens=20). The Gemma-4 modeling code computes position_ids assuming past_seen_tokens is real history, so position_ids ends up containing values beyond the position-embedding table size. The illegal memory access is the OOB lookup downstream.

So there are at least two peft + Gemma-4 incompatibilities:

1. Attention K-mask shape (with sdpa): [B, H, Q, Q] target vs [1, 1, Q, Q+num_virtual] actual → expand RuntimeError.
2. Position id overflow (with eager): torch.arange(Q) + past_seen_tokens produces OOB indices when past_seen_tokens was inflated by the prefix.

Same root cause class: Gemma-4's modeling code doesn't distinguish "real cached tokens" from "synthetic prefix tokens that peft prepended." The fix has to thread that distinction through both the attention-mask construction and the position-id construction.

Environment for this run: same as the original report (peft 0.19.1, transformers 5.6.2, torch 2.11.0, A100 80GB, Python 3.12). Reproducer is the same as the original — just swap attn_implementation="sdpa" → "eager".

Hope this helps narrow the patch scope. Still happy to test fixes on either attention impl.

[stharrold]

Update — PromptEncoderConfig (P-Tuning v2) also fails on Gemma-4

Tested PromptEncoderConfig(num_virtual_tokens=20) as an alternative to PrefixTuning — different injection point (input embedding layer vs past_key_values). Hit a different bug:

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 1690.50 GiB.
File "transformers/models/gemma4/modeling_gemma4.py", line 1635, in get_per_layer_inputs
    inputs_embeds[:, :, None, :]
File "transformers/models/gemma4/modeling_gemma4.py", line 2195, in forward
    per_layer_inputs = self.language_model.get_per_layer_inputs(llm_input_ids, llm_inputs_embeds)
File "peft/peft_model.py", line 1977, in forward

The 1.69 TB allocation is a shape-multiplication error. Gemma-4 has a Per-Layer Embeddings (PLE) feature that materializes per-layer-input tensors of shape (batch, seq_len, num_layers, hidden_size). With batch=1, seq~1652, num_layers=26, hidden=2048, bfloat16 → naive size ~177 MB. Getting 1.69 TB means a factor of ~10000 wrong somewhere — likely seq*hidden getting squared or a seq dim accidentally multiplying with itself when the PromptEncoder injects virtual tokens.

So that's now three distinct peft + Gemma-4 incompatibilities documented:

1. PrefixTuningConfig + sdpa → attention-mask expand mismatch (original report)
2. PrefixTuningConfig + eager → position_ids overflow at modeling_gemma4.py:2262 (prior comment)
3. PromptEncoderConfig (this comment) → 1.69 TB shape error at modeling_gemma4.py:1635

Common thread: Gemma-4's modeling code (and its PLE feature in particular) doesn't compose with peft's input-extension mechanisms, regardless of which PEFT method.

Could be worth filing #3 as a separate issue if maintainers prefer per-PEFT-method tracking — let me know. For now, posting here since the diagnostic chain came from the same investigation.

Same environment as original report (peft 0.19.1, transformers 5.6.2, torch 2.11.0, A100 80GB).

[BenjaminBossan]

Thanks for your detailed report @stharrold.

Regarding prefix tuning: I tried to dig into it a bit but always ran into another wall. A big issue is that head dim differs depending on whether sliding window attention is being used or not:

https://github.com/huggingface/transformers/blob/223fe5231b783fbfb25296bb0a243dad5d158cde/src/transformers/models/gemma4/modeling_gemma4.py#L1147

Prefix tuning cannot deal with different sizes here and thus we eventually get an error.

As for P-tuning, it works by injecting input embeds into the model. In the per-layer embedding code of transformers, it tries to reverse engineer the input IDs that correspond to those embeddings:

https://github.com/huggingface/transformers/blob/223fe5231b783fbfb25296bb0a243dad5d158cde/src/transformers/models/gemma4/modeling_gemma4.py#L1713-L1720

This cannot work, so even if the huge shape (and thus OOM) was avoided, we would still not get the correct result.

I think overall this means that Gemma 4 is just fundamentally incompatible with these prompt learning techniques. We have to consider those methods were developed with very specific model architectures in mind but new innovations in model architectures make these assumptions shaky. I don't see a good way to extend these PEFT methods to work with Gemma (maybe @zucchini-nlp has an idea) and would thus recommend to use different PEFT techniques instead, like LoRA.

[zucchini-nlp]

Yeah, gemma4 is a specif arch with different per-layer dimensions, and we don't have a standard way to encode them. I don't think PEFT can cleanly support it, and not sure if it is worth adding a dirty hack

[stharrold]

Thanks @BenjaminBossan and @zucchini-nlp for the careful diagnosis. The "different head dim depending on sliding window attention" pointer (modeling_gemma4.py:1147) and the per-layer-embedding input-ID reverse-engineering pointer (modeling_gemma4.py:1713-1720) make the architectural incompatibility concrete in a way the surface-level errors didn't surface clearly.

For our purposes we'll fall back to LoRA per your recommendation — the existing LoRA path was working for us prior to this evaluation, so it's a low-friction return. We'll separately evaluate gemma's own JAX PrefixTuning module (google-deepmind/gemma#631) since that path is authored with Gemma's per-layer-dim / sliding-window quirks in mind from the start. If anything from that experiment is upstream-relevant to peft, I'll comment back here.

Feel free to close this issue at your discretion — the diagnosis is the value, and there's no need to keep it open as a tracking burden if it's effectively WONTFIX from peft's side.

[BenjaminBossan]

@stharrold Hey, thanks for pointing to the Google implementation for prefix tuning. I checked if there is anything there we can emulate in PEFT. IIUC, what they do is to initialize the prefix for each layer individually. That is basically impossible to change in PEFT without breaking backwards compatibility, so it's not an option. But it did give me an idea: How about over-provisioning the shape of the projection matrix and then slice it to size if it's too big? It's not the most efficient, but it's better than nothing. I created a PR for this, #3205. Could you please check if it works for you?

For P-tuning, I still believe this can't be solved on the PEFT side.

[BenjaminBossan]

We provided a basic fix for prefix tuning with SDPA in #3205. Eager attention and P-tuning are considered WON'T FIX for now. Closing the PR but please post if something else comes regarding Gemma 4.