Summary

• Add source normalization in fixup_fused_lm_head() so the compiler's pattern 3 regex matches Gemma 4's ForConditionalGeneration.forward()
• Without this, logged training loss is inflated by ~gradient_accumulation_steps (e.g. loss shows ~52 instead of ~13 with GA=4)
• Training itself is correct (gradients are fine), only the logged metric is wrong

Problem

Gemma 4 uses different variable names and a chained .to() call in its loss code:

flat_logits = shift_logits.view(-1, self.config.get_text_config().vocab_size)
flat_labels = shift_labels.view(-1).to(shift_logits.device)
loss = loss_fct(flat_logits, flat_labels)

The compiler's cross_entropy_find_3 regex expects shift_logits/shift_labels and .view(-1) and .to(device) on separate lines. None of the 3 patterns match, so unsloth_fused_ce_loss is never applied.

Fix

Normalize Gemma 4's forward source in fixup_fused_lm_head() before the regex runs:

1. flat_logits = shift_logits.view(-1, ...) -> shift_logits = shift_logits.view(-1, ...)
2. Split flat_labels = shift_labels.view(-1).to(device) into two lines
3. loss = loss_fct(flat_logits, flat_labels) -> loss = loss_fct(shift_logits, shift_labels)

After normalization, pattern 3 matches and applies the fused cross entropy.

Test plan

• apply_fused_lm_head() on Gemma 4 forward source returns modified result
• Log output: [3/3 pattern] Successfully patched fast linear cross entropy for Gemma4ForConditionalGeneration
• Finetune Gemma 4 E2B-it with GA>1 and confirm loss is no longer inflated
• Verify other VLMs (Gemma 3, Llama 4, Qwen2.5-VL) still patch correctly