gemma-4-E2B-FlashNorm

google/gemma-4-E2B with the FlashNorm weight fold applied (Proposition 1 of the FlashNorm paper): 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 set to ones. The checkpoint loads and runs in stock HuggingFace Transformers and computes the same function as the base model.

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

Validation on this fold: float64 microcheck of the fold identity is exact to ~1e-16; at bf16, logit deviations from the base model are at the bf16 rounding level, consistent with the fp32 fold-only measurements reported for this architecture (max logit deviation 5.6e-4 at fp32).

Because Gemma 4 re-normalizes queries, keys, and values per head, the pre-attention RMSNorm of this checkpoint is additionally redundant at inference and can be cancelled entirely at runtime (Proposition 3): see flashNorm_cancel.py in the transformer-tricks repo. Measured on this model family in fp32: cancellation changes logits by at most 6.8e-4 (comparable to the exact fold-only control at 5.6e-4), perplexity by 0.0006%, with greedy generations identical and HellaSwag unchanged.

Post-attention, post-feedforward, and final norms are unchanged (they feed residual additions or tied embeddings and cannot be folded). Gemma is a trademark of Google; this derivative is distributed under the same Apache 2.0 license as the base model.