Repository: pytorch/pytorch

Issue:
Fused AdamW causes NaN loss
### 🐛 Describe the bug

There already has been an extended discussion of this issue over on the nanoGPT repository:
- karpathy/nanoGPT#167

I have been encouraged to submit a separate bug report here, since the issue seems to lie with pytorch. I have prepared a preconfigured [fork of nanoGPT](https://github.com/oddlama/nanoGPT_nan) where I provided my configuration and part of my dataset which causes the issue to appear. It now immediately produces NaNs after the first training step, but only if fused adam is in use.

The settings are basically the nanoGPT shakespeare configuration, but using a blocksize of 343 and vocab size of 2006. The data is quite sparse (only 1,3% is not 0, which significantly accelerates getting to the issue). I've included 1000 batches of my own real data so you will have the same conditions as I do. I'm training on a single 2080Ti, please look at a diff of the latest commit in this repository to see what I changed exactly - it's not much.

```bash
> git clone https://github.com/oddlama/nanoGPT_nan
> cd nanoGPT_nan
```
```bash
> python train.py config/nan.py --allow_fused=True
[...]
number of parameters: 11.39M
using fused AdamW: True
compiling the model... (takes a ~minute)
step 0: train loss 6.6010, val loss 6.6009
[2023-03-01 15:19:43,485] torch._inductor.utils: [WARNING] using triton random, expect difference from eager
/projects/venv/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py:1251: UserWarning: Your compiler for AOTAutograd is returning a a function that doesn't take boxed arguments. Please wrap it with functorch.compile.make_boxed_func or handle the boxed arguments yourself. See https://github.com/pytorch/pytorch/pull/83137#issuecomment-1211320670 for rationale.
  warnings.warn(
iter 0: loss 6.6793, time 22316.39ms, mfu -100.00%
iter 1: loss nan, time 3249.91ms, mfu -100.00%
```
```bash
> python train.py config/nan.py --allow_fused=False
[...]
number of parameters: 11.39M
using fused AdamW: False
compiling the model... (takes a ~minute)
step 0: train loss 6.6010, val loss 6.6009
[2023-03-01 15:20:26,110] torch._inductor.utils: [WARNING] using triton random, expect difference from eager
/projects/nanoGPT_nan/venv/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py:1251: UserWarning: Your compiler for AOTAutograd is returning a a function that doesn't take boxed arguments. Please wrap it with functorch.compile.make_boxed_func or handle the boxed arguments yourself. See https://github.com/pytorch/pytorch/pull/83137#issuecomment-1211320670 for rationale.
  warnings.warn(
iter 0: loss 6.6793, time 22286.14ms, mfu -100.00%
iter 1: loss 6.6811, time 3534.44ms, mfu -100.00%
iter 2: loss 6.6770, time 3365.06ms, mfu -100.00%
iter 3: loss 4.7208, time 3366.52ms, mfu -100.00%
iter 4: loss 3.0927, time 3365.22ms, mfu -100.00%
iter 5: loss 2.0694, time 3362.80ms, mfu 6.51%
iter 6: loss 1.4752, time 3360.81ms, mfu 6.51%
# continues to work fine
```

Hope this helps.


P.S.:

Can anyone explain what the warning means? Doesn't seem to be related to the issue, but is it a user error that should be fixed?
```
[2023-03-01 15:19:43,485] torch._inductor.utils: [WARNING] using triton random, expect difference from eager
/projects/venv/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py:1251: UserWarning: Your compiler for AOTAutograd is returning a a function that doesn't take boxed arguments. Please wrap it with functorch.compile.make_boxed_func or handle the boxed arguments yourself. See https://github.com/pytorch/pytorch/pull/83137#issuecomment-1211320670 for rationale.
  warnings.warn(
```


### Versions

```
Collecting environment information...
PyTorch version: 2.0.0.dev20230220+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Arch Linux (x86_64)
GCC version: (GCC) 12.2.1 20230201
Clang version: 15.0.7
CMake version: version 3.25.0
Libc version: glibc-2.37

Python version: 3.10.9 (main, Dec 19 2022, 17:35:49) [GCC 12.2.0] (64-bit runtime)
Python platform: Linux-5.15.79.1-microsoft-standard-WSL2-x86_64-with-glibc2.37
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 2080 Ti
Nvidia driver version: 516.94
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:                   48 bits physical, 48 bits virtual
Byte Order:                      Little Endian
CPU(s):                          24
On-line CPU(s) list:             0-23
Vendor ID:                       AuthenticAMD
Model name:                      AMD Ryzen 9 5900X 12-Core Processor
CPU family:                      25
Model:                           33
Thread(s) per core:              2
Core(s) per socket:              12
Socket(s):                       1
Stepping:                        0
BogoMIPS:                        7400.01
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl tsc_reliable nonstop_tsc cpuid extd_apicid pni pclmulqdq ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw topoext ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 erms rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves clzero xsaveerptr arat umip vaes vpclmulqdq rdpid fsrm
Hypervisor vendor:               Microsoft
Virtualization type:             full
L1d cache:                       384 KiB (12 instances)
L1i cache:                       384 KiB (12 instances)
L2 cache:                        6 MiB (12 instances)
L3 cache:                        32 MiB (1 instance)
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Not affected
Vulnerability Retbleed:          Not affected
Vulnerability Spec store bypass: Vulnerable
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP conditional, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.1
[pip3] pytorch-triton==2.0.0+d54c04abe2
[pip3] torch==2.0.0.dev20230220+cu118
[pip3] torchaudio==2.0.0.dev20230223+cu118
[pip3] torchinfo==1.7.2
[pip3] torchsummary==1.5.1
[pip3] torchvision==0.15.0.dev20230223+cu118
[conda] Could not collect
```

cc @ezyang @gchanan @zou3519 @vincentqb @jbschlosser @albanD @janeyx99


Discussion:
cc @crcrpar This is pretty concerning--do you have any thoughts? Also, we most likely want to add AMP x fused adam(w) unit tests as it looks to be a common source of wrongness.
I'll look into it for sure.
though just so you know, there's a case of fused ones with grad scalers: https://github.com/pytorch/pytorch/blob/e5a959a2d49e627905c2ad58314a85c9e2f2c1df/test/test_cuda.py#L2462
> Can anyone explain what the warning means? Doesn't seem to be related to the issue, but is it a user error that should be fixed?

```
[2023-03-01 15:19:43,485] torch._inductor.utils: [WARNING] using triton random, expect difference from eager
/projects/venv/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py:1251: UserWarning: Your compiler for AOTAutograd is returning a a function that doesn't take boxed arguments. Please wrap it with functorch.compile.make_boxed_func or handle the boxed arguments yourself. See https://github.com/pytorch/pytorch/pull/83137#issuecomment-1211320670 for rationale.
  warnings.warn(
```

@oddlama nope - the first warnings is just letting you know that the results of operators involving randomness won't give identical results across eager mode and `torch.compile`. The second warning is a bug that you can ignore - it was recently fixed at https://github.com/pytorch/pytorch/pull/95521

Resolving patch:
diff --git a/torch/optim/adam.py b/torch/optim/adam.py
--- a/torch/optim/adam.py
+++ b/torch/optim/adam.py
@@ -529,7 +529,6 @@ def _fused_adam(
     capturable: bool,  # Needed for consistency.
     differentiable: bool,
 ) -> None:
-    grouped_tensors = _group_tensors_by_device_and_dtype([params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps])
     grad_scale_dict = {grad_scale.device: grad_scale} if grad_scale is not None else None
     found_inf_dict = {found_inf.device: found_inf} if found_inf is not None else None
     grouped_tensors = _group_tensors_by_device_and_dtype([params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps])
@@ -542,16 +541,15 @@ def _fused_adam(
             device_max_exp_avg_sqs,
             device_state_steps,
         ) = grouped_tensors[(device, dtype)]
-        if grad_scale is not None and found_inf is not None:
+        device_grad_scale, device_found_inf = None, None
+        if grad_scale is not None:
             if device not in grad_scale_dict:
                 grad_scale_dict[device] = grad_scale.to(device, non_blocking=True)
+            device_grad_scale = grad_scale_dict[device]
+        if found_inf is not None:
             if found_inf not in found_inf_dict:
                 found_inf_dict[device] = found_inf.to(device, non_blocking=True)
-            device_grad_scale = grad_scale_dict[device]
             device_found_inf = found_inf_dict[device]
-        else:
-            device_grad_scale = None
-            device_found_inf = None
         torch._foreach_add_(device_state_steps, 1)
         torch._fused_adam_(
             device_params,
diff --git a/torch/optim/adamw.py b/torch/optim/adamw.py
--- a/torch/optim/adamw.py
+++ b/torch/optim/adamw.py
@@ -602,16 +602,15 @@ def _fused_adamw(
             device_max_exp_avg_sqs,
             device_state_steps,
         ) = grouped_tensors[(device, dtype)]
-        if grad_scale is not None and found_inf is not None:
+        device_grad_scale, device_found_inf = None, None
+        if grad_scale is not None:
             if device not in grad_scale_dict:
                 grad_scale_dict[device] = grad_scale.to(device, non_blocking=True)
+            device_grad_scale = grad_scale_dict[device]
+        if found_inf is not None:
             if found_inf not in found_inf_dict:
                 found_inf_dict[device] = found_inf.to(device, non_blocking=True)
-            device_grad_scale = grad_scale_dict[device]
             device_found_inf = found_inf_dict[device]
-        else:
-            device_grad_scale = None
-            device_found_inf = None
         torch._foreach_add_(device_state_steps, 1)
         torch._fused_adamw_(
             device_params,


Test patch:
diff --git a/test/test_optim.py b/test/test_optim.py
--- a/test/test_optim.py
+++ b/test/test_optim.py
@@ -1619,11 +1619,13 @@ def test_functional_fused_optimizer_with_foundinf(self):
         from torch.optim import adam, adamw
 
         num_tensors = 5
-        for functional_optim, amsgrad in itertools.product((adam.adam, adamw.adamw), (False, True)):
-            params, grads, exp_avgs, exp_avg_sqs = [[torch.ones((1,), device="cuda") for _ in range(num_tensors)] for _ in range(4)]
+        for functional_optim, amsgrad, no_grad_scale in itertools.product((adam.adam, adamw.adamw), (False, True), (False, True)):
+            params, grads, exp_avgs, exp_avg_sqs = [
+                [torch.ones((1,), device="cuda") for _ in range(num_tensors)] for _ in range(4)]
+            prev_params = [t.clone().detach() for t in params]
             max_exp_avg_sqs = [torch.ones((1,), device="cuda") for _ in range(num_tensors)] if amsgrad else []
             state_steps = [torch.ones((1,), dtype=torch.float32, device="cuda") for _ in range(num_tensors)]
-            grad_scale = torch.ones((1,), dtype=torch.float32, device="cuda")
+            grad_scale = None if no_grad_scale else torch.ones((1,), dtype=torch.float32, device="cuda")
             found_inf = torch.ones((1,), dtype=torch.float32, device="cuda")
 
             functional_optim(
@@ -1654,6 +1656,7 @@ def test_functional_fused_optimizer_with_foundinf(self):
                     for _ in range(num_tensors)
                 ],
             )
+            self.assertEqual(params, prev_params)
 
     def test_empty_grad(self):
         optimizers = [