Repository: facebookresearch/fairscale

Issue:
[ShardedDDP] Add a manual/ondemand reduction option
## 🚀 Feature
Add a method to enable ondemand grad reduction

## Motivation
Useful for checkpointing

## Pitch
Add a manual option to send to reduce the gradients to the appropriate ranks, nothing super complicated

## Alternatives
Not doing it, meaning that checkpointing support is more complicated

## Additional context
cc @prigoyal @min-xu-ai 


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
@@ -250,7 +250,9 @@ def refresh_trainable(self) -> None:
         """ If the module trainability has changed, update all the assumptions """
 
         # Make sure that this is not done while gradients are waiting to be reduced (if no_sync context for instance)
-        assert not functools.reduce(lambda x, y: x or y, self._grad_to_be_reduced, False), "Grads waiting to be reduced"
+        assert not functools.reduce(
+            lambda x, y: x or y, self._grad_to_be_reduced, False
+        ), "Grads waiting to be reduced: {}".format(self._grad_to_be_reduced)
 
         self._trainable_params = list(filter(lambda x: x.requires_grad, self._all_params))
         self._trainable_params.sort(key=lambda x: x.numel())
@@ -273,11 +275,21 @@ def refresh_trainable(self) -> None:
         self._setup_backward_hooks()
 
     def reduce(self) -> None:
-        """.. deprecated:: 0.0.4
-
-        This does not need to be called, the gradient reduction is done automatically during the BW pass
         """
-        logging.warning("This is not useful anymore, gradients have been reduced automatically with the backward pass")
+        This does not *need* to be called, the gradient reduction is done automatically during the BW pass.
+        Use this method to reduce the gradients manually
+        """
+
+        # Check that this is not a mistake, if there's nothing to reduce
+        assert functools.reduce(
+            lambda x, y: x or y, self._grad_to_be_reduced, False
+        ), "No grads waiting to be reduced, maybe that this was called twice or there was no BW pass ?"
+
+        # Trigger all the current BW hooks
+        _ = map(lambda x: x(), self._grad_accs)
+
+        # Make sure that all the futures are consumed
+        self._consume_work_handles()
 
     @torch.no_grad()
     def sync_buffers(self, blocking: bool = False) -> None:


Test patch:
diff --git a/tests/nn/data_parallel/test_sharded_ddp.py b/tests/nn/data_parallel/test_sharded_ddp.py
--- a/tests/nn/data_parallel/test_sharded_ddp.py
+++ b/tests/nn/data_parallel/test_sharded_ddp.py
@@ -137,10 +137,10 @@ def run_ddp_parity(rank, world_size, backend, temp_file_name):
     INPUTS = 2
     BATCH_SIZE = 32
 
-    def check_parity(amp: bool, accumulate: bool, change_train_graph: bool):
+    def check_parity(amp: bool, accumulate: bool, change_train_graph: bool, manual_reduction: bool):
 
         # The API should be the exact same in between the sharded and non-sharded variants, generic closure
-        def closure(model, scaler, input_tensor, should_accumulate):
+        def closure(model, scaler, input_tensor, should_accumulate, _manual_reduction=False):
             accumulate_steps = 3 if should_accumulate else 1
 
             model.zero_grad()
@@ -158,7 +158,13 @@ def step():
                 for _ in range(accumulate_steps - 1):
                     step()
 
-            step()
+            if not _manual_reduction:
+                step()
+            else:
+                with model.no_sync():
+                    step()
+
+                model.reduce()
 
         # Any model works. Add one different buffer per rank
         model = Sequential(Linear(INPUTS, 3), Linear(3, 3), Linear(3, 3), Linear(3, 3), Linear(3, 3), Linear(3, 3))
@@ -192,7 +198,9 @@ def closure_ddp(input_tensor=input_tensor):
                 return closure(ddp_model, ddp_scaler, input_tensor, accumulate)
 
             def closure_sharded(input_tensor=input_tensor):
-                return closure(sharded_ddp_model, sharded_ddp_scaler, input_tensor, accumulate)
+                return closure(
+                    sharded_ddp_model, sharded_ddp_scaler, input_tensor, accumulate, _manual_reduction=manual_reduction
+                )
 
             # Step/scale both
             if ddp_scaler is not None:
@@ -226,11 +234,18 @@ def closure_sharded(input_tensor=input_tensor):
 
     for accumulate in [False, True]:
         for change_train_graph in [False, True]:
-            for amp in amp_tests:
-                print(
-                    f"Checking configuration: accumulate {accumulate} - change train graph {change_train_graph} - amp {amp}"
-                )
-                check_parity(amp=amp, accumulate=accumulate, change_train_graph=change_train_graph)
+            manual_reductions = [False, True] if not accumulate and not change_train_graph else [False]
+            for manual_reduction in manual_reductions:
+                for amp in amp_tests:
+                    print(
+                        f"Checking configuration: accumulate {accumulate} - change train graph {change_train_graph} - amp {amp} - manual reduction {manual_reduction}"
+                    )
+                    check_parity(
+                        amp=amp,
+                        accumulate=accumulate,
+                        change_train_graph=change_train_graph,
+                        manual_reduction=manual_reduction,
+                    )
 
     dist.destroy_process_group()