rms_norm: fix bf16 accuracy (f32 normalize + conv_even) and parameterize epsilon#139
Open
atassis wants to merge 2 commits into
Open
rms_norm: fix bf16 accuracy (f32 normalize + conv_even) and parameterize epsilon#139atassis wants to merge 2 commits into
atassis wants to merge 2 commits into
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Two related changes to the aie2/aie2p bf16 RMSNorm kernel and its IRON operator:
round-to-nearest-even (
conv_even).epsilona parameter instead of a hardcoded1e-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 multipliedbf16*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.
dispatched kernel left in the core's global rounding state.
set_roundingisglobal per core, so in a fused graph a kernel that leaves floor/truncation set
poisons the ones that follow.
Separately,
epsilonwas hardcoded to1e-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, soepsilon 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 anaccfloataccumulator (mirroring
layer_norm.cc) and round once at the output; callset_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_normtake afloat epsilonargument.iron/operators/rms_norm/{design,design_weighted}.py: declare thenp.float32kernel argument and pass
epsilon(a per-instance compile-time constant, thesame mechanism the axpy op uses for its scalar factor).
iron/operators/rms_norm/op.py: add anepsilonfield (default 1e-5).iron/operators/rms_norm/reference.py: acceptepsand align the formula tothe kernel,
1/sqrt(mean(x^2) + eps)rather than1/(sqrt(mean(x^2)) + eps).Testing
On aie2p (Strix):
~0.0025.
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.ccaccfloat path.