Repository navigation
[Bug] [Relax] DecomposeOpsForTraining loses FP32 accumulation in FP16 LayerNorm #20909
Description
Expected behavior
Applying relax.transform.DecomposeOpsForTraining() to an FP16 LayerNorm should preserve its numerical behavior and FP16 output interface.
For an input of shape (1, 4096) alternating between -0.25 and 0.25, with gamma equal to one and beta equal to zero, the FP64 NumPy reference rounded to FP16 is [-1, 1, -1, 1, ...]. Compiling the original LayerNorm produces that result.
Actual behavior
The decomposed program returns [-1.4140625, 1.4140625, ...] instead. The maximum absolute error is 0.4140625, about 41.4% of the expected magnitude. Both original and transformed outputs are finite float16[1,4096] tensors.
DecomposeLayerNorm builds mean and variance directly on the FP16 data. After LegalizeOps, their reduction buffers use FP16, including the variance accumulator T_multiply_red. The original LayerNorm implementation promotes its reductions to FP32. Explicitly promoting the internal LayerNorm computation to FP32 and casting its result back to FP16 restores the expected result.
The FP32-input control also passes. Both modules use the ordinary LLVM compilation pipeline in the script below.
Environment
- TVM:
0.26.dev0+source.8f328e8 - TVM source snapshot:
8f328e802cfe5e41fcc8f5c17e7582b1c28bfce4 - OS: Ubuntu 22.04.5 LTS, x86_64; Linux
6.8.0-52-generic - CPU: 13th Gen Intel Core i9-13900K
- Python: 3.11.16; LLVM: 15.0.7
- Target:
llvm; Relax VM on CPU
Steps to reproduce
Save as repro.py and run python repro.py. The controls pass before the final assertion detects the decomposed result.
import numpy as np
import tvm
from tvm import relax as R
N = 4096
def make_module(dtype="float16", fp32_internal=False):
params = [
R.Var("x", R.TensorType([1, N], dtype)),
R.Var("gamma", R.TensorType([N], dtype)),
R.Var("beta", R.TensorType([N], dtype)),
]
bb = R.BlockBuilder()
with bb.function("main", params):
with bb.dataflow():
values = params
if fp32_internal:
values = [bb.emit(R.op.astype(v, "float32")) for v in params]
y = bb.emit(R.op.nn.layer_norm(
*values, axes=[-1], epsilon=1e-5))
if fp32_internal:
y = bb.emit(R.op.astype(y, dtype))
out = bb.emit_output(y)
bb.emit_func_output(out)
return bb.get()
def run(module, arrays):
assert R.analysis.check_well_formed(module)
with tvm.transform.PassContext(opt_level=3):
executable = tvm.compile(module, target="llvm")
vm = R.VirtualMachine(executable, tvm.cpu())
return vm["main"](*[
tvm.runtime.tensor(a.copy(), tvm.cpu()) for a in arrays
]).numpy()
x = np.tile(np.array([-0.25, 0.25], dtype="float16"),
N // 2).reshape(1, N)
arrays = [x, np.ones(N, dtype="float16"),
np.zeros(N, dtype="float16")]
data = x.astype("float64")
mean = data.mean(axis=-1, keepdims=True)
var = ((data - mean) ** 2).mean(axis=-1, keepdims=True)
reference = ((data - mean) / np.sqrt(var + 1e-5)).astype("float16")
decompose = R.transform.DecomposeOpsForTraining()
before = make_module()
after = decompose(before)
off = run(before, arrays)
on = run(after, arrays)
control = run(decompose(make_module(fp32_internal=True)), arrays)
fp32_arrays = [a.astype("float32") for a in arrays]
fp32 = run(decompose(make_module("float32")), fp32_arrays)
print("OFF:", off[0, :4], flush=True)
print("ON:", on[0, :4], flush=True)
print("FP32 internal control:", control[0, :4], flush=True)
print("FP32 input control:", fp32[0, :4], flush=True)
print("max absolute error:",
np.max(np.abs(on.astype("float32") -
reference.astype("float32"))), flush=True)
assert off.dtype == on.dtype == control.dtype == np.dtype("float16")
np.testing.assert_array_equal(off, reference)
np.testing.assert_array_equal(control, reference)
np.testing.assert_allclose(
fp32, (data / np.sqrt(0.0625 + 1e-5)).astype("float32"),
rtol=1e-5, atol=1e-6)
np.testing.assert_allclose(on, reference, rtol=0.02, atol=0.002)Output before the final assertion:
OFF: [-1. 1. -1. 1.]
ON: [-1.414 1.414 -1.414 1.414]
FP32 internal control: [-1. 1. -1. 1.]
FP32 input control: [-0.99992 0.99992 -0.99992 0.99992]
max absolute error: 0.4140625
The final assert_allclose fails with rtol=0.02 and atol=0.002.
Triage
- needs-triage
Activity
zzznnn-1 commentedon Nov 7, 2026
Hi, I'd like to work on this issue. I'll start by reproducing it on the current main branch, then investigate preserving FP32 accumulation in the FP16 LayerNorm decomposition while keeping the FP16 output dtype, and add regression tests.
If no one is already working on it, could you please assign it to me? Thanks!
pinpinpoo commentedon Nov 8, 2026
Hi @zzznnn-1,
Thanks for investigating and fixing this issue!
I reviewed your changes, and the FP32 promotion strategy for the normalization computation looks consistent with the expected behavior. Casting the normalized result back to FP16 before scale and center also preserves the original FP16 output interface.
I also appreciate the additional structural and numerical regression tests, particularly the coverage of different scale and center combinations.
Thanks again for addressing this issue and providing the fix so quickly!