pytorch - 💡(How to fix) Fix torch.compile LayerNorm produces NaN while eager is finite on valid float32 input [6 comments, 4 participants]

🐛 Describe the bug

I found a mismatch between eager and compiled execution for a single LayerNorm layer.

I isolated the problem to a minimal setup:

Input tensor is finite (float32), no NaN / no Inf LayerNorm weights are finite (float32) Eager output is finite Compiled output (torch.compile) contains NaN Observed logs:

input min=8.773236569945328e+31, max=3.4028234663852886e+38 weight min=0.04186137020587921, max=0.25266674160957336 input: has_nan=False, has_inf=False eager_output: has_nan=False, has_inf=False compile_output: has_nan=True, has_inf=False This suggests a compile-path numerical issue in LayerNorm, since the same input and weights are finite and eager remains stable.

from pathlib import Path
import torch
import torch._inductor.config as inductor_config
BASE_DIR = Path(__file__).resolve().parent
INPUT_PATH = BASE_DIR / "input.pt"
STATE_DICT_PATH = BASE_DIR / "sd.pt"
ORIG_DEVICE = "cuda"
ORIG_DTYPE = torch.float32
SWITCH_DEVICE = "cuda"
SWITCH_DTYPE = torch.float32
def build_layer(state_dict: dict, device: str, dtype: torch.dtype, compile_mode: bool) -> torch.nn.Module:
    weight = state_dict["weight"]
    layer = torch.nn.LayerNorm(normalized_shape=weight.shape[0], eps=1e-5, elementwise_affine=True)
    layer.load_state_dict(state_dict)
    layer = layer.to(device=device, dtype=dtype).eval()
    if compile_mode:
        layer = torch.compile(layer, dynamic=False)
    return layer


def tensor_nan_inf(name: str, tensor: torch.Tensor):
    has_nan = bool(torch.isnan(tensor).any().item())
    has_inf = bool(torch.isinf(tensor).any().item())
    print(f"[INFO] {name}: has_nan={has_nan}, has_inf={has_inf}")

def main():
    state_dict = torch.load(STATE_DICT_PATH, map_location="cpu")
    input_tensor = torch.load(INPUT_PATH, map_location="cpu").float()
    weight = state_dict["weight"].float()
    print(f"[INFO] input min={float(input_tensor.min().item())}, max={float(input_tensor.max().item())}")
    print(f"[INFO] weight min={float(weight.min().item())}, max={float(weight.max().item())}")
    tensor_nan_inf("input", input_tensor)
    eager_layer = build_layer(state_dict, ORIG_DEVICE, ORIG_DTYPE, compile_mode=False)
    orig_fallback = inductor_config.fallback_random
    try:
        inductor_config.fallback_random = True
        compile_layer = build_layer(state_dict, SWITCH_DEVICE, SWITCH_DTYPE, compile_mode=True)
    finally:
        inductor_config.fallback_random = orig_fallback
    with torch.no_grad():
        eager_output = eager_layer(input_tensor.to(device=ORIG_DEVICE, dtype=ORIG_DTYPE))
    with torch.no_grad():
        compile_output = compile_layer(input_tensor.to(device=SWITCH_DEVICE, dtype=SWITCH_DTYPE))
    eager_output = eager_output.detach().float().cpu()
    compile_output = compile_output.detach().float().cpu()
    tensor_nan_inf("eager_output", eager_output)
    tensor_nan_inf("compile_output", compile_output)


if __name__ == "__main__":
    main()

[INFO] input min=8.773236569945328e+31, max=3.4028234663852886e+38
[INFO] weight min=0.04186137020587921, max=0.25266674160957336
[INFO] input: has_nan=False, has_inf=False
[INFO] eager_output: has_nan=False, has_inf=False
[INFO] compile_output: has_nan=True, has_inf=False

Versions

PyTorch version: 2.10.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.10.19 | packaged by conda-forge | (main, Jan 26 2026, 23:45:08) [GCC 14.3.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: 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 Is XPU available: False HIP runtime version: N/A MIOpen runtime version: N/A Is XNNPACK available: True Caching allocator config: N/A

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==1.26.4 [pip3] nvidia-cublas-cu12==12.6.4.1 [pip3] nvidia-cuda-cupti-cu12==12.6.80 [pip3] nvidia-cuda-nvrtc-cu12==12.6.77 [pip3] nvidia-cuda-runtime-cu12==12.6.77 [pip3] nvidia-cudnn-cu12==9.10.2.21 [pip3] nvidia-cufft-cu12==11.3.0.4 [pip3] nvidia-curand-cu12==10.3.7.77 [pip3] nvidia-cusolver-cu12==11.7.1.2 [pip3] nvidia-cusparse-cu12==12.5.4.2 [pip3] nvidia-cusparselt-cu12==0.7.1 [pip3] nvidia-nccl-cu12==2.27.5 [pip3] nvidia-nvjitlink-cu12==12.6.85 [pip3] nvidia-nvtx-cu12==12.6.77 [pip3] onnxruntime-gpu==1.23.2 [pip3] optree==0.18.0 [pip3] pytorch-triton==3.2.0+git4b3bb1f8 [pip3] torch==2.10.0+cu126 [pip3] torchaudio==2.11.0.dev20260127+cu126 [pip3] torchvision==0.25.0+cu126 [pip3] triton==3.6.0+git9844da95 [conda] numpy 1.26.4 pypi_0 pypi [conda] nvidia-cublas-cu12 12.6.4.1 pypi_0 pypi [conda] nvidia-cuda-cupti-cu12 12.6.80 pypi_0 pypi [conda] nvidia-cuda-nvrtc-cu12 12.6.77 pypi_0 pypi [conda] nvidia-cuda-runtime-cu12 12.6.77 pypi_0 pypi [conda] nvidia-cudnn-cu12 9.10.2.21 pypi_0 pypi [conda] nvidia-cufft-cu12 11.3.0.4 pypi_0 pypi [conda] nvidia-curand-cu12 10.3.7.77 pypi_0 pypi [conda] nvidia-cusolver-cu12 11.7.1.2 pypi_0 pypi [conda] nvidia-cusparse-cu12 12.5.4.2 pypi_0 pypi [conda] nvidia-cusparselt-cu12 0.7.1 pypi_0 pypi [conda] nvidia-nccl-cu12 2.27.5 pypi_0 pypi [conda] nvidia-nvjitlink-cu12 12.6.85 pypi_0 pypi [conda] nvidia-nvtx-cu12 12.6.77 pypi_0 pypi [conda] optree 0.18.0 pypi_0 pypi [conda] pytorch-triton 3.2.0+git4b3bb1f8 pypi_0 pypi [conda] torch 2.10.0+cu126 pypi_0 pypi [conda] torchaudio 2.11.0.dev20260127+cu126 pypi_0 pypi [conda] torchvision 0.25.0+cu126 pypi_0 pypi [conda] triton 3.6.0+git9844da95 pypi_0 pypi

cc @chauhang @penguinwu