Add gradient_accumulation_no_sync flag to AutoUnit (#1057)

Summary:
Pull Request resolved: #1057

 ### Context

TorchTNT's AutoUnit.train_step calls set_requires_gradient_sync(False) on FSDP2 modules during gradient accumulation micro-batches. This tells FSDP2 to skip the gradient reduce-scatter in backward, keeping full unsharded gradients in memory on each rank. For large models this is catastrophic:

 - 26B model on 16 GPUs: sharded gradients = 3.25 GB/rank, unsharded = ~52 GB/rank. Observed memory jump from 38% → 86% GPU utilization — a 48% increase from gradient accumulation alone.
 - Small models (2B-4B): only 5-8% increase, because the absolute gradient size is small relative to 80 GB GPU memory.

The FSDP2 API already supports the memory-efficient alternative: simply don't disable gradient sync. Gradients reduce-scatter every micro-batch and accumulate in .grad in sharded form. This is exactly what torchtitan does. The only tradeoff is N× more reduce-scatter communication (N = gradient_accumulation_steps).

Goal: Add a single flag to AutoUnit that controls whether no_sync is used during gradient accumulation, enabling users to trade communication for memory.

 ### Approach

 Add gradient_accumulation_no_sync: bool = True to AutoUnit.__init__. When False, skip both:
 - FSDP2: set_requires_gradient_sync(False) call
 - DDP/FSDP1: module.no_sync() context manager

 Gradients then reduce-scatter (FSDP2) or all-reduce (DDP) on every micro-batch and accumulate in sharded .grad. The math is identical —
 reduce is associative over micro-batches.

 No changes needed in HuggingfaceSFTUnit or MediaAutoUnit — the flag flows through **kwargs.

Reviewed By: galrotem

Differential Revision: D99706122

fbshipit-source-id: 321a130aaaacfa98af9551f9d4579273823e7c23

Author: diego-urgell

tests/framework/test_auto_unit.py

@@ -781,6 +781,27 @@ def test_gradient_accumulation_fsdp2(self, _) -> None:
 
             auto_unit.train_progress.increment_step()
 
+    @patch("torchtnt.framework.auto_unit._is_fsdp2_module", return_value=True)
+    def test_gradient_accumulation_fsdp2_sync_true(self, _) -> None:
+        """When gradient_accumulation_sync=True, set_requires_gradient_sync
+        should never be called — gradients sync on every micro-batch."""
+        auto_unit = DummyAutoUnit(
+            module=torch.nn.Linear(1, 1),
+            gradient_accumulation_steps=3,
+            gradient_accumulation_sync=True,
+        )
+
+        fsdp_module_mock = MagicMock()
+        auto_unit.module.set_requires_gradient_sync = fsdp_module_mock
+        auto_unit._is_last_batch = False
+
+        state = get_dummy_train_state()
+
+        for _step in range(4):
+            auto_unit.train_step(state, (torch.rand(1, 1), torch.rand(1, 1)))
+            fsdp_module_mock.assert_not_called()
+            auto_unit.train_progress.increment_step()
+
     @patch("torchtnt.framework.auto_unit.prepare_module")
     def test_global_mesh(self, mock_prepare_module: Mock) -> None:
         """

torchtnt/framework/auto_unit.py

@@ -474,6 +474,11 @@ class AutoUnit(
         step_lr_interval: whether to step lr_scheduler every step or every epoch. Defaults to every epoch.
         precision: the precision to use in training/evaluation (using automatic mixed precision), as either a string or a torch.dtype. Acceptable strings are ``'fp32'``, ``'fp16'``, and ``'bf16'``.
         gradient_accumulation_steps: how many batches to accumulate gradients over.
+        gradient_accumulation_sync: if True, gradients are synchronized (reduce-scatter for FSDP2,
+            all-reduce for DDP) on every micro-batch during gradient accumulation and accumulated in sharded form.
+            This uses more communication but significantly less memory — recommended for large models.
+            If False (default), synchronization is skipped during accumulation micro-batches using ``no_sync`` / ``set_requires_gradient_sync(False)``,
+            which reduces communication but keeps full unsharded gradients in memory.
         detect_anomaly: whether to enable anomaly detection for the autograd engine https://pytorch.org/docs/stable/autograd.html#anomaly-detection
         clip_grad_norm: max norm of the gradients for clipping https://pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html
         clip_grad_value: max value of the gradients for clipping https://pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_value_.html
@@ -515,6 +520,7 @@ def __init__(
         step_lr_interval: Literal["step", "epoch"] = "epoch",
         precision: Optional[Union[str, torch.dtype]] = None,
         gradient_accumulation_steps: int = 1,
+        gradient_accumulation_sync: bool = False,
         detect_anomaly: Optional[bool] = None,
         clip_grad_norm: Optional[float] = None,
         clip_grad_value: Optional[float] = None,
@@ -585,6 +591,7 @@ def __init__(
         self.step_lr_interval = step_lr_interval
 
         self.gradient_accumulation_steps = gradient_accumulation_steps
+        self.gradient_accumulation_sync = gradient_accumulation_sync
 
         self.clip_grad_norm = clip_grad_norm
         self.clip_grad_value = clip_grad_value
@@ -697,11 +704,12 @@ def train_step(self, state: State, data: TData) -> Tuple[torch.Tensor, Any]:
         maybe_no_sync = (
             module.no_sync()
             if not should_update_weights
+            and not self.gradient_accumulation_sync
             and (isinstance(module, DDP) or isinstance(module, FSDP))
             else contextlib.nullcontext()
         )
         # fsdp2 has separate way of disabling gradient sync
-        if _is_fsdp2_module(module):
+        if _is_fsdp2_module(module) and not self.gradient_accumulation_sync:
             if not should_update_weights:
                 cast(FSDPModule, module).set_requires_gradient_sync(False)
             elif should_update_weights and self.gradient_accumulation_steps > 1: