Update README.md

README.md

@@ -1,3 +1,363 @@
 ---
 license: apache-2.0
 ---
+
+# On Gemma 4, `v_proj` does not exist on 5 of the 30 layers — and your LoRA config does not know it
+
+> This is a report created by Claude after hours and days spent training my LoRas for the Goetia merge. You might find this text useful when training your LoRas for the Gemma 4 MoE family.
+
+If you fine-tune `google/gemma-4-26B-A4B` with a `target_modules` list containing
+the string `"v_proj"`, that adapter attaches to **25 layers, not 30**. PEFT does not
+warn you, because it only raises when *nothing* matched. Your adapter is smaller than
+you think and asymmetric across depth, and the only visible sign is a
+trainable-parameter count you probably did not hand-verify.
+
+That is the short version. The long version is more interesting, because on the five
+layers where `v_proj` is absent, `k_proj` *is* the value matrix — which means an
+adapter you believe is "queries and keys only" is editing the value path, and an
+adapter you believe is "values and output only" cannot reach values there at all.
+
+I found this the hard way, by running a controlled experiment that turned out to be
+measuring something other than what I designed it to measure. Details at the end.
+
+## What is already known
+
+The mechanism itself is not my discovery, and I want to be precise about that before
+adding anything.
+
+The Gemma 4 technical report states it in one sentence, under long-context efficiency:
+"We improve memory efficiency by re-using keys as values in the global attention
+layers (except in E2B and E4B), i.e., values=keys." The official model card puts it as:
+"To optimize memory for long contexts, global layers feature unified Keys and Values,
+and apply Proportional RoPE (p-RoPE)."
+
+Devansh has gone furthest in public, and states the consequence explicitly:
+"Gemma 4 eliminates the V projection in global layers. The key projection is computed,
+then reused directly as the value, with only RMSNorm applied on the value side as a
+differentiator in the forward pass." Maarten Grootendorst's visual guide covers K=V as
+a KV-cache trick. idlemachines noted the normalization asymmetry: "values get
+normalised too, but magnitude-only, with no learned scale."
+
+And the flag is documented in `transformers`. `Gemma4TextConfig` carries a docstring
+for `attention_k_eq_v`: "Whether keys and values share the same projection weights.
+When `True`, the key projection output is reused as the value projection." One line,
+in the API reference, with no consequences drawn.
+
+One clarification worth making, because the public write-ups blur it: Gemma 4 has
+**two independent** KV-saving mechanisms, and only one of them removes `v_proj`.
+
+- `num_kv_shared_layers` — *cross-layer* sharing. Later layers reuse KV tensors from
+  an earlier non-shared layer. This is the one the official HF launch post and
+  Sebastian Raschka describe. In `26B-A4B` it is set to **0**, i.e. off.
+- `attention_k_eq_v` — *within-layer* sharing. On non-sliding layers, values are the
+  key projection. This is the one that sets `v_proj` to `None`. In `26B-A4B` it is
+  **true**.
+
+If you read about Gemma 4 KV sharing and concluded it does not affect your adapter,
+you may have read about the wrong mechanism.
+
+## The part nobody seems to have written down
+
+What I could not find anywhere is what this does to a LoRA config. Three consequences,
+and they compound.
+
+### The layout
+
+```json
+"attention_k_eq_v": true,
+"layer_types": ["sliding_attention", ..., "full_attention", ...]
+```
+
+`layer_types` puts `full_attention` at exactly indices **5, 11, 17, 23, 29** — every
+sixth layer, and always the last. The other 25 are `sliding_attention` with a
+1024-token window. Print the loaded model and the attention blocks are not uniform:
+
+| | layers 0–4, 6–10, 12–16, 18–22, 24–28 | layers **5, 11, 17, 23, 29** |
+|---|---|---|
+| `q_proj` | 2816 → 4096 | 2816 → **8192** |
+| `k_proj` | 2816 → 2048 | 2816 → **1024** |
+| `v_proj` | 2816 → 2048 | **absent** |
+| `o_proj` | 4096 → 2816 | **8192** → 2816 |
+
+That is 115 attention projections, not 4 × 30 = 120. The shapes follow from
+`head_dim: 256` / `num_key_value_heads: 8` on sliding layers versus
+`global_head_dim: 512` / `num_global_key_value_heads: 2` on global ones.
+
+This is not a broken checkpoint. `model.safetensors.index.json` of
+`google/gemma-4-26B-A4B` itself has no `self_attn.v_proj.weight` key for those five
+layers. From `modeling_gemma4.py`:
+
+```python
+self.use_alternative_attention = config.attention_k_eq_v and not self.is_sliding
+self.v_proj = (
+    nn.Linear(config.hidden_size, num_key_value_heads * self.head_dim, bias=config.attention_bias)
+    if not self.use_alternative_attention
+    else None
+)
+```
+
+Note `and not self.is_sliding`. The flag is global, its effect is not.
+
+### Consequence 1: silent partial match
+
+PEFT matches `target_modules` strings by module-name suffix. There is no
+`...layers.5.self_attn.v_proj` to match, so nothing matches, and nothing is reported —
+PEFT raises only when the whole list found nothing. So the popular seven-name list
+gives you `v_proj` on 25 layers and `q/k/o_proj` on 30.
+
+This is not hypothetical. Current Gemma 4 fine-tuning guides recommend exactly
+`["q_proj", "o_proj", "k_proj", "v_proj", "gate_proj", "up_proj", "down_proj"]`
+with no caveat about layer coverage. Community adapters use regexes like
+`(mlp|self_attn)\.(up|down|gate|q|k|v|o)_proj` that treat `v` uniformly across depth.
+Unsloth's guide uses `target_modules="all-linear"`, which sidesteps the problem by
+accident — it enumerates what exists rather than what you named — but does not
+explain it. (Per oxen.ai, recent PEFT ships default Gemma 4 target modules scoped to
+the language model via regex; that fixes vision-tower leakage, not the `v_proj` count.)
+
+### Consequence 2: on those layers, `k_proj` is the value matrix
+
+Here is the exact order of operations in `forward`:
+
+```python
+key_states   = self.k_proj(hidden_states).view(hidden_shape)
+value_states = self.v_proj(hidden_states).view(hidden_shape) if self.v_proj is not None else key_states
+
+key_states   = self.k_norm(key_states)
+key_states   = apply_rotary_pos_emb(key_states, cos, sin, unsqueeze_dim=2)
+key_states   = key_states.transpose(1, 2)
+
+value_states = self.v_norm(value_states)
+value_states = value_states.transpose(1, 2)
+```
+
+Look at *where* the fallback happens. `value_states` takes the **raw** output of
+`k_proj`, before `k_norm` and before RoPE. Then it goes through `v_norm`, an RMSNorm
+with `with_scale=False`. So one projection feeds two paths, normalized differently,
+and positional information is applied to the key path only.
+
+For anyone trying to reason about attention in terms of separable circuits — where the
+query-key product decides *where* to attend and the value-output product decides *what*
+gets written into the residual stream — this matters:
+
+- An adapter on `q_proj + k_proj` is **not** query-key only. On those five layers it
+  edits values.
+- An adapter on `v_proj + o_proj` has **no** access to values there. Only `o_proj`.
+- On those five layers the two paths **cannot be separated at all**. They share one
+  matrix.
+
+And the five layers are the only ones that see the whole context; the other 25 are
+windowed at 1024 tokens. So anything you care about that involves long context —
+instruction following deep into a chat, re-reading a system prompt every turn,
+recalling something from 20k tokens back — lives precisely where the separation you
+are testing does not exist.
+
+### Consequence 3: QK-norm changes what a LoRA delta can even do
+
+`q_norm` and `k_norm` are RMSNorm over `head_dim`, applied **after** the projection
+and **before** RoPE:
+
+```python
+query_states = self.q_proj(hidden_states).view(hidden_shape)
+query_states = self.q_norm(query_states)
+query_states = apply_rotary_pos_emb(query_states, cos, sin, unsqueeze_dim=2)
+```
+
+RMSNorm rescales each head's vector to unit RMS. So a LoRA delta on `q_proj` or
+`k_proj` can change the **direction** of queries and keys but not their **magnitude** —
+the normalization discards it. `o_proj` has no equivalent per-head constraint.
+
+If you are comparing "adapt q/k" against "adapt v/o" at equal parameter budget on
+any QK-norm architecture — Gemma 3 and 4, Qwen3, OLMo 2/3 — this asymmetry is part
+of your result whether you account for it or not. I could not find any discussion of
+QK-norm interacting with LoRA. The closest published work is on controlling attention
+logits during pretraining (Anson & Aitchison 2025; Zhai et al., σReparam, ICML 2023),
+which treats the coupled magnitudes of Q and K as the thing to control — but says
+nothing about adapters.
+
+## Two smaller landmines in the same area
+
+Both are mine as far as I can tell, and both are cheap to avoid:
+
+- `Gemma4TextRouter.proj` (2816 → 128) is a real `nn.Linear`, so a loose regex like
+  `.*proj$` will catch the MoE **router**. Adapting expert routing is a far less
+  predictable edit than adjusting attention. Exclude it explicitly.
+- `gate_proj` / `up_proj` / `down_proj` in `target_modules` land on the *dense*
+  `Gemma4TextMLP` (2816 → 2112) sitting next to the experts — that is the single
+  shared expert — and not on the 128 routed ones. The names are absorbed by the wrong
+  module, which is why the trainable-parameter count comes out plausible-looking but
+  wrong.
+
+For completeness, the neighbouring traps that **are** already well documented, so you
+do not have to rediscover them: `Gemma4TextExperts` stores weights as stacked
+`nn.Parameter`, so bitsandbytes cannot quantize them (Axolotl's expert-quantization
+docs; bitsandbytes #1849) and PEFT needs `target_parameters` rather than
+`target_modules` to reach them (PEFT docs; unsloth #4907 for the
+"abnormally low trainable parameter count" symptom). The vision and audio towers reuse
+the same leaf names, so an unanchored list leaks the adapter into them (oxen.ai;
+Axolotl multimodal docs). On my first run part of the adapter landed on the vision
+encoder and the loss flattened almost immediately.
+
+## Building the target list so it does what you wrote
+
+Stop passing projection-name strings. Read the module paths off the live model and
+assert your assumptions:
+
+```python
+import re
+
+PROJ = ("q_proj", "k_proj", "v_proj", "o_proj")
+LAYER_RE = re.compile(r"language_model\.layers\.(\d+)\.")
+
+layer_mods = {}
+for name, _ in model.named_modules():
+    if "language_model.layers." not in name:      # anchor: excludes vision/audio towers
+        continue
+    if not name.endswith(PROJ):
+        continue
+    li = int(LAYER_RE.search(name).group(1))
+    layer_mods.setdefault(li, {})[name.rsplit(".", 1)[-1]] = name
+
+global_layers = sorted(i for i, v in layer_mods.items() if "v_proj" not in v)
+assert global_layers == [5, 11, 17, 23, 29], f"layer plan changed: {global_layers}"
+assert sum(len(v) for v in layer_mods.values()) == 115
+
+# A genuinely query-key-only arm: skip k_proj on global layers,
+# where k_proj is also the value matrix.
+targets = [
+    layer_mods[li][p]
+    for li in sorted(layer_mods)
+    for p in ("q_proj", "k_proj")
+    if p in layer_mods[li] and not (p == "k_proj" and li in global_layers)
+]
+
+FORBIDDEN = ("vision", "audio", "router", "experts", "embed", "lm_head",
+             "gate_proj", "up_proj", "down_proj")
+assert not [t for t in targets if any(b in t.lower() for b in FORBIDDEN)]
+```
+
+Then re-audit **after** `get_peft_model`, because that is where a config can still
+surprise you: check that the number of trainable tensors is exactly twice the number
+of targets, and that the set of touched layers and projection types matches your plan.
+Make both `assert`, not `print`. A printed warning scrolls off screen, and an hour of
+A100 time goes with it. PEFT also ships `get_model_status()` / `get_layer_status()`,
+which is the supported way to see what actually got wrapped.
+
+## How I ran into this, and why my own numbers do not settle anything
+
+I wanted to know which half of attention carries writing style and which half is
+responsible for a fine-tune losing its grip on output format. So: two adapters, same
+data (1415 train / 74 eval), same seed, r=32, alpha=64, lr 2e-5, 2 epochs, 354 steps,
+QLoRA 4-bit, one A100 80GB. One on `v_proj + o_proj`, one on `q_proj + k_proj`.
+
+| | A: `v_proj + o_proj` | B: `q_proj + k_proj` |
+|---|---|---|
+| targets | 55 | 60 |
+| trainable params | 11,182,080 | 11,796,480 |
+| final eval loss | **1.9703** | **2.2056** |
+| mean token accuracy | 54.96 % | 51.38 % |
+
+A is below B at every eval checkpoint, monotonically; both plateau; B has 5.5 % *more*
+trainable parameters and still loses.
+
+I am not asking you to believe that means anything, for five reasons.
+
+**It reproduces a 2021 result.** "The value/output side beats the query/key side at
+equal budget" is Table 5 of the original LoRA paper: on WikiSQL, Wq 70.4 / Wk 70.0
+versus Wv 73.0 / Wo 73.2, and Wq+Wk 71.4 versus Wq+Wv 73.7. Yao et al. (IJCAI 2025)
+added the mechanism: the gradient with respect to W_K contains W_Q, which is near zero
+early in training, so Q and K are multiplicatively suppressed while V is not.
+
+**The split is not the split I thought it was.** That is this whole article. Arm B
+edited values on five layers; arm A never reached values there.
+
+**The gap is over-determined.** Beyond that, QK-norm handicaps arm B structurally, and
+my base started at loss 7.43 on this data — very far off. When the base is that far
+from the target distribution, the run mostly measures which arm can move the output
+distribution fastest, and that favours `o_proj`, which writes straight into the
+residual stream, over q/k, which only reshape a softmax.
+
+**My base was not clean, and the contamination is exactly on the seam I was testing.**
+The base is a 15-way MoE merge followed by abliteration. I went back and checked what
+the abliteration tool actually modifies: `attn.o_proj` only, on layers 14–26, by
+unconstrained L-BFGS optimization of the matrix rather than a rank-1 projection. So
+`o_proj` had been surgically rewritten in 13 of 30 layers before I started, while
+q/k/v were untouched. Arm A trains on top of rewritten matrices, arm B on top of
+pristine ones. I cannot predict the direction of that bias — a rewritten `o_proj`
+could be easier or harder to adapt further — but a comparison with a systematic
+asymmetry like that is not a fair one. Worth stating plainly: if you benchmark
+anything about attention on an abliterated model, find out which matrices were
+abliterated first.
+
+(On the other hand, the merge did not lose tensors: the merged checkpoint and
+`google/gemma-4-26B-A4B` have byte-identical `total_size` — 51,611,872,412 — and the
+shard files differ by 584 bytes, which is the size difference of the safetensors JSON
+headers. The missing `v_proj` really is architectural.)
+
+**And my behavioural hypothesis was wrong.** I expected the value/output half to carry
+style and the query/key half to be responsible for breaking format. What I saw was the
+opposite arrangement: the value/output arm carries the style *and* breaks structured
+output sooner, while the query/key arm holds formatting but barely transfers style —
+it describes a character's voice instead of speaking in it. I report that as a negative
+result rather than dropping it, with the caveat it deserves: those behavioural
+observations are single generations per setting at one context depth, judged by me,
+compared at equal adapter weight rather than equal effect size. Equal weight is the
+wrong normalization when one arm is simply a stronger intervention per unit of weight.
+That is an anecdote, not a measurement.
+
+## What I would actually like to know
+
+- Does the loss gap survive on a clean `google/gemma-4-26B-A4B`, with three seeds, and
+  with the arms rebuilt so that `k_proj` on global layers goes to neither side?
+- How much of a LoRA delta on `q_proj` / `k_proj` survives `q_norm` / `k_norm`? If the
+  answer is "not much", then a chunk of the folklore about which projections matter is
+  really a statement about where the normalization sits — and that folklore predates
+  QK-norm becoming standard.
+- Those five global layers, where W_V is literally W_K: good place to adapt, or bad?
+  It is the only spot in the model where the two paths are physically tied, and I have
+  no intuition for what a low-rank edit there does.
+
+If you have hit consequence 1 without noticing, or if you have run any of this on a
+clean base, I would like to hear about it.
+
+## References
+
+Architecture and code:
+
+- Gemma 4 Technical Report — https://arxiv.org/html/2607.02770v1
+- Gemma 4 model card (Google) — https://ai.google.dev/gemma/docs/core/model_card_4
+- `google/gemma-4-26B-A4B` config — https://huggingface.co/google/gemma-4-26B-A4B/blob/main/config.json
+- `modeling_gemma4.py` — https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma4/modeling_gemma4.py
+- `configuration_gemma4.py` (`attention_k_eq_v` docstring) — https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma4/configuration_gemma4.py
+- Gemma4 model docs — https://huggingface.co/docs/transformers/model_doc/gemma4
+
+Prior public description of K=V:
+
+- Devansh, "Google's Gemma 4 is Weirder than you Realize" — https://machine-learning-made-simple.medium.com/googles-gemma-4-is-weirder-than-you-realize-17d00d95b0d5
+- Maarten Grootendorst, "A Visual Guide to Gemma 4" — https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-gemma-4
+- idlemachines, "Gemma 4 is not your standard transformer" — https://idlemachines.co.uk/essays/gemma4-architecture
+
+LoRA target-module selection:
+
+- Hu et al., LoRA, 2021, Table 5 / §7.1 — https://arxiv.org/abs/2106.09685
+- Yao et al., IJCAI 2025, unequal importance of attention matrices — https://arxiv.org/abs/2410.02247
+- Elhage et al., 2021, query-key and output-value circuits — https://transformer-circuits.pub/2021/framework/index.html
+- PEFT LoRA developer guide — https://huggingface.co/docs/peft/developer_guides/lora
+
+Already-documented neighbouring traps:
+
+- Axolotl, MoE expert quantization — https://docs.axolotl.ai/docs/expert_quantization.html
+- bitsandbytes #1849, fused MoE weights — https://github.com/bitsandbytes-foundation/bitsandbytes/issues/1849
+- unsloth #4907, low trainable-parameter count on Gemma 4 MoE — https://github.com/unslothai/unsloth/issues/4907
+- oxen.ai, Gemma 4 fine-tuning pipeline — https://ghost.oxen.ai/writing-a-fine-tuning-and-deployment-pipeline-isnt-as-easy-as-it-looks-gemma-4-version/
+- Axolotl multimodal docs — https://docs.axolotl.ai/docs/multimodal.html
+
+Attention-logit control (background for consequence 3):
+
+- Anson & Aitchison, "Controlling changes to attention logits", 2025 — https://arxiv.org/abs/2511.21377
+- Zhai et al., σReparam, ICML 2023 — https://proceedings.mlr.press/v202/zhai23a/zhai23a.pdf
+
+---
+
+Code excerpts are from `huggingface/transformers`, Apache License 2.0. Gemma is
+provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms.
+"Gemma 4" is used descriptively; this article is not affiliated with or endorsed by Google.