Gemma 4 E4B support (gemma4 arch, MatFormer/PLE): inference + full-GPU LoRA training

[3ndetz]

Summary

Adds support for Gemma-4-E4B (gemma3n-style MatFormer with KV-sharing). Upstream QVAC declares LLM_ARCH_GEMMA4 but cannot load E4B GGUFs (missing tensor 'blk.N.attn_k.weight'). This PR fixes inference loading and gets training to run end-to-end on the CPU backend.

Status: draft / WIP.

What works after this PR

1. Inference loading of Gemma-4-E4B ✅

   • Loader required attn_k/attn_k_norm for every layer, but E4B shares KV across the last attention.shared_kv_layers layers (those have no own attn_k/attn_v/attn_k_norm). hparams.n_layer_kv_from_start and the KV reuse callback already exist — only the loader was over-strict. Mark wk/attn_k_norm NOT_REQUIRED for shared-KV layers.
   • Verified: E4B Q4_K_M loads + generates in llama-cli (was a hard failure).

2. TANH backward ✅ — d/dx tanh = 1-tanh², via existing ops (all backends). Gemma uses tanh for soft-capping.

3. GELU backward ✅ — exact tanh-approx GELU derivative matching ggml_gelu_backward_f32, built from existing ops (scale_bias/sqr/mul/tanh/add), no new kernels.

4. rms_norm_back contiguity ✅ — gemma4 feeds a non-contiguous grad into rms_norm_back; CPU op asserts contiguous. Force ggml_cont.

Verified

• Gemma-4-E4B trains end-to-end on the CPU backend and saves a valid GGUF LoRA adapter (168 tensors). Previously: would not even load.
• Generic backward ops (TANH/GELU) are numerically correct — Qwen3-0.6B loss converges with the same patched binary.

Remaining work

• E4B loss diverges (rises instead of falling) while Qwen converges → a gemma4-specific backward numerical issue remains. Prime suspects: attention logit/soft-cap tanh path, Gemma's (1+weight) RMSNorm convention in the backward, or the altup / per-layer-embedding path. Needs per-op gradient checking on the gemma4 graph.
• CUDA out_prod raises CUBLAS_STATUS_INVALID_VALUE on gemma4's variable per-layer KV dims (ggml-cuda/out-prod.cu:113). CPU backend is unaffected; GPU training needs an out_prod shape fix.
• dataset prep underflows in ggml_opt_dataset_init when corpus < context window (ndata wraps to ~2^64). Workaround: corpus > -c. A clamp/error would be friendlier.

Guidance welcome on the gemma4-specific backward (esp. whether soft-cap/altup have intended backward paths) and on out_prod for variable-KV shapes.

[3ndetz]

Update: Gemma-4-E4B LoRA training now converges end-to-end.

The last blocker was a NaN in the GELU backward I added: for Gemma's large pre-GELU activations, (1+3a*x^2)→Inf while (1-tanh^2(g))→0, giving Inf*0=NaN → NaN gradients → adapter never learned. (Qwen never hit this — it uses SiLU, not GELU.) Fixed by clamping x to [-30,30] in the derivative (gelu' is saturated there, so it's numerically faithful).

Evidence on a Gemma-4-E4B Q4_K_M model:

• val loss converges monotonically: 3.24 → 2.22 → 1.89 → 1.61 (was frozen before)
• saved adapter lora_b: 0 NaN, 42 finite-nonzero (was 42 NaN)
• inference output changes with vs without the trained adapter

Remaining for full GPU speed: the CUDA out_prod kernel raises CUBLAS_INVALID_VALUE on gemma4's variable per-layer KV dims; I worked around it locally by falling back that op to CPU. A proper out_prod fix for those shapes would unlock full-GPU E4B training.

[3ndetz]

Pinpointed the CUDA out_prod blocker for full-GPU gemma4 training:

The failing cublasSgemm (out-prod.cu:113) is the attention-internal out_prod in the backward (head-batched, ne2=8), where src0 is a transposed/permuted view. Captured dims at failure:

ne0=96 ne1=96 ne01=512 lda=48 ldb=4096 ldc=96 ne2=8 src1_T=1

lda=48 < m=ne0=96 violates cuBLAS's requirement (for OP_N, lda >= max(1,m)) → CUBLAS_STATUS_INVALID_VALUE. It comes from lda = nb01/sizeof(float) on a non-contiguous src0.

Wrapping the generic MUL_MAT-backward out_prod args in ggml_cont did not fix it — this out_prod originates inside the attention backward, not the generic path. Proper fix: in ggml_cuda_out_prod, when src0 is non-contiguous (or nb01/sizeof(float) < ne00), copy it to a contiguous f32 buffer (e.g. via ggml_cuda_cpy into a pool-allocated contiguous tensor) before the cuBLAS call, like the existing dequantize path already does for quantized src0.

Local workaround that makes gemma4 training work today: fall back GGML_OP_OUT_PROD to CPU (return false in the CUDA supports_op), letting the scheduler split the graph — correct, just slower.

[3ndetz]

Resolved — out_prod now works on CUDA, Gemma-4-E4B trains fully on GPU.

Root cause was simpler than the transposed-src0 theory: ggml_cuda_out_prod only converted quantized operands to F32 (type != F32 && type != F16), so F16 operands fell through and were read as raw float* — wrong pointer interpretation and lda = nb01/sizeof(float) = ne00/2 < m → CUBLAS_STATUS_INVALID_VALUE. Gemma-4-E4B's attention backward feeds F16 operands here.

Fix (pushed): convert any non-F32 operand (F16 included) via the existing ggml_get_to_fp32_cuda path. No CPU fallback needed anymore.

Result: 4 epochs in ~64s on an RTX 3090, val loss converges 3.18 → 2.28 → 1.93 → 1.62. Full chain (loader + dataset guard + TANH/GELU/rms backward + GELU clamp + this F16 out_prod fix) gives working full-GPU LoRA training for Gemma-4-E4B.

[3ndetz]

Added examples/lumi-serve — a single resident process that serves /completion (inference) AND /train (LoRA training) on the same loaded model, eliminating model reloads between serving and training.

Training uses a context with training=true on the shared model (weights not reloaded); the trained adapter is applied to the inference context in-process via llama_set_adapters_lora.

On Gemma-4-E4B Q4_K_M (RTX 3090): in-process /train ~11s (vs ~180s previously with the separate llama-finetune-lora binary + server restarts), inference ~1s. Verified memory effect: after training on a short conversation, a cold probe shifts behaviour accordingly; /clear_memory reverts it. This demonstrates the zero-reload on-device personalization loop the E4B fixes enable.

[gianni-cor]

Hello @3ndetz , sorry for the late reply. We rebased our fork on llama.cpp v8828. is the issue still there? Also if you wanna add lora-finetuning support for a new model architecture Vulkan and Metal support is required

[3ndetz]

Thanks for the review, @gianni-cor. Let me clear up the model identity first, since the PR title was ambiguous (my fault):

It's gemma-4-E4B (google/gemma-4-E4B-it, general.architecture = gemma4) — not gemma-3n. The "(gemma3n MatFormer)" in the title was meant to point at the architecture lineage, which reads as if the model is gemma-3n. It isn't. gemma-4-E4B uses the MatFormer + Per-Layer-Embeddings design that gemma-3n introduced (E4B = "effective 4B", ~7.5B raw; tensors like per_layer_token_embd, per_layer_proj_norm, and gemma4.attention.shared_kv_layers / gemma4.embedding_length_per_layer_input in the metadata). So this PR adds support for the gemma-4 generation that shares gemma-3n's nested/PLE architecture — not gemma-3n itself. I'll rename the title to avoid the confusion.

Validation (not just a smoke test): with these changes our build loads gemma-4-E4B (Q4_K_M) and runs full-GPU LoRA finetuning end-to-end. I trained a complete persona/belief LoRA on it over a ~25k-doc corpus (loss converges over thousands of steps), merged it with llama-export-lora, requantized to Q4_K_M, and the merged model loads and runs in LM Studio — text and vision (mmproj/image understanding) both work. So the shared-KV inference-loader fix, the TANH/GELU backward ops, and the CUDA out-prod F16→F32 fix are all exercised in a real training+deployment run.

Re the v8828 rebase: fair question — our fork is not rebased onto v8828 yet, so I can't claim the loader fix is still required there. I'll rebase onto v8828, rebuild, and re-verify whether gemma-4-E4B still needs the shared-KV loader patch (newer llama.cpp may already cover part of it), then report back here with the concrete result.

Re Vulkan/Metal: noted — the training-backward additions (TANH/GELU backward, out-prod fix) are CUDA-only at the moment; Vulkan/Metal kernels would be a separate follow-up.

[3ndetz]

Followed up on the v8828 rebase question — checked temp-8828 directly:

The inference-loading issue is resolved upstream on v8828; my loader fix is redundant there. temp-8828
already has native gemma-4-E4B support: LLM_ARCH_GEMMA4 is registered, and llama-model.cpp handles the
shared-KV layers via LLM_KV_ATTENTION_SHARED_KV_LAYERS → hparams.n_layer_kv_from_start = n_layer - n_kv_shared_layers (plus f_attention_scale = 1.0, since Gemma-4 uses scaling = 1.0). So gemma-4-E4B loads
and generates on v8828 without the loader patch from this PR — the missing tensor 'blk.N.attn_k.weight' failure
no longer reproduces there.

What may still be relevant from this PR is the training side — the unary backward ops (TANH/GELU backward),
the GELU-backward NaN clamp, and the CUDA out-prod F16→F32 fix needed for full-GPU LoRA finetuning of
gemma-4-E4B. I'll rebase just those training changes onto the latest temp-934x and confirm whether LoRA
training still needs them, then trim this PR down to the training delta (and keep it CUDA-scoped for now;
Vulkan/Metal kernels as a separate follow-up).

Thanks for the pointer — closing the inference part as obsoleted by the rebase.

[3ndetz]

Followup — built temp-8828 from source and confirmed gemma-4-E4B inference runs natively on it.

I built the temp-8828 branch (the v8828 llama.cpp rebase) locally (Windows, Ninja + CUDA arch 86, Release) and ran the resulting llama-cli on gemma-4-E4B-it (Q4_K_M GGUF; general.architecture = gemma4, MatFormer + Per-Layer-Embeddings):

llama-cli -m gemma-4-E4B-it-Q4_K_M.gguf \
  -p "Question: What is the capital of France? Answer:" -n 40 -ngl 999 -c 2048

ggml_cuda_init: found 1 CUDA device (RTX 3090, 24 GB)
...
Paris
[ Prompt: 105.5 t/s | Generation: 127.3 t/s ]

It loads with full GPU offload (-ngl 999, ~5.5 GB VRAM) and generates correctly — no "missing tensor attn_k" / loader error. So the inference loader half of this PR is redundant on temp-8828: that branch already carries LLM_ARCH_GEMMA4 + the shared-KV loader (n_layer_kv_from_start = n_layer - shared_kv_layers, f_attention_scale = 1.0).

The remaining unique value of this PR is the gemma-4 LoRA training-backward (GELU/TANH backward + CUDA outprod F16→F32 fix), which I don't see in any tetherto branch (fabric-llm-finetune has no gemma4 arch). Happy to trim this PR down to just that training delta and rebase it onto the current temp-93xx line if that's preferred.