[cuda] Numerical Inconsistency: NaN in Eager CUDA vs Valid Output in Compiled CPU for nn.LSTM
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.5k
- PR merge metrics
- PR metrics pending
Description
### 🐛 Describe the bug
I found a numerical stability issue where `nn.LSTM` produces NaNs when running on CUDA (Eager mode), but works correctly on CPU (Compiled with Inductor) using the same input and weights. The input tensor is verified to be valid (no NaNs or Infs), and the relative L2 difference between inputs across devices is zero.
- Mode A (Eager CUDA): Produces NaN output.
- Mode B (Compiled CPU): Produces valid numerical output.
This divergence suggests a potential stability issue or bug in the CUDA kernel implementation for LSTM under specific input conditions.
```python
import torch
import torch.nn as nn
import os
import contextlib
@contextlib.contextmanager
def set_inductor_env_vars(env_str: str):
original_env = os.environ.copy()
try:
s = (env_str or "0000")[:4].ljust(4, "0")
os.environ["INDUCTOR_FALLBACK_RANDOM"] = "1" if s[0] == "1" else "0"
os.environ["INDUCTOR_EPILOGUE_FUSION"] = "1" if s[1] == "1" else "0"
os.environ["INDUCTOR_SHAPE_PADDING"] = "1" if s[2] == "1" else "0"
os.environ["INDUCTOR_DYNAMIC"] = "1" if s[3] == "1" else "0"
yield
finally:
os.environ.clear()
os.environ.update(original_env)
def get_rel_l2(a, b):
a = a.detach().cpu().float()
b = b.detach().cpu().float()
diff_norm = torch.norm(a - b)
base_norm = torch.norm(a) + 1e-12
return (diff_norm / base_norm).item()
def reproduce_lstm_nan_comparison():
bundle_path = "bundle.pt"
weight_path = "lstm1_weights.pt"
bundle = torch.load(bundle_path, map_location="cpu")
mini_sd = torch.load(weight_path, map_location="cpu")
x_input = bundle["input_cpu_fp32"]
print("\n>>> Running Mode A (cuda/float32/0000) ...")
lstm_a = nn.LSTM(input_size=50, hidden_size=50, num_layers=1, batch_first=True, bidirectional=False).to("cuda").to(torch.float32)
lstm_a.load_state_dict(mini_sd)
lstm_a.eval()
with torch.no_grad():
x_a = x_input.to("cuda").to(torch.float32)
has_nan_input_a = torch.isnan(x_a).any().item()
output_a, _ = lstm_a(x_a)
has_nan_a = torch.isnan(output_a).any().item()
print(f"Mode A input has NaN: {has_nan_input_a}, output has NaN: {has_nan_a}")
print("\n>>> Running Mode B (cpu/float32/1000) ...")
lstm_b = nn.LSTM(input_size=50, hidden_size=50, num_layers=1, batch_first=True, bidirectional=False).to("cpu").to(torch.float32)
lstm_b.load_state_dict(mini_sd)
lstm_b.eval()
with set_inductor_env_vars("1000"):
run_b = torch.compile(lstm_b)
with torch.no_grad():
x_b = x_input.to("cpu").to(torch.float32)
has_nan_input_b = torch.isnan(x_b).any().item()
output_b, _ = run_b(x_b)
has_nan_b = torch.isnan(output_b).any().item()
print(f"Mode B input has NaN: {has_nan_input_b}, output has NaN: {has_nan_b}")
print("\n" + "="*50)
print(f"Input Relative Difference (Rel L2): {get_rel_l2(x_a, x_b):.4e}")
print("="*50)
if __name__ == "__main__":
reproduce_lstm_nan_comparison()
```
```
>>> Running Mode A (cuda/float32/0000) ...
Mode A input has NaN: False, output has NaN: True
>>> Running Mode B (cpu/float32/1000) ...
Mode B input has NaN: False, output has NaN: False
==================================================
Input Relative Difference (Rel L2): 0.0000e+00
==================================================
```
[lstm_nan_repro.zip.zip](https://github.com/user-attachments/files/24844808/lstm_nan_repro.zip.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
Versions of relevant libraries:
[pip3] numpy==2.0.2
[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.5.1.17
[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.6.3
[pip3] nvidia-nccl-cu12==2.21.5
[pip3] nvidia-nvjitlink-cu12==12.6.85
[pip3] nvidia-nvtx-cu12==12.6.77
[pip3] torch==2.6.0+cu126
[pip3] torchaudio==2.6.0+cu126
[pip3] torchvision==0.21.0+cu126
[pip3] triton==3.2.0
cc @mikaylagawarecki @ptrblck @msaroufim @eqy @jerryzh168 @tinglvv @nWEIdia
Contributor guide
Assessment
This issue has not been assessed yet.