Repository: facebookresearch/fairscale

Issue:
[OSS] Handle torch AMP auto skip inf/nan steps
## 🚀 Feature
Using torch amp can lead to step() being skipped in the presence of inf or nans. An issue with that with OSS is that each rank will see a different step, meaning that some ranks could skip and others try to step (typically vanishing gradient on the way up and underflow with fp16). We need to handle that properly

## Motivation
Without that using torch amp with OSS can easily deadlock, see for instance https://github.com/facebookresearch/fairscale/pull/179

## Pitch
A somewhat elegant and not too far reaching solution to that could be to add a config flag to OSS, following which each rank will check for inf/nan on its grads, and signal all ranks that they should skip step() in that case. Users willing to apply OSS and torch amp would be advised to switch this flag on to prevent deadlocks

## Alternatives
Prevent mixed precision and OSS, but that makes OSS a lot less compelling

## Additional context
#179 and https://pytorch.org/docs/stable/notes/amp_examples.html?highlight=mixed%20precision#id2


Discussion:
hmm, I'm wondering now if subclassing Scaler would not be a better idea, so that we check on all grads before hypothetically skipping the step()
Hey @blefaudeux I think we're in sync on this, I'm currently looking at subclassing the scaler and ensuring the inf check is reduced correctly:

```python
import torch.distributed as dist
from torch.cuda.amp import GradScaler


class ShardedGradScaler(GradScaler):
    def step(self, optimizer, *args, **kwargs):
        optimizer_state = self._per_optimizer_states[id(optimizer)]

        for v in optimizer_state["found_inf_per_device"].values():
            dist.all_reduce(v)

        return super().step(optimizer, *args, **kwargs)
```
Agreed, so if OSS can expose a distributed gradient check (not too different from the clip_grad_norm) the ShardedGradScaler code would be fairly minimal and easy to grok
looks like there are some subtleties to keep in mind for speed (for instance prevent the upstream GradScaler from spending too long checking for infs again in the `super().step()`), but doable. I should have a working example along these lines in the next couple of hours
> Hey @blefaudeux I think we're in sync on this, I'm currently looking at subclassing the scaler and ensuring the inf check is reduced correctly:
> 
> ```python
> import torch.distributed as dist
> from torch.cuda.amp import GradScaler
> 
> 
> class ShardedGradScaler(GradScaler):
>     def step(self, optimizer, *args, **kwargs):
>         optimizer_state = self._per_optimizer_states[id(optimizer)]
> 
>         for v in optimizer_state["found_inf_per_device"].values():
>             dist.all_reduce(v)
> 
>         return super().step(optimizer, *args, **kwargs)
> ```

looks like this take is the cleanest, I just spend some time getting up to speed with GradScaler and it looks like the most idiomatic way indeed. 

Resolving patch:
diff --git a/benchmarks/oss.py b/benchmarks/oss.py
--- a/benchmarks/oss.py
+++ b/benchmarks/oss.py
@@ -14,6 +14,7 @@
 import numpy as np
 import torch
 import torch.autograd.profiler as profiler
+from torch.cuda.amp import GradScaler as TorchGradScaler
 import torch.distributed as dist
 import torch.multiprocessing as mp
 import torch.nn as nn
@@ -25,6 +26,7 @@
 
 from fairscale.nn.data_parallel import ShardedDataParallel as ShardedDDP
 from fairscale.optim import OSS
+from fairscale.optim.grad_scaler import ShardedGradScaler
 
 OPTIM = torch.optim.RMSprop
 TEMPDIR = tempfile.gettempdir()
@@ -92,6 +94,7 @@ def train(
     # Shard the optimizer
     optimizer: Optional[torch.optim.Optimizer] = None
     model = cast(nn.Module, model)
+    scaler = (TorchGradScaler() if args.optim_type == OptimType.vanilla else ShardedGradScaler()) if args.amp else None
 
     if optim_type == OptimType.oss_sharded_ddp:
         model = ShardedDDP(
@@ -102,7 +105,6 @@ def train(
             broadcast_buffers=True,
         )
         optimizer = model.sharded_optimizer
-
     else:
         if args.cpu:
             device_ids = None
@@ -136,7 +138,7 @@ def train(
         for batch in dataloader:
             batch__start = time.monotonic()
 
-            def closure(data=batch):
+            def closure(data=batch, grad_scaler=None):
                 model.zero_grad()
                 if args.debug and rank == 0 and next(model.parameters()).grad is not None:
                     logging.debug(
@@ -144,16 +146,18 @@ def closure(data=batch):
                             next(model.parameters()).norm().item(), next(model.parameters()).grad.norm().item()
                         )
                     )
-                if not args.cpu and args.amp:
+                if grad_scaler is not None:
                     # Automatically computes the FW pass in half precision
                     with torch.cuda.amp.autocast():
                         outputs = model(data["inputs"])
                         loss = loss_fn(outputs, data["label"])
+
+                        # Accumulates scaled gradients.
+                        grad_scaler.scale(loss).backward()
                 else:
                     outputs = model(data["inputs"])
                     loss = loss_fn(outputs, data["label"])
-
-                loss.backward()
+                    loss.backward()
 
                 if optim_type == OptimType.oss_sharded_ddp:
                     model.reduce()
@@ -170,16 +174,24 @@ def closure(data=batch):
                 logging.info("Profiling the run")
                 with profiler.profile(use_cuda=True, record_shapes=True, profile_memory=True) as prof:  # type: ignore
                     with profiler.record_function("batch"):
-                        final_loss = optimizer.step(closure)
-                        logging.info("profiling done")
+                        if scaler is not None:
+                            final_loss = closure(grad_scaler=scaler)  # AMP scaler.step does not support closures
+                            scaler.step(optimizer)
+                            scaler.update()
+                        else:
+                            final_loss = optimizer.step(closure)
 
-                if rank == 0:
-                    prof.export_chrome_trace(f"{optim_type}_trace.json")
+                        prof.export_chrome_trace(f"{optim_type}_trace_rank_{rank}.json")
 
                 need_profiling = False  # only profile once
 
             else:
-                final_loss = optimizer.step(closure)
+                if scaler is not None:
+                    final_loss = closure(grad_scaler=scaler)  # AMP scaler.step does not support closures
+                    scaler.step(optimizer)
+                    scaler.update()
+                else:
+                    final_loss = optimizer.step(closure)
 
             if args.debug and rank == 0:
                 logging.debug("buffer: {}".format(next(model.buffers()).norm().item()))
diff --git a/examples/tutorial_oss.py b/examples/tutorial_oss.py
--- a/examples/tutorial_oss.py
+++ b/examples/tutorial_oss.py
@@ -57,7 +57,8 @@ def train(rank: int, world_size: int, epochs: int, use_oss: bool):
     training_start = time.monotonic()
     # Any relevant training loop, nothing specific to OSS. For example:
     model.train()
-    for e in range(epochs):
+
+    for _ in range(epochs):
         for (data, target) in dataloader:
             data, target = data.to(rank), target.to(rank)
 
diff --git a/fairscale/optim/grad_scaler.py b/fairscale/optim/grad_scaler.py
--- a/fairscale/optim/grad_scaler.py
+++ b/fairscale/optim/grad_scaler.py
@@ -3,15 +3,44 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Dict
+from typing import Any, Dict, Optional
 
-from torch import Tensor, device
+import torch
 from torch.cuda.amp import GradScaler as TorchGradScaler
+import torch.distributed as dist
 from torch.optim import Optimizer
 
+from .oss import OSS
+
 
 class GradScaler(TorchGradScaler):
     def _unscale_grads_(
-        self, optimizer: Optimizer, inv_scale: Tensor, found_inf: Tensor, allow_fp16: bool
-    ) -> Dict[device, Tensor]:
+        self, optimizer: Optimizer, inv_scale: torch.Tensor, found_inf: torch.Tensor, allow_fp16: bool
+    ) -> Dict[torch.device, torch.Tensor]:
         return super()._unscale_grads_(optimizer, inv_scale, found_inf, True)
+
+
+class ShardedGradScaler(TorchGradScaler):
+    """
+    A shard-aware :class:`GradScaler<torch.cuda.amp.GradScaler>`, to be used in conjunction with
+    :class:`OSS` and :class:`ShardedOptimizer`.
+
+    Interface and usecases are not changed, more explanations can be found in the corresponding pytorch
+    documentation https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.GradScaler
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def step(self, optimizer: Optimizer, *args: Any, **kwargs: Any) -> Optional[float]:
+        assert isinstance(optimizer, OSS), "ShardedGradScaler is to be used in combination with a sharded optimizer"
+
+        # Re-use the GradSCaler machinery, but make sure that the status is sync'ed in between the ranks
+        optimizer_state = self._per_optimizer_states[id(optimizer)]
+        handles = [dist.all_reduce(v, async_op=True) for v in optimizer_state["found_inf_per_device"].values()]
+
+        # Make sure that the calls are done before moving out
+        _ = list(map(lambda x: x.wait(), handles))
+
+        # Call Torch's GradScaler in turn, states have been synchronized across ranks
+        return super().step(optimizer, *args, **kwargs)


Test patch: