Repository: pytorch/pytorch

Issue:
Fused AdamW has worse loss than Apex and unfused AdamW for fp16/AMP
### 🐛 Describe the bug

@stas00 recently benchmarked on an 80GB A100 and discovered fused AdamW is worse than both the Apex equivalent and the default implementation. Moreover, fused AdamW is not more performant in this case.
 
| Variation                 |     Train samples per second |   Diff % |   Train loss |
|:--------------------------|----------:|-------:|--------:|
| --optim adamw_torch_fused |    389.41 |      0 |    2.66 |
| --optim adamw_torch       |    389.37 |      0 |    2.55 |
| --optim adamw_apex_fused  |    399.27 |      3 |    2.53 |

Repro:
See @stas00's comment below!
The batch size (--per_device_train_batch_size) can be adjusted to be smaller.

cc @ezyang @gchanan @zou3519 @vincentqb @jbschlosser @albanD @crcrpar

### Versions

Master/latest nightly


Discussion:

Resolving patch:
diff --git a/torch/cuda/amp/grad_scaler.py b/torch/cuda/amp/grad_scaler.py
--- a/torch/cuda/amp/grad_scaler.py
+++ b/torch/cuda/amp/grad_scaler.py
@@ -286,7 +286,9 @@ def unscale_(self, optimizer):
 
     def _maybe_opt_step(self, optimizer, optimizer_state, *args, **kwargs):
         retval = None
-        if not sum(v.item() for v in optimizer_state["found_inf_per_device"].values()):
+        # NOTE(crcrpar): Gradients could be inf/nan after `GradScaler.unscale_(optimizer)`
+        # especially when gradient clipping is applied.
+        if not sum(v.item() for v in self._check_inf_per_device(optimizer).values()):
             retval = optimizer.step(*args, **kwargs)
         return retval
 


Test patch:
diff --git a/test/test_cuda.py b/test/test_cuda.py
--- a/test/test_cuda.py
+++ b/test/test_cuda.py
@@ -2512,6 +2512,43 @@ def _grad_scaling_autocast_fused_optimizers(self, optimizer_ctor, optimizer_kwar
                         actual = actual.squeeze()
                     self.assertEqual(state_control[k], actual)
 
+    def test_grads_invalidated_between_unscale_and_step(self):
+        for optimizer_ctor, optimizer_kwargs in product(
+            (torch.optim.Adam, torch.optim.AdamW),
+            (
+                {"foreach": False, "fused": False},
+                {"foreach": True, "fused": False},
+                {"foreach": False, "fused": True},
+            ),
+        ):
+            with self.subTest(optimizer=optimizer_ctor, optimizer_kwargs=optimizer_kwargs):
+                self._test_grads_invalidated_between_unscale_and_step(optimizer_ctor, optimizer_kwargs)
+
+    def _test_grads_invalidated_between_unscale_and_step(self, optimizer_ctor, optimizer_kwargs):
+        model, _, optimizer, _, data, loss_fn, _ = self._create_scaling_case(
+            optimizer_ctor=optimizer_ctor, optimizer_kwargs=optimizer_kwargs,
+        )
+        scaler = torch.cuda.amp.GradScaler(init_scale=128.0)
+
+        orig_params = [p.clone().detach() for p in model.parameters()]
+
+        for i, (input, target) in enumerate(data):
+            optimizer.zero_grad()
+            with torch.autocast('cuda', enabled=True):
+                output = model(input)
+                loss = loss_fn(output, target)
+            scaler.scale(loss).backward()
+            scaler.unscale_(optimizer)
+
+            # deliberately break grads
+            for j, param in enumerate(model.parameters()):
+                param.grad.copy_(torch.inf if j % 2 else torch.nan)
+
+            scaler.step(optimizer)
+            scaler.update()
+
+        self.assertEqual(orig_params, list(model.parameters()))
+
     def test_grad_scaling_clipping(self):
         def run(data, model, optimizer, scaler, loss_fn, skip_iter, try_scaling_api):
             max_norm = 0.2  # A reasonable value that actually has an effect, based on printouts of grads