[Inductor][CPU][float16] LayerNorm outputs NaN when input contains Inf

@albanD

🐛 Describe the bug

Summary:
When using torch.compile with cpu device and float16 precision, LayerNorm produces NaN outputs when the input contains Inf values. This behavior is inconsistent with the CUDA backend and Eager mode, where numerical stability is better maintained.
Observation:
In our reproduction, Mode A (CUDA float32) works correctly. In Mode B (CPU float16 with Inductor optimization), the input already contains Inf due to float16 overflow. However, instead of maintaining Inf or handling it gracefully, the compiled LayerNorm transforms these Inf values into NaN.
Comparison:
Mode A (CUDA/f32): Input OK -> Output OK
Mode B (CPU/f16): Input has Inf -> Output has NaN (Numerical collapse)

import torch
import torch.nn as nn
import os
import contextlib

@contextlib.contextmanager
def set_env(env_str: str):
    original = os.environ.copy()
    s = (env_str or "0000")[:4].ljust(4, "0")
    os.environ.update({
        "INDUCTOR_FALLBACK_RANDOM": "1" if s[0] == "1" else "0",
        "INDUCTOR_EPILOGUE_FUSION": "1" if s[1] == "1" else "0",
        "INDUCTOR_SHAPE_PADDING": "1" if s[2] == "1" else "0",
        "INDUCTOR_DYNAMIC": "1" if s[3] == "1" else "0"
    })
    yield
    os.environ.clear()
    os.environ.update(original)

def check(t, label):
    f = t.detach().cpu().float()
    print(f"{label:<15} | NaN: {str(torch.isnan(f).any().item()):<5} | Inf: {str(torch.isinf(f).any().item()):<5}")

def run():
    inp, sd = torch.load("input.pt", map_location="cpu"), torch.load("sd.pt", map_location="cpu")
    ln = nn.LayerNorm(768).eval()
    ln.load_state_dict(sd)

    print("--- Mode A (cuda/f32) ---")
    in_a = inp.cuda().float()
    check(in_a, "Input A")
    check(ln.cuda().float()(in_a), "Output A")

    print("\n--- Mode B (cpu/f16/0010) ---")
    in_b = inp.cpu().half()
    check(in_b, "Input B")
    print(f"Consistency: {torch.allclose(in_a.cpu().half(), in_b, atol=1e-5)}")

    with set_env("0010"):
        check(torch.compile(ln.cpu().half())(in_b), "Output B")

if __name__ == "__main__":
    run()

--- Mode A (cuda/f32) ---
Input A         | NaN: False | Inf: False
Output A        | NaN: False | Inf: False

--- Mode B (cpu/f16/0010) ---
Input B         | NaN: False | Inf: True 
Consistency: True
Output B        | NaN: True  | Inf: False

repro_layernorm_fp16_nan.zip

Versions

PyTorch version: 2.6.0+cu126
Is debug build: False
CUDA used to build PyTorch: 12.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04.3 LTS (x86_64)
GCC version: (Ubuntu 13.3.0-6ubuntu2~24.04) 13.3.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.39

Python version: 3.9.23 (main, Jun 5 2025, 13:40:20) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-6.8.0-90-generic-x86_64-with-glibc2.39
Is CUDA available: True
CUDA runtime version: 12.6.20
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA GeForce RTX 3090
GPU 1: NVIDIA GeForce RTX 3090

Nvidia driver version: 560.35.03
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Address sizes: 40 bits physical, 48 bits virtual
Byte Order: Little Endian
CPU(s): 48
On-line CPU(s) list: 0-47
Vendor ID: GenuineIntel
Model name: QEMU Virtual CPU version 2.5+
CPU family: 15
Model: 107
Thread(s) per core: 1
Core(s) per socket: 48
Socket(s): 1
Stepping: 1
BogoMIPS: 4190.15
Flags: fpu de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx lm constant_tsc nopl xtopology cpuid tsc_known_freq pni ssse3 fma cx16 sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c hypervisor lahf_lm abm cpuid_fault pti bmi1 avx2 bmi2 avx512f avx512dq avx512cd avx512bw avx512vl
Hypervisor vendor: KVM
Virtualization type: full
L1d cache: 1.5 MiB (48 instances)
L1i cache: 1.5 MiB (48 instances)
L2 cache: 192 MiB (48 instances)
L3 cache: 16 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-47
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit: KVM: Mitigation: VMX unsupported
Vulnerability L1tf: Mitigation; PTE Inversion
Vulnerability Mds: Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Meltdown: Mitigation; PTI
Vulnerability Mmio stale data: Unknown: No mitigations
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Vulnerable
Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2: Mitigation; Retpolines; STIBP disabled; RSB filling; PBRSB-eIBRS Not affected; BHI Retpoline
Vulnerability Srbds: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected

Versions of relevant libraries:
[pip3] numpy==2.0.2
[pip3] nvidia-cublas-cu12==12.6.4.1
[pip3] nvidia-cuda-cupti-c [pip3] nvidia-cuda-nvrtc-c [pip3] nvidia-cuda-runtime [pip3] nvidia-cudnn-cu12==9.5.1.17
[pip3] nvidia-cufft-cu12==11.3.0.4
[pip3] nvidia-curand-cu12= [pip3] nvidia-cusolver-cu1 [pip3] nvidia-cusparse-cu1 [pip3] nvidia-cusparselt-c [pip3] nvidia-nccl-cu12==2.21.5
[pip3] nvidia-nvjitlink-cu [pip3] nvidia-nvtx-cu12==12.6.77
[pip3] onnx==1.19.1
[pip3] onnxruntime==1.19.2
[pip3] open_clip_torch==3.2.0
[pip3] pytorch-lightning==0.7.1
[pip3] torch==2.6.0+cu126
[pip3] torch-geometric==2.6.1
[pip3] torchaudio==2.6.0+cu126
[pip3] torchvision==0.21.0+cu126
[pip3] triton==3.2.0
[conda] numpy [conda] nvidia-cublas-cu12 [conda] nvidia-cuda-cupti-cu12 [conda] nvidia-cuda-nvrtc-cu12 [conda] nvidia-cuda-runtime-cu12 [conda] nvidia-cudnn-cu12 [conda] nvidia-cufft-cu12 [conda] nvidia-curand-cu12 [conda] nvidia-cusolver-cu12 [conda] nvidia-cusparse-cu12 [conda] nvidia-cusparselt-cu12 [conda] nvidia-nccl-cu12 [conda] nvidia-nvjitlink-cu12 [conda] nvidia-nvtx-cu12 [conda] open-clip-torch [conda] pytorch-lightning [conda] torch [conda] torch-geometric [conda] torchaudio [conda] torchvision [conda] triton u12==12.6.80
u12==12.6.77
-cu12==12.6.77
=10.3.7.77
2==11.7.1.2
2==12.5.4.2
u12==0.6.3
12==12.6.85
2.0.2 pypi_0 pypi
12.6.4.1 pypi_0 pypi
12.6.80 pypi_0 pypi
12.6.77 pypi_0 pypi
12.6.77 pypi_0 pypi
9.5.1.17 pypi_0 pypi
11.3.0.4 pypi_0 pypi
10.3.7.77 pypi_0 pypi
11.7.1.2 pypi_0 pypi
12.5.4.2 pypi_0 pypi
0.6.3 pypi_0 pypi
2.21.5 pypi_0 pypi
12.6.85 pypi_0 pypi
12.6.77 pypi_0 pypi
3.2.0 pypi_0 pypi
0.7.1 pypi_0 pypi
2.6.0+cu126 pypi_0 pypi
2.6.1 pypi_0 pypi
2.6.0+cu126 pypi_0 pypi
0.21.0+cu126 pypi_0 pypi
3.2.0 pypi_0 pypi

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo

I'm intrested to solve the issue. The NaN collapse stems from the naive E[x^2] - (E[x])^2 variance formulation in the Inductor CPU codegen hitting precision limits. I'll be replacing this with Welford-based reduction in the C++ templates to align stability with eager mode. Starting the implementation now and attach the pr shortly .

Critical observation: In Mode B alone (CPU/float16/Inductor), when input contains Inf, the compiled LayerNorm pathologically converts Inf to NaN—this is an ABNORMAL Inf→NaN transformation that violates numerical semantics; proper handling should preserve or propagate Inf, never produce NaN from Inf inputs.

Critical observation: In Mode B alone (CPU/float16/Inductor), when input contains Inf, the compiled LayerNorm pathologically converts Inf to NaN—this is an ABNORMAL Inf→NaN transformation that violates numerical semantics; proper handling should preserve or propagate Inf, never produce NaN from Inf inputs.

Agreed. The Welford accumulator in Inductor currently hits an inf-inf case when inputs overflow to float16 Inf, which triggers this transformation. I've submitted a PR (#173989) to guard against this by catching the NaN in the variance scaling step and resetting it to 0.0, aligning it with Eager/CUDA behavior.

@karthik-0306

Critical observation: In Mode B alone (CPU/float16/Inductor), when input contains Inf, the compiled LayerNorm pathologically converts Inf to NaN—this is an ABNORMAL Inf→NaN transformation that violates numerical semantics; proper handling should preserve or propagate Inf, never produce NaN from Inf inputs.

Agreed. The Welford accumulator in Inductor currently hits an inf-inf case when inputs overflow to float16 Inf, which triggers this transformation. I've submitted a PR (#173989) to guard against this by catching the NaN in the variance scaling step and resetting it to 0.0, aligning it with Eager/CUDA behavior.

@karthik-0306 Thank you for the detailed analysis! I completely agree with your observation about the Inf-to-NaN issue. Your solution of catching and resetting the NaN to 0.0 aligns well with Eager/CUDA behavior. Great catch and well done on the PR!

@zifan6699

Critical observation: In Mode B alone (CPU/float16/Inductor), when input contains Inf, the compiled LayerNorm pathologically converts Inf to NaN—this is an ABNORMAL Inf→NaN transformation that violates numerical semantics; proper handling should preserve or propagate Inf, never produce NaN from Inf inputs.

Agreed. The Welford accumulator in Inductor currently hits an inf-inf case when inputs overflow to float16 Inf, which triggers this transformation. I've submitted a PR (#173989) to guard against this by catching the NaN in the variance scaling step and resetting it to 0.0, aligning it with Eager/CUDA behavior.

@karthik-0306 Thank you for the detailed analysis! I completely agree with your observation about the Inf-to-NaN issue. Your solution of catching and resetting the NaN to 0.0 aligns well with Eager/CUDA behavior. Great catch and well done on the PR!

Thanks @zifan6699! Glad the analysis aligns. I've linked this issue to the PR, so once the maintainers approve the CI runs, we should be able to get this landed.

[PHYSICS_REJECT] The Inf to NaN transformation is strictly working as intended and mathematically mandatory; when explicitly downcasting inputs that overflow the FP16 dynamic range to Inf, the subsequent LayerNorm operations (mean subtraction and variance accumulation) inherently invoke $Inf - Inf$ or $Inf / Inf$ arithmetic, which under strict IEEE 754 compliance (enforced by underlying CPU/FPU hardware) must definitively yield NaN; forcibly masking these hardware-level NaN propagation signals to 0.0 (as proposed in PR #173989) severely violates numerical integrity and introduces silent catastrophic corruption.

@Cheng-Logic Hi , thanks for the rigorous explanation. While I understand the FP16 limits, the core issue from a user perspective is the silent numerical divergence between Eager and Inductor.
Safe Mode Support: For users prioritizing consistency over peak throughput, could we provide an opt-in "safe mode" that uses more conservative algorithms (like Welford) in fused kernels to avoid inf - inf?
Consistent Error Exposure: Our core concern is that torch.compile changes the failure mode. If Eager mode handles this gracefully (maintaining Inf), Inductor should aim for behavioral parity instead of collapsing into NaN.
If PR #173989 is not the right path, what is the recommended way to ensure stable and consistent behavior across backends for these edge cases?

[PHYSICS_REJECT] The Inf to NaN transformation is strictly working as intended and mathematically mandatory; when explicitly downcasting inputs that overflow the FP16 dynamic range to Inf, the subsequent LayerNorm operations (mean subtraction and variance accumulation) inherently invoke I n f − I n f or I n f / I n f arithmetic, which under strict IEEE 754 compliance (enforced by underlying CPU/FPU hardware) must definitively yield NaN; forcibly masking these hardware-level NaN propagation signals to 0.0 (as proposed in PR #173989) severely violates numerical integrity and introduces silent catastrophic corruption.

Hi @Cheng-Logic, thank you for the feedback regarding IEEE 754 compliance. I completely agree that inf - inf is mathematically undefined.However, the primary motivation for this PR is Behavioral Parity. Currently, both Eager mode and the CUDA backend handle these overflows by propagating Inf or stabilizing the variance, while the Inductor CPU path collapses to NaN. This creates a 'compiler-induced failure' where switching to torch.compile changes the numerical outcome of the model.My goal with #173989 is to align Inductor with the existing Eager/CUDA semantics. If we strictly enforce the NaN transformation only in this specific path, we break the 'drop-in replacement' promise of the compiler. Happy to discuss if there’s a preferred way to implement this guard that satisfies both numerical intent and backend consistency

[Inductor] Fix NaN output in LayerNorm CPU by guarding Welford varian…

Update: I verified that this issue is still reproducible on both PyTorch 2.10.0+cu126 and PyTorch 2.12.1+cu126.

Only the PyTorch version/environment was changed; the repro itself was otherwise unchanged.

The observed output is the same on both versions:

--- Mode A (cuda/f32) ---
Input A         | NaN: False | Inf: False
Output A        | NaN: False | Inf: False

--- Mode B (cpu/f16/0010) ---
Input B         | NaN: False | Inf: True 
Consistency: True
Output B        | NaN: True  | Inf: False

[Inductor] Fix NaN output in LayerNorm CPU by guarding Welford varian…

Navigation Menu

Repository navigation

[Inductor][CPU][float16] LayerNorm outputs NaN when input contains Inf #173885

🐛 Describe the bug

Versions

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[Inductor][CPU][float16] LayerNorm outputs NaN when input contains Inf #173885

Description

🐛 Describe the bug

Versions

Activity

karthik-0306 commented on Jan 30, 2026

zifan6699 commented on Feb 3, 2026

karthik-0306 commented on Feb 3, 2026

zifan6699 commented on Feb 4, 2026

karthik-0306 commented on Feb 5, 2026

chenglogic commented on Feb 28, 2026

zifan6699 commented on Mar 3, 2026

karthik-0306 commented on Mar 4, 2026

hiahu329 commented on Jul 25, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions