The rounding quirk hiding inside RMSNorm

A type conversion in our RMSNorm implementation did nothing. The value was already float32, so converting it to float32 left it unchanged. Our tests passed.

With bfloat16, the same line had a job to do. It rounded an intermediate result before multiplying by a learned weight. Our reference kept that result in float32 until after the multiplication. The equations matched; the two implementations were making a different numerical choice.

That is a useful detail to know when moving a model between platforms or evaluating a quantization. Before attributing a discrepancy to the quantized weights, it helps to check what the surrounding operations are doing. We followed this one through a small experiment and the source code of several implementations.

A familiar operation

RMSNorm (root mean square normalization) keeps the scale of an activation vector under control. It appears in many language models, including Llama and Gemma. For a vector with $n$ elements, it first divides each value by the root mean square, with a small $\epsilon$ added for numerical stability:

$$\hat{x}_i = \frac{x_i}{\sqrt{\frac{1}{n}\sum_{j=1}^{n} x_j^2 + \epsilon}}$$

A learned weight, often called the gain, then scales each element:

$$y_i = \hat{x}_i \cdot \gamma_i$$

The formula leaves one practical question unanswered: how much precision do we keep along the way?

Our example uses bfloat16, or bf16, a 16-bit floating-point format commonly used for model weights and activations. It covers roughly the same range of magnitudes as float32, but has fewer bits for distinguishing nearby values. Moving a float32 result into bf16 can therefore round it, even when the number is far from overflowing.

Two places to round

Suppose normalization has produced $\hat{x}_i$ in float32, and the gain $\gamma_i$ is already stored in bf16. Two conventions are possible.

Round once: multiply the normalized value by the gain in float32, then round the result to bf16 for storage.

$$y_i = \operatorname{bf16}(\hat{x}_i \cdot \gamma_i)$$

Round twice: first round the normalized value to bf16, then multiply by the gain and round the product to bf16.

$$y_i = \operatorname{bf16}(\operatorname{bf16}(\hat{x}_i) \cdot \gamma_i)$$

Here, “once” and “twice” count the roundings to bf16 after normalization. The float32 arithmetic has its own rounding too. These expressions describe where precision is lost, without prescribing how a processor executes the multiplication.

The extra pair of parentheses is small enough to miss in a code review. With float32 inputs and outputs, that intermediate cast is a no-op. A float32-only test can pass without testing the choice at all.

A difference you can reproduce

We can isolate the choice without running a model. The following NumPy helper rounds finite float32 values to bf16 precision, using round-to-nearest-even. It keeps those rounded values in float32 arrays so we can inspect them. It is sufficient for this example, rather than a complete bf16 implementation for special values such as NaNs.

import numpy as np

def bf16(value):
    # Round finite float32 values to bf16 precision.
    value = np.asarray(value, dtype=np.float32)
    bits = value.view(np.uint32)
    rounded = bits + np.uint32(0x7FFF) + ((bits >> 16) & np.uint32(1))
    return (rounded & np.uint32(0xFFFF0000)).view(np.float32)

def round_once(x_hat, gamma):
    return bf16(x_hat * bf16(gamma))

def round_twice(x_hat, gamma):
    return bf16(bf16(x_hat) * bf16(gamma))

x_hat = np.float32(1.1)
gamma = bf16(1.05)  # The stored gain is actually 1.046875.

print(round_once(x_hat, gamma))   # 1.1484375
print(round_twice(x_hat, gamma))  # 1.15625

A few examples show the same effect. The gain column contains the actual bf16 values used in the multiplication:

The extra intermediate rounding can move the final result in either direction. In these examples the outputs are neighbouring bf16 values.

How often do they disagree?

For a larger illustration, we sampled two million synthetic values for $\hat{x}$ from a standard normal distribution and gains uniformly between 0.9 and 1.1. These are stand-ins for the values entering the final multiplication, not activations collected from a model. The experiment isolates the cast; it does not compute the preceding RMS normalization.

Run this after the first code block:¹

rng = np.random.default_rng(42)
x_hat = rng.standard_normal(2_000_000).astype(np.float32)
gamma = rng.uniform(0.9, 1.1, x_hat.size).astype(np.float32)

once = round_once(x_hat, gamma)
twice = round_twice(x_hat, gamma)
different = np.count_nonzero(once != twice)
nonzero = once != 0
relative = np.abs(once[nonzero] - twice[nonzero]) / np.abs(once[nonzero])

print(f"Different outputs: {different:,} / {x_hat.size:,}")
print(f"Share: {100 * different / x_hat.size:.2f}%")
print(f"Largest relative difference: {100 * relative.max():.5f}%")

The outputs differ in 495,539 cases, or 24.78%. The largest relative difference in this sample is 0.78125%, measured against the round-once result.

That is a difference between two calculations, not a measured loss of model quality. It does not mean a quarter of a model's answers change, or that a fixed percentage of error accumulates at every layer. Those questions need tests on the model itself.

Where implementations differ

We found both conventions in established implementations. This is a comparison of particular paths, with sources and scope recorded below, rather than a rule for every operation a framework offers.

The Transformer Engine kernel also shows why the boundary matters for FP8, an 8-bit floating-point format. In its per-tensor FP8 output path, normalization, gain multiplication and output scaling stay in float32 before the final conversion to FP8. There is no intermediate bf16 result to round again. This observation concerns that kernel, not every FP8 recipe or backend.³

Follow the cast in MLX

The boundary is visible in these lines of MLX's Metal kernel. Here is the output expression:

out[i] =
    w[w_stride * i] * static_cast<T>(thread_x[i] * local_inv_mean[0]);

The declarations above it supply the missing context: thread_x and local_inv_mean hold float32 values, while the gain w and output out use T. For the bf16 path, T is bf16:

1. The product inside static_cast<T>(...) is float32. The cast rounds that normalized value to bf16.
2. The gain and the cast result are now both bf16. Their multiplication produces a bf16 result, rounding again. Storing that result does not add a third rounding.

The looped kernel repeats the same pattern, using x[r + i] instead of thread_x[i]. Its multiplication by the float32 inverse RMS also produces float32 before the explicit cast.

We implemented an opt-in variant in Precisit's MLX fork. Here is how it selects between the two expressions:

out[i] = PRECISE
    ? static_cast<T>(
          static_cast<float>(w[w_stride * i]) * thread_x[i] *
          local_inv_mean[0])
    : w[w_stride * i] * static_cast<T>(thread_x[i] * local_inv_mean[0]);

With PRECISE enabled, both multiplications stay in float32 until the final cast. "Round once" refers to conversion to bf16 here; float32 arithmetic still has its own rounding. The expression also groups the products as (gain * x) * inverse_rms, rather than gain * (x * inverse_rms). Removing the bf16 intermediate does not guarantee identical float32 intermediates.

Which convention belongs to the model?

Hugging Face's Llama and Gemma implementations document this distinction. In the reference versions linked below, Llama casts the normalized value before multiplying by the weight. Gemma multiplies in float32 and casts afterward. Gemma also uses $1 + \gamma$ as its gain, which is a separate formula choice.⁶

So the early cast in MLX follows the Llama convention at this boundary. Replacing it with a later cast would change that behaviour. Fewer rounding steps can bring a calculation closer to its real-valued formula while moving it away from the implementation a checkpoint is meant to reproduce.

This distinction was already discussed in the Transformers project in 2024.⁶ Finding it in our own tests was a reminder to read the model code alongside the framework primitive. A familiar operation name can hide an unfamiliar assumption.

Give the test the same precision as the model

For a model port or a quantization comparison, a useful check is to feed the same already-rounded inputs and gains to both implementations. Use the dtype intended for deployment, keep the normalization epsilon fixed, and inspect the intermediate casts. That separates an operation's rounding choices from differences introduced while preparing its inputs.

If the question is fidelity, compare against the model's intended reference convention. If the question is numerical error, use a higher-precision calculation of the same rounded inputs as an oracle. They answer different questions.

Beyond the multiplication grouping shown above, reduction order and reciprocal-square-root calculations can also differ. Matching the bf16 rounding boundary therefore does not guarantee bit-identical output. Backward passes need their own checks too.

Conclusion: know your numbers

Moving a model between platforms also moves it between numerical conventions. To make sense of a comparison, we need to know which precision each step uses, where values are rounded and which reference we intend to match. A small output difference is a clue to investigate; its effect on model quality needs a separate measurement.

Huawei's Ascend Transformer Boost (ATB) offers a useful precedent for making such choices visible. Its RMSNorm precisionMode selects float32 or float16 intermediate calculations for a supported fp16 path. This is separate from the Llama/Gemma modelType setting, and has a different scope from the bf16 rounding boundary examined here.⁷

Our MLX fork makes the later cast selectable through MLX_RMS_PRECISE=1. The change covers all four forward output expressions in the single-row and looped Metal kernels, including partial rows. It preserves the default behaviour and leaves the backward kernels unchanged.⁸

When using a build of the fork, set the variable before starting the process. The choice is read at the first Metal RMSNorm evaluation and retained for that process. It is a process-wide setting; the upstream revision we inspected does not expose it. Record the chosen convention alongside the model reference when comparing results.

Our float32 test was answering a narrower question than we had realised. The lower-precision case made that visible. Knowing our numbers means understanding how they came to be, as well as what they say. Next time they almost agree, we have one more place to look.

Footnotes

1. Reproduced with NumPy 2.5.3. default_rng(42) draws float64 samples, which the code then converts to float32; the gain is further rounded to bf16 by the helper. Relative difference is abs(once - twice) / abs(once), excluding zero denominators. The supplied code specifies the distribution, draw order and rounding. This is a synthetic arithmetic experiment, not a hardware or model benchmark. ↩

2. In our recorded comparison with PyTorch 2.14.0 on CPU, torch.nn.RMSNorm and torch.rms_norm matched the float32-composed, final-cast reference for the tested bf16 inputs (4,096 rows, width 512). This does not establish the behaviour of every shape or CUDA path. PyTorch RMSNorm documentation. ↩

3. NVIDIA Transformer Engine, rmsnorm_fwd_kernels.cuh, commit 5f6105b. The inspected kernel uses float32 compute_t; its per-tensor FP8 branch scales before the output cast. Source inspection, not a GPU execution result here. ↩ ↩²

4. Huawei CANN ops-nn, norm/rms_norm/op_kernel/rms_norm_split_d.h, commit fac671fe. In the inspected bf16 ComputeY path, the normalized value and gain are multiplied in float32 before the output cast. Source inspection, not an Ascend hardware test. ↩

5. MLX, rms_norm.metal, commit 6c0f02a. The Metal kernel casts the normalized value to T before multiplying by the gain. The table concerns bf16 input, gain and output. ↩

6. Hugging Face Transformers v4.44.2, commit 17489028: LlamaRMSNorm and GemmaRMSNorm. The comparison assumes weights stored in the input dtype. Gemma's source explicitly points to the earlier discussion of the rounding difference. ↩ ↩²

7. Huawei ATB, RmsNormParam::NormParam, commit 34c40fd9, and the RMSNorm API documentation. HIGH_PRECISION_MODE is the default and uses float32 intermediates; HIGH_PERFORMANCE_MODE uses float16. In the inspected header, the precision selector supports float16 inputs, excludes quantized operation, and cannot be combined with the non-default modelType or rstd modes. It is a related API design, not evidence of a selectable bf16 round-once/round-twice boundary or bitwise parity with another backend. These are documented settings; we did not benchmark the modes on Ascend hardware. ↩

8. Precisit's public MLX fork, commit 4121382 on precise-norm. A PRECISE template parameter selects the forward expression. The Metal dispatch code reads MLX_RMS_PRECISE once and chooses the corresponding kernel. The option applies to RMSNorm's Metal forward path; it does not add LayerNorm or CUDA variants. ↩