enable gradient accumulation in SDD (#3462)

Summary:

Context: gradient accumulation is still not available in TorchRec, especially for the SDD pipeline which is being used by many recommendation models.

This diff: implements a UI enabling gradient accumulation for SDD.

Differential Revision: D84915986

Author: Zheng-Yong (Arsa) Ang

torchrec/distributed/train_pipeline/train_pipelines.py

@@ -420,6 +420,10 @@ class TrainPipelineSparseDist(TrainPipeline[In, Out]):
             (applicable to 2D sharding only)
             if set and DMP collection is enabled for 2D sharding,
             sync DMPs every N batches (default to 1, i.e. every batch, None to disable)
+        gradient_accumulation_steps (int): number of steps to accumulate gradients before
+            performing backward pass and optimizer update. Default is 1 (no accumulation).
+        gradient_accumulation_start_step (Optional[int]): the step to start gradient accumulation.
+            if not set, gradient accumulation will start from the second step.
     """
 
     # The PipelinedForward class that is used in _rewrite_model
@@ -440,6 +444,8 @@ def __init__(
         ] = None,
         dmp_collection_sync_interval_batches: Optional[int] = 1,
         enqueue_batch_after_forward: bool = False,
+        gradient_accumulation_steps: int = 1,
+        gradient_accumulation_start_step: Optional[int] = None,
     ) -> None:
         self._model = model
         self._optimizer = optimizer
@@ -505,6 +511,14 @@ def __init__(
             dmp_collection_sync_interval_batches
         )
 
+        # NOTE: the first step cannot be no_sync when DDP.static_graph = True, due to an unfortunate restriction in torch.distributed.
+        #       thus, unless explicitly set, we set gradient_accumulation_start_step to 1 if gradient_accumulation_steps > 1.
+        self._pre_accumulation_steps: int = gradient_accumulation_start_step or (
+            1 if gradient_accumulation_steps > 1 else 0
+        )
+        self._accumulation_steps: int = gradient_accumulation_steps
+        self._accumulation_step_count: int = 0
+
         if self._dmp_collection_sync_interval_batches is not None:
             logger.info(
                 f"{self.__class__.__name__}: [Sparse 2D] DMP collection will sync every "
@@ -682,7 +696,10 @@ def progress(self, dataloader_iter: Iterator[In]) -> Out:
         # TODO: Remove once Bulk Eval migrated (needed for bwd compat, this class only)
         self._set_module_context(self.contexts[0])
 
-        if self._model.training:
+        # only zero grad at the start of each accumulation
+        if self._model.training and (
+            self._pre_accumulation_steps > 0 or self._accumulation_step_count == 0
+        ):
             with record_function("## zero_grad ##"):
                 self._optimizer.zero_grad()
 
@@ -698,35 +715,50 @@ def progress(self, dataloader_iter: Iterator[In]) -> Out:
             # batch i+2: load data and copy to gpu, the dataload iter will first exhaust here
             self.enqueue_batch(dataloader_iter)
 
-        # forward
-        with record_function(f"## forward {self.contexts[0].index} ##"):
-            self._state = PipelineState.CALL_FWD
-            losses, output = self._model_fwd(self.batches[0])
+        no_sync = self._pre_accumulation_steps <= 0 and (
+            self._model.training
+            and self._accumulation_step_count + 1 < self._accumulation_steps
+        )
+        with (
+            self._model._dmp_wrapped_module.no_sync()  # pyre-ignore[16]
+            if no_sync
+            else contextlib.nullcontext()
+        ):
+            # forward
+            with record_function(f"## forward {self.contexts[0].index} ##"):
+                self._state = PipelineState.CALL_FWD
+                losses, output = self._model_fwd(self.batches[0])
 
-        if self._enqueue_batch_after_forward:
-            # batch i+2: load data and copy to gpu, the dataload iter will first exhaust here.
-            # Start this step after the forward of batch i, so that the H2D copy doesn't compete
-            # for pcie bandwidth with embedding lookup from UVM/UVM_CACHING.
-            self.enqueue_batch(dataloader_iter)
+            if self._enqueue_batch_after_forward:
+                # batch i+2: load data and copy to gpu, the dataload iter will first exhaust here.
+                # Start this step after the forward of batch i, so that the H2D copy doesn't compete
+                # for pcie bandwidth with embedding lookup from UVM/UVM_CACHING.
+                self.enqueue_batch(dataloader_iter)
 
-        if len(self.batches) >= 2:
-            # invoke data (values, lengths, etc.) all_to_all comms (second part of input_dist)
-            self.wait_sparse_data_dist(self.contexts[1])
+            if len(self.batches) >= 2:
+                # invoke data (values, lengths, etc.) all_to_all comms (second part of input_dist)
+                self.wait_sparse_data_dist(self.contexts[1])
 
-        if self._model.training:
             # backward
-            self._state = PipelineState.CALL_BWD
-            self._backward(losses)
-
-            self.sync_embeddings(
-                self._model,
-                self._dmp_collection_sync_interval_batches,
-                self.contexts[0],
-            )
+            if self._model.training:
+                self._state = PipelineState.CALL_BWD
+                self._backward(losses)
 
-            # update
-            with record_function(f"## optimizer {self.contexts[0].index} ##"):
-                self._optimizer.step()
+                if no_sync:
+                    self._accumulation_step_count += 1
+                else:
+                    self.sync_embeddings(
+                        self._model,
+                        self._dmp_collection_sync_interval_batches,
+                        self.contexts[0],
+                    )
+                    # update
+                    with record_function(f"## optimizer {self.contexts[0].index} ##"):
+                        self._optimizer.step()
+                    if self._pre_accumulation_steps > 0:
+                        self._pre_accumulation_steps -= 1
+                    else:
+                        self._accumulation_step_count = 0
 
         self.dequeue_batch()
         return output