Gemma 4 A4B 98-Expert v5 (20.8B) — v4 + Shared FFN α=1.2
A research checkpoint that takes
98e v4 and applies a
single architectural delta: every
mlp.down_proj.weight in the 30 decoder layers is scaled by
α=1.2, in BF16, before any quantization. Expert layout, router, attention, embeddings, and norms are identical to v4.
| 98e v4 | 98e v5 (this model) |
|---|
| Total params | 20.8B | ~20.8B |
| Experts per layer | 98 (30 dropped) | 98 (same drop) |
| Drop map | multi-class CD-map (max), p16 | same as v4 |
| Shared FFN α | 1.0 (none) | 1.2 on down_proj, all 30 layers |
Eval status — complete. All 9 canonical benches have landed on NVFP4A16 under
omnimergekit/eval/EVAL_PROTOCOL.md (greedy sampler,
thinking_token_budget=12288, vLLM 0.20.2 stock with
--reasoning-parser gemma4 + Fix-A reasoning_content fallback). Numbers below are NVFP4A16. ARC, GPQA, and HE+ were rescored on the stack-pinned vLLM after the Fix-A patch; the others ran on the post-revert stack on solidpc (IFEval-100 ran on stack-pinned L40 pod 37006213).
What changed vs v4
The single delta is a uniform 1.2× scaling of every layer's expert-mixture output projection:
1for layer_idx in range(30):
2 state[f"model.language_model.layers.{layer_idx}.mlp.down_proj.weight"].mul_(1.2)
That's it. The MoE router still picks the same top-8 experts (its weights are untouched), the gate+up projections that combine before down_proj are untouched, the norm scales are untouched. v5 differs from v4 only in how loudly each layer's mixed expert signal speaks into the residual stream.
Independent numerical check on layers.0.mlp.down_proj.weight (BF16 → FP32):
| abs(max) | abs(mean) | std |
|---|
| v4 | 0.398438 | 0.026069 | 0.035111 |
| v5 | 0.478516 | 0.031283 | 0.042133 |
| ratio (v5/v4) | 1.2010 | 1.2000 | 1.2000 |
The same 1.2× ratio holds for all 30 layers (sampled L0/L14/L29).
Why α=1.2?
The MoE down_proj is the only path through which the per-expert FFN output enters the residual stream. Scaling it up by α is equivalent to a learned "FFN-mixer gain" of α — a single scalar per architecture, not per-layer or per-expert. The empirical sweep (scripts/shared_alpha_sweep.py) over α ∈ {1.0, 1.05, 1.1, 1.15, 1.2, 1.25, 1.3} on a longer-smoke probe (HE-30 chat + HE+30 chat + LCB-30) picked α = 1.2 as the best non-degenerate point: above 1.2 the model gets louder but routing distribution becomes noisier; below 1.2 the recovery against v4 is marginal.
α=1.0 is v4 by construction. α=1.2 lifts the mixed-expert contribution roughly in proportion to the experts removed (98/128 ≈ 76.6%; 1/0.766 ≈ 1.305) without quite saturating that ratio — a deliberate "less than full rescaling" choice so the expert signal is amplified without overwhelming attention.
Eval — NVFP4A16, vLLM, greedy
This NVFP4A16 quant was produced with
nvidia-modelopt 0.43.0 directly on the v5 BF16 weights (i.e. the 1.2× scaling is absorbed into the per-tensor
weight_scale_2; on solidPC this measures
weight_scale_2_v5 / weight_scale_2_v4 ≈ 1.2010). Served via vLLM 0.20.2 with
--reasoning-parser gemma4,
enable_thinking=true,
thinking_token_budget=12288. Sampler: greedy (
T=0, top_p=1, top_k=0, do_sample=false) — the canonical Gemma 4 9-bench recipe.
| Bench (n) | 128e ref | 98e v4 | 98e v5 | Δ (v5 − v4) |
|---|
| ARC-Challenge-chat (1172) | 95.99% | 95.99% | 95.82% | −0.17 |
| GPQA Diamond (198) flex | 73.23% | 69.19% | 68.18% | −1.01 |
| GSM8K-100 | 91.00% | 86.00% | 91.00% | +5.00 |
| MATH-500-100 | 89.00% | 89.00% | 90.00% | +1.00 |
| AIME 2024 (30) | 36.67% | 36.67% | 70.00% | +33.33 |
| IFEval-100 (prompt_strict) | 95.00% | 93.00% | 89.00% | −4.00 |
| HumanEval-164 chat | 96.95% | 96.95% | 93.29% | −3.66 |
| HumanEval+ chat (164) | 92.07% | 91.46% | 87.20% | −4.26 |
| LCB-medium-55 v4 | 87.27% | 78.18% | 80.00% | +1.82 |
Read this carefully: α=1.2 amplifies the same signal — it does not change
which experts fire or
what they emit. It lifts benches where v4 was already on a correct reasoning trajectory but emitting at low confidence; it cannot recover benches where v4's pruning damaged the underlying expert knowledge (e.g. organic chemistry on GPQA, see
v4 card §GPQA-Diamond domain breakdown). The full picture:
strong math wins (AIME +33.33 pp, GSM8K +5.00 pp, MATH-500 +1.00 pp), a code-reasoning win (LCB-medium +1.82 pp),
regressions on code-generation (HE −3.66 pp, HE+ −4.26 pp) and instruction-following (IFEval −4.00 pp), and flat reasoning/general benches (ARC, GPQA). The pattern is consistent with α=1.2 amplifying the v4 expert mixture: it lets the math-specialist experts speak more confidently into the residual, but the same amplification adds noise to code-completion and instruction-format adherence where v4 was already competent.
Method (mechanical detail)
- Start from v4 BF16 weights (
google/gemma-4-A4B-98e-v4-it after the multi-class CD-map drop).
- Load shards, multiply every
model.language_model.layers.{i}.mlp.down_proj.weight by 1.2 in BF16 in-place.
- Save with
max_shard_size="10GB" (per memory/feedback_save_pretrained_shard_oom.md — 50 GB default cgroup-OOMs on tight pods).
- Copy
config.json, processor_config.json, tokenizer*, chat_template.jinja, expert_drop_metadata.json unchanged.
- Quantize to NVFP4A16 with
modelopt==0.43.0 (the pinned working version — see memory/feedback_modelopt_pin_0_43.md — 0.44.0 has two Gemma 4 regressions).
No retraining, no calibration, no router refit. The single .mul_(1.2) is the whole recipe.
Usage
Transformers
1from transformers import AutoModelForCausalLM, AutoTokenizer
2import torch
3
4model = AutoModelForCausalLM.from_pretrained(
5 "ManniX-ITA/gemma-4-A4B-98e-v5-it",
6 torch_dtype=torch.bfloat16,
7 device_map="auto",
8 attn_implementation="eager", # Gemma 4 head_dim=512 — FA2 not supported
9)
10tok = AutoTokenizer.from_pretrained("ManniX-ITA/gemma-4-A4B-98e-v5-it")
11
12msgs = [{"role": "user", "content": "Explain the Heisenberg uncertainty principle."}]
13inputs = tok.apply_chat_template(msgs, return_tensors="pt", add_generation_prompt=True).to(model.device)
14out = model.generate(inputs, max_new_tokens=512, do_sample=False) # greedy, canonical recipe
15print(tok.decode(out[0][inputs.shape[1]:], skip_special_tokens=True))
vLLM (NVFP4A16, canonical eval recipe)
1python -m vllm.entrypoints.openai.api_server \
2 --model ManniX-ITA/gemma-4-A4B-98e-v5-NVFP4A16 \
3 --served-model-name 98e_v5_nvfp4a16 \
4 --port 8099 \
5 --gpu-memory-utilization 0.55 \
6 --max-model-len 32768 \
7 --max-num-batched-tokens 8192 \
8 --dtype bfloat16 \
9 --trust-remote-code \
10 --reasoning-parser gemma4 \
11 --default-chat-template-kwargs '{"enable_thinking": true}'
Related Models
Recipe + Code
OmniMergeKit is the canonical home. The scripts that produced this model:
scripts/shared_alpha_apply.py — applies mlp.down_proj.weight *= α across all decoder layers (deterministic, in-place, dtype-preserving).scripts/shared_alpha_sweep.py — the longer-smoke probe over α ∈ {1.0..1.3}.omnimergekit/scripts/quantize_any.py --method nvfp4a16 — the canonical NVFP4A16 producer with modelopt==0.43.0 pinned. (Do not use one-off recipes; see memory/feedback_use_omnimergekit_canonical.md.)eval/EVAL_PROTOCOL.md — locked greedy methodology for the 9-bench suite.
License
This model inherits the
Gemma license from the base model.
Acknowledgements
- Google for the base Gemma 4 26B-A4B-it model
- The GPQA Diamond benchmark (Rein et al., 2023)
- The OmniMergeKit project for the surgery + eval toolkit
- The vLLM and modelopt teams for the NVFP4A16 serving / quantization pipeline
