Skip to content

Commit a5d7261

Browse files
karthik-0306aws-kingrj
authored andcommitted
[Inductor] Fix NaN output in LayerNorm CPU by guarding Welford variance. (pytorch#173989)
Running torch.compile on CPU with float16 precision causes LayerNorm to produce NaN outputs when the input contains Inf values. This happens because the Welford variance calculation in lowering.py hits a numerical collapse ($\infty - \infty$) in the $M2$ accumulator, which then propagates as a NaN.This PR adds an ops.where guard in var_mean_welford_ to catch these non-finite results and reset them to 0.0. This ensures that even if the input overflows the precision range, the LayerNorm output remains numerically stable, matching the behavior of the CUDA backend and Eager mode. **Fixes**: pytorch#173885 Pull Request resolved: pytorch#173989 Approved by: https://github.com/jansel
1 parent 48f9fce commit a5d7261

4 files changed

Lines changed: 49 additions & 2 deletions

File tree

‎test/inductor/test_cpu_repro.py‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6624,6 +6624,35 @@ def f(buf, idx):
66246624
actual = torch.compile(f, backend="inductor")(buf, idx)
66256625
self.assertEqual(actual, expected)
66266626

6627+
def test_layernorm_nan_with_inf_cpu_float16(self):
6628+
# https://github.com/pytorch/pytorch/issues/173885
6629+
# LayerNorm with torch.compile on CPU produces NaN when input has Inf
6630+
# values due to inf - inf = NaN in Welford variance calculation.
6631+
for dtype in _lowp_fp_dtypes:
6632+
torch._dynamo.reset()
6633+
6634+
class Model(torch.nn.Module):
6635+
def __init__(self):
6636+
super().__init__()
6637+
self.ln = torch.nn.LayerNorm(4)
6638+
6639+
def forward(self, x):
6640+
return self.ln(x)
6641+
6642+
mod = Model().to(dtype).eval()
6643+
# Create input that overflows during Welford M2 accumulation
6644+
x = torch.tensor(
6645+
[[65504.0, 65504.0, -65504.0, -65504.0]], dtype=torch.float32
6646+
).to(dtype)
6647+
6648+
with torch.no_grad():
6649+
eager_out = mod(x)
6650+
compiled_mod = torch.compile(mod, backend="inductor")
6651+
compiled_out = compiled_mod(x)
6652+
# Compiled output should match eager mode (whether Eager produces
6653+
# inf or finite numbers, it shouldn't degrade into unexpected NaNs).
6654+
self.assertEqual(eager_out, compiled_out)
6655+
66276656

66286657
if __name__ == "__main__":
66296658
from torch._inductor.test_case import run_tests

‎torch/_inductor/ir.py‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1257,7 +1257,14 @@ def welford_combine_fn(
12571257
a_mean, a_m2, a_weight = a
12581258
b_mean, b_m2, b_weight = b
12591259

1260-
delta = b_mean - a_mean
1260+
# Guard against inf - inf = NaN when both means are infinite and
1261+
# equal. This occurs during FP16/BF16 LayerNorm when inputs
1262+
# overflow to inf.
1263+
delta = ops.where(
1264+
ops.logical_and(ops.isinf(a_mean), ops.eq(a_mean, b_mean)),
1265+
ops.constant(0.0, torch.float32),
1266+
b_mean - a_mean,
1267+
)
12611268
new_weight = a_weight + b_weight
12621269
w2_over_w = b_weight / new_weight
12631270
return (

‎torch/_inductor/runtime/triton_helpers.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -262,7 +262,9 @@ def welford_reduce(value, mean, m2, weight, first_iteration):
262262

263263
@triton.jit
264264
def welford_combine(mean_1, m2_1, weight_1, mean_2, m2_2, weight_2):
265-
delta = mean_2 - mean_1
265+
# Guard against inf - inf = NaN when both means are infinite and equal.
266+
# This occurs during FP16/BF16 LayerNorm when inputs overflow to inf.
267+
delta = tl.where(mean_1 == mean_2, 0.0, mean_2 - mean_1)
266268
new_weight = weight_1 + weight_2
267269
w2_over_w = tl.where(new_weight == 0.0, 0.0, weight_2 / new_weight)
268270
return (

‎torch/csrc/inductor/cpp_prefix.h‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -192,7 +192,16 @@ Welford<T> welford_combine(
192192
if (b.index == 0) {
193193
return a;
194194
}
195+
// Guard against inf - inf = NaN when both means are infinite and equal.
196+
// This occurs during FP16/BF16 LayerNorm when inputs overflow to inf.
195197
auto delta = b.mean - a.mean;
198+
if constexpr (IsVecType<T>::value) {
199+
delta = T::blendv(delta, T(0), a.mean == b.mean);
200+
} else {
201+
if (std::isinf(a.mean) && a.mean == b.mean) {
202+
delta = T(0);
203+
}
204+
}
196205
auto a_weight = use_index ? T(a.index) : a.weight;
197206
auto b_weight = use_index ? T(b.index) : b.weight;
198207
auto new_weight = a_weight + b_weight;

0 commit comments

Comments
 (0)