Skip to content

rms_norm: fix bf16 accuracy (f32 normalize + conv_even) and parameterize epsilon#139

Open
atassis wants to merge 2 commits into
amd:develfrom
atassis:rms-norm-accuracy-and-eps
Open

rms_norm: fix bf16 accuracy (f32 normalize + conv_even) and parameterize epsilon#139
atassis wants to merge 2 commits into
amd:develfrom
atassis:rms-norm-accuracy-and-eps

Conversation

@atassis

@atassis atassis commented Jul 19, 2026

Copy link
Copy Markdown
Contributor

Summary

Two related changes to the aie2/aie2p bf16 RMSNorm kernel and its IRON operator:

  1. An accuracy fix: normalize in f32 and round once, and set the rounding mode to
    round-to-nearest-even (conv_even).
  2. Make epsilon a parameter instead of a hardcoded 1e-5f.

The default behavior is unchanged (epsilon still defaults to 1e-5). Callers opt
into a different value with RMSNorm(..., epsilon=...).

Motivation

I hit this bringing a small LLM decode (Gemma-3-270M, 18 fused layers) up on
aie2p. The generated tokens matched a host reference for the first few decode
steps and then flipped on a step where two candidate tokens were close in logit
space. The graph itself was correct: an f32 (and a faithful round-to-nearest
bf16) reference of the same graph reproduced the reference output. The device was
noisier than a correct bf16 forward should be, and the error grew roughly
linearly with depth rather than as sqrt(depth).

Two things in the RMSNorm kernel caused that:

  • inv_rms (the reciprocal RMS scalar) was cast to bf16 and then multiplied
    bf16*bf16, so the entire normalized vector was scaled by one imprecise scalar
    per norm. That is a coherent, same-direction error per norm. Over a deep
    residual stream it accumulates roughly linearly, which is what made it large
    enough to flip a near-tie argmax; unbiased round-to-nearest error would grow
    only as sqrt(N) and partly cancel.
  • The kernel set no rounding mode, so it ran under whatever a previously
    dispatched kernel left in the core's global rounding state. set_rounding is
    global per core, so in a fused graph a kernel that leaves floor/truncation set
    poisons the ones that follow.

Separately, epsilon was hardcoded to 1e-5f (Llama's value). Gemma uses 1e-6,
and it is not a rounding detail there: Gemma's deep post-feedforward sandwich
norms operate on near-zero inputs where mean(x^2) is comparable to epsilon, so
epsilon sets the normalization scale.

What changed

Commit 1, "normalize in f32 and set conv_even rounding":

  • aie_kernels/{aie2,aie2p}/rms_norm.cc: normalize through an accfloat
    accumulator (mirroring layer_norm.cc) and round once at the output; call
    set_rounding(conv_even) at both entry points.

Commit 2, "parameterize epsilon":

  • aie_kernels/{aie2,aie2p}/rms_norm.cc: rms_norm_general /
    rms_norm_bf16_vector / weighted_rms_norm take a float epsilon argument.
  • iron/operators/rms_norm/{design,design_weighted}.py: declare the np.float32
    kernel argument and pass epsilon (a per-instance compile-time constant, the
    same mechanism the axpy op uses for its scalar factor).
  • iron/operators/rms_norm/op.py: add an epsilon field (default 1e-5).
  • iron/operators/rms_norm/reference.py: accept eps and align the formula to
    the kernel, 1/sqrt(mean(x^2) + eps) rather than 1/(sqrt(mean(x^2)) + eps).

Testing

On aie2p (Strix):

  • Isolated weighted RMSNorm rel-L2 vs a host reference improved from ~0.008 to
    ~0.0025.
  • A fused 18-layer Gemma-3-270M decode that was previously flipping a token now
    matches its host reference exactly (8 of 8 greedy tokens), deterministic across
    repeated runs, with the op passing epsilon=1e-6.

The aie2 (Phoenix) kernel gets the same source change and compiles; I do not have
aie2 hardware, so it is compile-verified, not silicon-verified. The change mirrors
the aie2p one and the existing aie2 layer_norm.cc accfloat path.

atassis added 2 commits July 19, 2026 22:47
The aie2/aie2p bf16 RMSNorm kernel rounded the reciprocal-RMS scalar to bf16
(`broadcast<T>(static_cast<T>(inv_rms))`) and then did a bf16*bf16 multiply, so
the whole normalized vector was scaled by one imprecise scalar per norm. That is
a coherent (same-direction) per-norm error; over a deep residual stream (e.g. an
LLM decode with tens of layers) it accumulates roughly linearly and is large
enough to flip near-tie argmaxes, unlike the ~sqrt(N) growth of unbiased
round-to-nearest error.

Normalize in f32 via an accfloat accumulator and round once at the output
(mirrors layer_norm.cc). Also set the rounding mode to conv_even at the kernel
entry points: the kernel previously set no mode and inherited whatever a prior
kernel left in the core's global rounding state (often floor/truncation).

Verified on aie2p: isolated RMSNorm rel-L2 0.008 -> 0.0025, and a fused
Gemma-3-270M decode that was flipping a token now matches its host reference.
epsilon was hardcoded to 1e-5 in the kernel, so one build could not serve models
with different rms_norm_eps (Llama uses 1e-5, Gemma uses 1e-6). This matters:
Gemma's deep post-feedforward sandwich norms operate on near-zero inputs where
eps dominates the normalization scale.

Thread epsilon from the op through the design into the kernel as a float scalar
argument (a compile-time constant per op instance, the same mechanism the axpy
op uses for its scalar factor). Defaults to 1e-5 so existing behavior is
unchanged; a caller opts into another value via RMSNorm(..., epsilon=...).

Also align the CPU reference to the kernel formula, 1/sqrt(mean(x^2) + eps)
rather than 1/(sqrt(mean(x^2)) + eps), and thread eps through it.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant