Repository: facebookresearch/fairscale

Issue:
[ShardedDDP] Reduce the gradient bubble
## 🚀 Feature
Right now the backward pass is allowed to proceed all the way before the gradients are cleaned up, which means that all the gradients materialize. The peak memory use is bigger than what it should be if we clean them on the fly as the BW pass proceeds

## Motivation
Free more memory

## Pitch
See attached PR, look up the futures dequeue when firing up a BW hook, and cleanup whatever is ready (non blocking, important for perf)

## Alternatives
Not doing that, works but memory waste




Discussion:

Resolving patch:
diff --git a/fairscale/nn/data_parallel/sharded_ddp.py b/fairscale/nn/data_parallel/sharded_ddp.py
--- a/fairscale/nn/data_parallel/sharded_ddp.py
+++ b/fairscale/nn/data_parallel/sharded_ddp.py
@@ -11,7 +11,7 @@
 import contextlib
 from itertools import chain
 import logging
-from typing import Any, Callable, Generator, List, Tuple, Union
+from typing import Any, Callable, Dict, Generator, List, Tuple, Union
 
 import torch
 from torch import nn
@@ -92,6 +92,10 @@ def __init__(
         # we build an iterator which goes through all the parameters involved globally
         self._param_iterator = chain(*[optim.should_bucket_param.keys() for optim in self.sharded_optimizers])
         self._grad_to_be_reduced = [True for _ in self._param_iterator]
+        self._reduced_grads: Dict[OSS, int] = {}
+        self._reduced_grads_max = {o: len(o.param_to_rank.values()) for o in self.sharded_optimizers}
+        self._clear_counters()
+
         self._grad_accs: List[Callable] = []
         self._setup_backward_hooks()
 
@@ -110,7 +114,7 @@ def forward(self, *inputs: Any, **kwargs: Any) -> Any:
             self.sync_buffers(blocking=True)
 
         # Reset all the grad reduce and bucket state flags
-        self._grad_to_be_reduced = [True] * len(self._grad_to_be_reduced)
+        self._clear_counters()
 
         # Normal FW on the base model
         return self.module(*inputs, **kwargs)
@@ -122,18 +126,6 @@ def reduce(self) -> None:
         """
         logging.warning("This is not useful anymore, gradients have been reduced automatically with the backward pass")
 
-    def _sync_params_and_buffers(self) -> None:
-        """
-        Sync the complete model states in between the ranks
-        """
-        with torch.no_grad():
-            work_handles = [
-                dist.broadcast(t, src=self.reference_global_rank, group=self.process_group, async_op=True)
-                for t in self.module.state_dict().values()
-            ]
-
-            _ = list(map(lambda x: x.wait(), work_handles))
-
     def sync_buffers(self, blocking: bool = False) -> None:
         """
         Sync all the param buffers in between ranks (including for instance batch norm statistics).
@@ -155,6 +147,12 @@ def no_sync(self) -> Generator:
         yield
         self.should_accumulate_grads = old_should_accumulate_grads
 
+    def _clear_counters(self) -> None:
+        """ Reset all the grad reduce and call counters
+        """
+        self._grad_to_be_reduced = [True for _ in self._grad_to_be_reduced]
+        self._reduced_grads = {o: 0 for o in self.sharded_optimizers}
+
     def _find_rank(self, param: Parameter) -> Tuple[OSS, int]:
         """ Look up where this parameter belongs to """
         for optim in self.sharded_optimizers:
@@ -164,9 +162,7 @@ def _find_rank(self, param: Parameter) -> Tuple[OSS, int]:
         assert False, "This parameter is not present in an optimizer, this should not happen"
         return (None, -1)
 
-    def _get_reduce_fn(
-        self, index: int, param: torch.Tensor, should_bucket: bool, dst_rank: int, optimizer: OSS
-    ) -> Callable:
+    def _get_reduce_fn(self, index: int, param: torch.Tensor, dst_rank: int, optimizer: OSS) -> Callable:
         """
         Two possible backward hooks for a given parameter: either directly reduce to the appropriate rank,
         or contribute to a bucket and reduce when the bucket is full.
@@ -174,7 +170,7 @@ def _get_reduce_fn(
         Either way a delayed action is necessary and is passed as a callback.
         """
 
-        def reduce_direct(*_: Any) -> None:
+        def reduce(*_: Any) -> None:
             # Skip gradient reduction, do not alter status flags
             if not self.should_accumulate_grads and self._grad_to_be_reduced[index]:
                 assert param.grad is not None, "Reducing gradients during backward pass, cannot be None"
@@ -197,50 +193,18 @@ def cleanup() -> None:
                         callback=cleanup,
                     )
                 )
+                self._reduced_grads[optimizer] += 1
 
-                # If all the reduce operations have been called,
-                # make sure that all the asynchronous calls have concluded before moving on
-                # and execute the delayed actions (release gradients, unroll the buckets)
-                if len(optimizer.work_handles) == optimizer._max_work_handles:
-                    optimizer._consume_work_handles()
-
-        # Bucket, update status, and possibly unroll the results
-        def reduce_bucket(*_: Any) -> None:
-            # Skip gradient reduction, do not alter status flags
-            if not self.should_accumulate_grads and self._grad_to_be_reduced[index]:
-                assert param.grad is not None, "Reducing gradients during backward pass, cannot be None"
-
-                # Make sure that this is not fired twice
-                self._grad_to_be_reduced[index] = False
-
-                # Copy to the flat buffer, update the buffer state
-                bucket = optimizer.buckets[param.device][dst_rank]
-
-                assert bucket.append(param, use_gradient=True), "Bucket overflow: max %s - current %s - adding %s" % (
-                    bucket.max_size,
-                    bucket.current_offset,
-                    param.grad.numel(),
-                )
-
-                if bucket.full():
-                    bucket.buffer /= self.world_size
-
-                    optimizer.work_handles.append(
-                        Workhandle(
-                            handle=dist.reduce(
-                                tensor=bucket.buffer, dst=dst_rank, group=self.process_group, async_op=True,
-                            ),
-                            callback=bucket.unroll,
-                        )
-                    )
+                # Opportunistically try to empty the queue
+                optimizer._try_consume_work_handle()
 
                 # If all the reduce operations have been called,
                 # make sure that all the asynchronous calls have concluded before moving on
                 # and execute the delayed actions (release gradients, unroll the buckets)
-                if len(optimizer.work_handles) == optimizer._max_work_handles:
+                if self._reduced_grads[optimizer] == self._reduced_grads_max[optimizer]:
                     optimizer._consume_work_handles()
 
-        return reduce_bucket if should_bucket else reduce_direct
+        return reduce
 
     def _setup_backward_hooks(self) -> None:
         """
@@ -250,7 +214,7 @@ def _setup_backward_hooks(self) -> None:
 
         # Go through the parameters, attach the hook
         for sharded_optimizer in self.sharded_optimizers:
-            for param, should_bucket in sharded_optimizer.should_bucket_param.items():
+            for param, _ in sharded_optimizer.should_bucket_param.items():
                 if param.grad is not None and param.grad.requires_grad:
                     raise RuntimeError("ShardedDataParallel only works with gradients that don't require grad")
 
@@ -262,5 +226,17 @@ def _setup_backward_hooks(self) -> None:
                 dst_rank = sharded_optimizer.param_to_rank[param]
                 index = len(self._grad_accs)
 
-                grad_acc.register_hook(self._get_reduce_fn(index, param, should_bucket, dst_rank, sharded_optimizer))
+                grad_acc.register_hook(self._get_reduce_fn(index, param, dst_rank, sharded_optimizer))
                 self._grad_accs.append(grad_acc)  # keep this function in scope
+
+    def _sync_params_and_buffers(self) -> None:
+        """
+        Sync the complete model states in between the ranks
+        """
+        with torch.no_grad():
+            work_handles = [
+                dist.broadcast(t, src=self.reference_global_rank, group=self.process_group, async_op=True)
+                for t in self.module.state_dict().values()
+            ]
+
+            _ = list(map(lambda x: x.wait(), work_handles))
diff --git a/fairscale/optim/oss.py b/fairscale/optim/oss.py
--- a/fairscale/optim/oss.py
+++ b/fairscale/optim/oss.py
@@ -3,14 +3,14 @@
 # This source code is licensed under the BSD license found in the
 # LICENSE file in the root directory of this source tree.
 
-from collections import OrderedDict
+from collections import OrderedDict, deque
 import copy
 from enum import Enum, auto
 import itertools
 from itertools import chain
 import logging
 from math import inf
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Type, Union
+from typing import TYPE_CHECKING, Any, Callable, Deque, Dict, List, Optional, Tuple, Type, Union
 
 import torch
 import torch.distributed as dist
@@ -74,8 +74,6 @@ def __init__(
         broadcast_buffer_size: int = 2 ** 17,
         **default: Any,
     ):
-        # logging.warning("Disabling bucketing for now, error prone for some models")
-        broadcast_buffer_size = 0
 
         # Hold all the model params in the root .param_groups
         self.in_super_constructor = True
@@ -92,7 +90,6 @@ def __init__(
         self.world_size = dist.get_world_size(self.group)
         self.rank = dist.get_rank(self.group)
         self.global_rank = self.get_global_rank(self.group, self.rank)
-
         self.optim = optim(self.partition_parameters()[self.rank], **default)
 
         # - Sync local and global param_groups keys
@@ -107,6 +104,11 @@ def __init__(
         # Current default device is set by the parameters allocated to this rank
         self._device = list(self.per_device_params.keys())[0]
         self.buckets: Dict[torch.device, List[Bucket]] = {}
+
+        # if torch.cuda.is_available() and self.world_size <= torch.cuda.device_count():
+        #     broadcast_buffer_size = 0
+        #     logging.warning("Assuming single node job, bucketing is disabled")
+
         self.bucket_size = broadcast_buffer_size
         for device, per_device in self.per_device_params.items():
             # Allocate one buffer per rank and per device to group the small parameters
@@ -115,7 +117,7 @@ def __init__(
                 for _ in range(len(per_device))
             ]
         self.should_bucket_param: Dict[torch.Tensor, bool] = {}
-        self.work_handles: List[Workhandle] = []
+        self.work_handles: Deque[Workhandle] = deque()
         self._max_work_handles = -1
         self._setup_bucket_strategy()
 
@@ -554,12 +556,19 @@ def _consume_work_handles(self) -> None:
         We start from the first/older ones, since they are the most likely to be ready and non-blocking
         """
 
-        for work_handle in self.work_handles:
+        while len(self.work_handles) > 0:
+            work_handle = self.work_handles.popleft()
             work_handle.handle.wait()
             if work_handle.callback is not None:
                 work_handle.callback()
 
-        self.work_handles.clear()
+    def _try_consume_work_handle(self) -> None:
+        """ Try to consume the oldest future. This is non blocking, if not ready we'll pass
+        """
+        while len(self.work_handles) > 0 and self.work_handles[0].handle.is_completed():
+            work_handle = self.work_handles.popleft()
+            if work_handle.callback is not None:
+                work_handle.callback()
 
     def _handle_trailing_buckets(self, flush_type: BucketFlush) -> None:
         """
@@ -594,6 +603,10 @@ def _setup_bucket_strategy(self) -> None:
         network requests have been issued.
         """
 
+        # Determine the max work handles in flight:
+        # - count all the buckets on the fly
+        self._max_work_handles = 0
+
         for device, per_rank_params in self.per_device_params.items():
             for dst_rank, params in enumerate(per_rank_params):
                 offset = 0
@@ -604,6 +617,11 @@ def _setup_bucket_strategy(self) -> None:
                     # - enough room in the bucket
                     if (offset + param.numel()) < self.buckets[device][dst_rank].max_size:
                         self.should_bucket_param[param] = True
+
+                        if offset == 0:
+                            # count this bucket, only once
+                            self._max_work_handles += 1
+
                         offset += param.numel()
                     else:
                         self.should_bucket_param[param] = False
@@ -615,8 +633,4 @@ def _setup_bucket_strategy(self) -> None:
 
         # Determine the max work handles in flight:
         # - all the direct reduce/broadcast
-        self._max_work_handles = sum(not value for value in self.should_bucket_param.values())
-
-        # - if we're bucketing, this means more work handles: one per rank and per device
-        if self.bucket_size > 0:
-            self._max_work_handles += len(self.per_device_params.keys()) * self.world_size
+        self._max_work_handles += sum(not value for value in self.should_bucket_param.values())


Test patch: