Fix Gemma 4 LoRA training: vision backward NaN + audio_tower freeze leak

[john-rocky]

Problem

Fine-tuning the multimodal google/gemma-4-E2B-it checkpoint with the current SFT + LoRA pipeline hits two Gemma-4-specific bugs that together make training impossible or highly wasteful.

---

1. vision_tower backward produces NaN on padded images (blocker)

When pixel_values contain fewer real patches than max_patches (any non-square input or smaller-than-max resolution), padded query rows in the bidirectional attention mask are entirely -inf. softmax of an all-(-inf) row is finite in forward (those positions are zeroed out later in the pooler) but during backward produces NaN — exp(-inf - -inf) = exp(NaN) — so every LoRA/LR gradient ends up NaN.

Minimal repro on vanilla google/gemma-4-E2B-it:

```python
import mlx.core as mx, mlx.nn as nn
from mlx_vlm.utils import load
from mlx_vlm.trainer.utils import freeze_model

model, _ = load("google/gemma-4-E2B-it", processor_config={"trust_remote_code": True})
freeze_model(model)
model.vision_tower.unfreeze()

any pixel_values where num_real < max_patches, e.g. a random 960x624 tensor

pv = mx.random.normal(shape=(1, 3, 960, 624))

def f(x):
return model.vision_tower(x).sum()

v_direct = f(pv) # finite
v_vg, _ = nn.value_and_grad(model.vision_tower, f)(pv) # NaN
```

The forward alone is fine, so the symptom only appears once a training loop wraps the forward in nn.value_and_grad — which makes it easy to misdiagnose as a LoRA or dtype bug.

Fix: fill masked positions with a large finite negative (-1e4) instead of -inf. Forward behaviour is unchanged (exp(-1e4) is subnormal and the padded outputs are still zeroed downstream by the pooler), but the softmax gradient stays finite, so LoRA / full FT on Gemma 4 converges normally.

---

2. freeze_model leaks audio_tower / embed_audio / embed_vision

freeze_model hard-codes a top-level freeze list that pre-dates Gemma 4's audio branch. So after get_peft_model hands the model to the SFT trainer, every audio_tower.* weight (plus embed_audio, embed_vision) is still trainable. Two user-visible consequences:

• save_adapter dumps model.trainable_parameters(), so adapter files balloon from ~100 MB to ~700 MB per checkpoint with unrelated audio weights.
• Those weights silently receive gradients (even on image-only batches where no audio flows through), which is surprising and wastes optimizer memory.

Fix: add audio_tower / embed_audio / embed_vision to the top-level freeze list. A fallback is included because some audio sub-modules (notably AudioRelativePositionEmbedding, whose _inv_timescales is a plain buffer) trip nn.Module.freeze(); on exception we walk leaf modules instead.

---

Verified

With both fixes applied, Gemma 4 E2B LoRA fine-tuning on Unsplash / Pexels photos runs to completion on an M4 Max (500 iters, batch 1, bf16):

• loss 10.6 → 1.1 (no NaN)
• adapter file size: ~110 MB per checkpoint (was ~700 MB)
• LoRA merge + HF-format round-trip works for downstream CoreML conversion

---

Scope

Two surgical changes; no API changes:

• mlx_vlm/models/gemma4/vision.py — mask fill value
• mlx_vlm/trainer/utils.py — freeze-list additions + leaf-walk fallback

Context: I maintain a separate repo where this sequence (LoRA-train Gemma 4 → fuse → CoreML bundle → iPhone) is productized; happy to split into two PRs if you prefer.

[Blaizzy]

This looks good, could you trim the comments? Then we can merge

[john-rocky]

Trimmed the comments in b01c289 — kept only the non-obvious WHY (the NaN-in-backward rationale on the mask, and the freeze() fallback). PTAL.

[Blaizzy]

LGTM, thanks!