[Gemma 4] Gemma4UnifiedForConditionalGeneration text-only inference produces degenerate output (token repetition collapse)
Author: synapticode-ai
Pull request: no

- `transformers`: latest `main` (`5.10.0.dev0`, tested 2026-06-08)
- `torch`: 2.7.0
- Platform: macOS, Apple M4 Pro (reproduced on **both MPS and CPU**)
- Python: 3.12

The collapse is **device-, dtype-, and attention-implementation-independent** — it reproduces on CPU with `float32` and `attn_implementation="eager"`, so it is not an MPS or a fused-kernel artifact.

### Who can help?

@ArthurZucker @molbap (Gemma)

### Reproduction

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "google/gemma-4-12B-it"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, torch_dtype=torch.bfloat16, device_map="cpu"
).eval()

prompt = "The capital of France is"
ids = tok(prompt, return_tensors="pt")

# 1) generation collapses to a repeated token
out = model.generate(**ids, max_new_tokens=12, do_sample=False)
print(tok.decode(out[0], skip_special_tokens=True))
# -> "The capital of France is111111111111"

# 2) prompt cross-entropy is far too high for a coherent 12B
with torch.no_grad():
    logits = model(**ids).logits
loss = torch.nn.functional.cross_entropy(
    logits[0, :-1], ids["input_ids"][0, 1:]
)
print(float(loss))   # ~17.9  (a coherent 12B sits at ~2-3)
```

**Expected:** coherent continuation (e.g. "… Paris."), prompt loss in the ~2–3 range.
**Actual:** the model emits a fixed repeated token (`"1111…"`); prompt cross-entropy ≈ **17.9**; greedy top-5 next-token candidates are `['1', '-', '.', '0', '_']`.

### The weights are good — the fault is in the `transformers` forward

The **same weights**, converted to a `Q4_K_M` GGUF, generate coherent text under `llama.cpp` (build `9430`). So this is not a checkpoint problem; the `Gemma4Unified` forward path in `transformers` is producing corrupted hidden states for **text-only** input.

### Diagnostic findings

We ran an in-process diagnosis to localize the fault. Summary:

1. **Multimodal context partially rescues the text path.** With a *real* image in the input, the text-token cross-entropy drops from **17.9 → 5.95**. Dummy / zero-embedding prepends do **not** rescue it (best 13.7). Setting `mm_token_type_ids` to all-zeros has **no effect**. This points to a fault on the **text-only path** that the image-content branch happens to avoid.
2. **`language_model.forward` receives well-formed inputs in both cases.** Per-type attention masks, sequential `position_ids`, and correctly scaled input embeddings are all structurally identical between text-only and with-image runs at the `language_model` input boundary — so this is **not** a missing-input-state problem at that boundary; the corruption is inside the decoder stack.
3. **Eleven components individually ruled out** as the sole cause (each toggled/normalized and re-measured, no recovery): per-type RoPE dispatch, per-type attention-mask dispatch, attention scaling (two variants), `embed_scale`, `final_logit_softcapping`, the RMSNorm `+1` convention, `layer_scalar` (three variants), and `v_norm`.
4. **Attention weights look healthy** (max attention probability 0.28–0.71, entropy 0.9–1.9 nats across layers) — no attention-collapse signature.
5. **Hidden-state norms stay alive and varied across all 48 layers** — magnitude is preserved while *content* is corrupted, consistent with a subtle per-layer transform/dispatch error rather than a NaN/overflow blowup.

### Possibly related

- #45200 — `mm_token_type_ids` handling for fine-tuning. **Tested; no effect on this inference bug.**
- google-deepmind/gemma#622 — reported repetition on the 31B / 26B Gemma 4 variants (same family; possibly a shared root cause).

### Affected

- `google/gemma-4-12B-it` **text-only** inference (`Gemma4UnifiedForConditionalGeneration`).
- Multimodal (with-image) inference is **partially** functional (loss ~5.95).

---

*Draft prepared 2026-06-08. Probe numbers from `transformers` `5.10.0.dev0` on M4 Pro.*

Comments:
The model was trained with specific format in mind, and applying a chat template I got a sensible output

```
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "google/gemma-4-12B-it"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="cuda")

chat = [
  {"role": "user", "content": "The capital of France is"},
]
# try disabling thinking mode if not needed
ids = tok.apply_chat_template(chat, tokenize=True, return_dict=True, add_generation_prompt=True, return_tensors="pt", enable_thinking=True).to("cuda")

out = model.generate(**ids, max_new_tokens=50, do_sample=False)
print(tok.batch_decode(out))
>>> ['<bos><|turn>system\n
\n<turn|>\n<|turn>user\nThe capital of France is<turn|>\n<|turn>model\n<|channel>thought\nThe user is asking for the capital of France']
```

Thanks @zucchini-nlp — really appreciate the fast response.

That makes sense. The -it model was trained on chat-formatted input, so raw
text without apply_chat_template is effectively out-of-distribution. The
severity of the collapse (total degeneration to a single repeated token
rather than just degraded quality) threw us off — we assumed a forward-path
fault rather than a format dependency.

Your example with the chat template produces coherent output on our end
too. We'll move our evaluation to the base model (google/gemma-4-12B) for
raw-text perplexity, which is the right comparator for our use case.

Thanks again for the quick clarification — closing this as resolved.
Kind Regards

*Rob* Lakelin

On Thu, 11 Jun 2026 at 01:24, Raushan Turganbay ***@***.***>
wrote:

> *zucchini-nlp* left a comment (huggingface/transformers#46531)
> <https://github.com/huggingface/transformers/issues/46531#issuecomment-4671778174>
>
> The model was trained with specific format in mind, and applying a chat
> template I got a sensible output
>
> import torch
> from transformers import AutoModelForCausalLM, AutoTokenizer
>
> model_id = "google/gemma-4-12B-it"
> tok = AutoTokenizer.from_pretrained(model_id)
> model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="cuda")
>
> chat = [
>   {"role": "user", "content": "The capital of France is"},
> ]
> # try disabling thinking mode if not needed
> ids = tok.apply_chat_template(chat, tokenize=True, return_dict=True, add_generation_prompt=True, return_tensors="pt", enable_thinking=True).to("cuda")
>
> out = model.generate(**ids, max_new_tokens=50, do_sample=False)
> print(tok.batch_decode(out))
> >>> ['<bos><|turn>system\n
\n<turn|>\n<|turn>user\nThe capital of France is<turn|>\n<|turn>model\n<|channel>thought\nThe user is asking for the capital of France']
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/huggingface/transformers/issues/46531?email_source=notifications&email_token=B6YOFCDVK5HMJGOATH2QQ3T47F4TPA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTINRXGE3TOOBRG42KM4TFMFZW63VGMF2XI2DPOKSWK5TFNZ2KYZTPN52GK4S7MNWGSY3L#issuecomment-4671778174>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/B6YOFCH52PDZURJWV63HRU347F4TPAVCNFSNUABFKJSXA33TNF2G64TZHMYTKNJSGIYDMNBRHNEXG43VMU5TINRSHA3TCNJUGUYKC5QC>
> .
> Triage notifications, keep track of coding agent tasks and review pull
> requests on the go with GitHub Mobile for iOS
> <https://github.com/notifications/mobile/ios/B6YOFCHV673K4ZPTBHAWPZD47F4TPA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTINRXGE3TOOBRG42KM4TFMFZW63VGMF2XI2DPOKSWK5TFNZ2KUZTPN52GK4S7NFXXG>
> and Android
> <https://github.com/notifications/mobile/android/B6YOFCBWPJQSC56F4UD3LFD47F4TPA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTINRXGE3TOOBRG42KM4TFMFZW63VGMF2XI2DPOKSWK5TFNZ2K4ZTPN52GK4S7MFXGI4TPNFSA>.
> Download it today!
> You are receiving this because you authored the thread.Message ID:
> ***@***.***>
>