gemma-4-A4B-98e-v5-it

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.

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:

for layer_idx in range(30):
    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):

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.

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)

1. Start from v4 BF16 weights (google/gemma-4-A4B-98e-v4-it after the multi-class CD-map drop).
2. Load shards, multiply every model.language_model.layers.{i}.mlp.down_proj.weight by 1.2 in BF16 in-place.
3. Save with max_shard_size="10GB" (per memory/feedback_save_pretrained_shard_oom.md — 50 GB default cgroup-OOMs on tight pods).
4. Copy config.json, processor_config.json, tokenizer*, chat_template.jinja, expert_drop_metadata.json unchanged.
5. 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

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model = AutoModelForCausalLM.from_pretrained(
    "ManniX-ITA/gemma-4-A4B-98e-v5-it",
    torch_dtype=torch.bfloat16,
    device_map="auto",
    attn_implementation="eager",  # Gemma 4 head_dim=512 — FA2 not supported
)
tok = AutoTokenizer.from_pretrained("ManniX-ITA/gemma-4-A4B-98e-v5-it")

msgs = [{"role": "user", "content": "Explain the Heisenberg uncertainty principle."}]
inputs = tok.apply_chat_template(msgs, return_tensors="pt", add_generation_prompt=True).to(model.device)
out = model.generate(inputs, max_new_tokens=512, do_sample=False)  # greedy, canonical recipe
print(tok.decode(out[0][inputs.shape[1]:], skip_special_tokens=True))

vLLM (NVFP4A16, canonical eval recipe)

python -m vllm.entrypoints.openai.api_server \
    --model ManniX-ITA/gemma-4-A4B-98e-v5-NVFP4A16 \
    --served-model-name 98e_v5_nvfp4a16 \
    --port 8099 \
    --gpu-memory-utilization 0.55 \
    --max-model-len 32768 \
    --max-num-batched-tokens 8192 \
    --dtype bfloat16 \
    --trust-remote-code \
    --reasoning-parser gemma4 \
    --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