gemma-4-E2B-FlashNorm-test

The stock-loading compatibility companion of gemma-4-E2B-FlashNorm (naming follows the FlashNorm convention: the plain -FlashNorm repo is the strict artifact, -test is the companion that loads everywhere). The gains of input_layernorm are folded into the q/k/v projections (q only on KV-shared layers) and the gains of pre_feedforward_layernorm into the gate/up projections, in float64, with those norm weights kept as all-ones tensors so the checkpoint loads in stock HuggingFace Transformers with no custom code and computes the same function as the base model:

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained('open-machine/gemma-4-E2B-FlashNorm-test')

Which norm tensors remain, and why

Norm (per layer) State here Why
input_layernorm all-ones (folded into q/k/v) removed entirely in the strict repo
pre_feedforward_layernorm all-ones (folded into gate/up) weightless in the strict repo
q_norm, k_norm unchanged per-head gains, foldable into the q/k projections (possible follow-up); the per-head division itself stays
post_attention_layernorm, post_feedforward_layernorm, post_per_layer_input_norm unchanged feed residual additions, not foldable
final model.norm unchanged lm_head is tied to the embeddings, folding would alter the embedding table
v_norm no weight tensor scale-free in Gemma 4

Validation

Float64 microcheck of the fold identity is exact to ~1e-14. At bf16, logit deviations from the base model are at the bf16 rounding level, consistent with the fp32 fold-only measurements for this architecture (max logit deviation 5.6e-4 at fp32). WikiText-2 (wikitext-2-raw-v1, test) perplexity changes are below 0.001%. Runtime cancellation of the pre-attention norm (Proposition 3) applies directly to this checkpoint via flashNorm_cancel.py; the strict repo ships that cancellation built into its modeling code.

Gemma is a trademark of Google; distributed under the same Apache 2.0 license as the base model.

Downloads last month
8
Safetensors
Model size
5B params
Tensor type
BF16
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for open-machine/gemma-4-E2B-FlashNorm-test

Finetuned
(111)
this model