Repository: facebookresearch/fairscale

Issue:
Race condition in gradient reduction
Following up on the previous issue, as I couldn't reopen: https://github.com/facebookresearch/fairscale/issues/777

There seems to be a more fundamental issue here unrelated to mixed precision, and the race condition I mentioned in that issue appears real.

When `method FC` is used, enabling `maintain_forward_counter` in `checkpoint_wrapper` is not optional, as I would expect.

The forward counters seem like a trade-off. With them, you can reduce network overhead at the cost of more VRAM, by retaining the full parameters until the last backward. But in my use case, I have multiple large modules called like Sequence[A, B, C, .., A, B, C...]. I want to reshard after both forward and backward, so using the counters is undesirable.

The unit test fails to capture this error, because the weight tensors are so small that the race condition I was concerned about actually happens. If you make the weights a little larger or force synchronous gradient reduction, the test will fail, without the forward counters. See the diffs below for how to trigger it.

These errors could be avoided with stronger consistency in FSDP's internal state. In my code, I partially avoid the problem by using `method CF` with some small hacks, but the race condition is still technically present. Perhaps a better solution would use `register_full_backward_hook` or similar to detect when we have completed the backward pass for the FSDP module, and reset the appropriate state. Also have a barrier in place to ensure any previous async reduction work is done before the next pre-backward call. This was discussed in the previous issue as well.

Diff for unit test
```
diff --git a/tests/nn/data_parallel/test_fsdp_multiple_forward_checkpoint.py b/tests/nn/data_parallel/test_fsdp_multiple_forward_checkpoint.py
index b4127bc..6776e79 100644
--- a/tests/nn/data_parallel/test_fsdp_multiple_forward_checkpoint.py
+++ b/tests/nn/data_parallel/test_fsdp_multiple_forward_checkpoint.py
@@ -56,10 +56,10 @@ class Model2(nn.Module):
 
     def __init__(self):
         super().__init__()
-        self.block1 = nn.Sequential(nn.Conv2d(3, 4, kernel_size=3), nn.BatchNorm2d(4), nn.ReLU(inplace=True))
-        self.block2 = nn.Sequential(nn.Conv2d(4, 4, kernel_size=3), nn.BatchNorm2d(4), nn.ReLU(inplace=False))
-        self.block3 = nn.Sequential(nn.Conv2d(4, 8, kernel_size=3), nn.BatchNorm2d(8), nn.ReLU(inplace=True))
-        self.head = nn.Sequential(nn.AdaptiveAvgPool2d(output_size=(1, 1)), nn.Flatten(), nn.Linear(8, 10))
+        self.block1 = nn.Sequential(nn.Conv2d(3, 400, kernel_size=3), nn.BatchNorm2d(400), nn.ReLU(inplace=True))
+        self.block2 = nn.Sequential(nn.Conv2d(400, 400, kernel_size=3), nn.BatchNorm2d(400), nn.ReLU(inplace=False))
+        self.block3 = nn.Sequential(nn.Conv2d(400, 800, kernel_size=3), nn.BatchNorm2d(800), nn.ReLU(inplace=True))
+        self.head = nn.Sequential(nn.AdaptiveAvgPool2d(output_size=(1, 1)), nn.Flatten(), nn.Linear(800, 10))
 
     def forward(self, x):
         if isinstance(x, torch.Tensor):
@@ -91,9 +91,9 @@ def _create_model(
             if with_model2:
                 model.block3 = auto_wrap_bn(model.block3, single_rank_pg=False, fsdp_config=fsdp_config)
         if with_checkpoint:
-            model.block2 = checkpoint_wrapper(model.block2, maintain_forward_counter=True)
+            model.block2 = checkpoint_wrapper(model.block2, maintain_forward_counter=False)
             if with_model2:
-                model.block3 = checkpoint_wrapper(model.block3, maintain_forward_counter=True)
+                model.block3 = checkpoint_wrapper(model.block3, maintain_forward_counter=False)
         with enable_wrap(
             wrapper_cls=FSDP,
             flatten_parameters=flatten,
```

Alternatively, force synchronous reduction, and the unit test will fail with small weights.
```
diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index f9a75c8..b77bc08 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -1360,6 +1360,8 @@ class FullyShardedDataParallel(nn.Module):
                 assert self._reducer is not None
                 grad_chunks = chunk_and_pad(param.grad.data, self.world_size)
                 self._reducer.reduce_scatter_async(grad_chunks, group=self.process_group, callback_fn=callback_fn)
+                self._reducer.flush()
+                self._streams["post_backward"].synchronize()
             else:
                 # Currently the only way for _is_sharded to be False is if
                 # world_size == 1. This could be relaxed in the future, in which
```

Example error:

`NCCL_P2P_DISABLE=1 python -m pytest -s tests/nn/data_parallel/test_fsdp_multiple_forward_checkpoint.py::test_multiple_forward_checkpoint[bn-model2-no_auto_wrap_bn-flatten-full]`
```
E       torch.multiprocessing.spawn.ProcessRaisedException: 
E       
E       -- Process 1 terminated with the following error:
E       Traceback (most recent call last):
E         File "/home/darryl/miniconda3/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 59, in _wrap
E           fn(i, *args)
E         File "/home/darryl/Projects/fairscale/tests/nn/data_parallel/test_fsdp_multiple_forward_checkpoint.py", line 200, in _distributed_worker
E           loss.backward()
E         File "/home/darryl/miniconda3/lib/python3.8/site-packages/torch/_tensor.py", line 255, in backward
E           torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
E         File "/home/darryl/miniconda3/lib/python3.8/site-packages/torch/autograd/__init__.py", line 147, in backward
E           Variable._execution_engine.run_backward(
E         File "/home/darryl/miniconda3/lib/python3.8/site-packages/torch/autograd/function.py", line 87, in apply
E           return self._forward_cls.backward(self, *args)  # type: ignore[attr-defined]
E         File "/home/darryl/Projects/fairscale/fairscale/nn/checkpoint/checkpoint_activations.py", line 323, in backward
E           torch.autograd.backward(outputs_with_grad, args_with_grad)
E         File "/home/darryl/miniconda3/lib/python3.8/site-packages/torch/autograd/__init__.py", line 147, in backward
E           Variable._execution_engine.run_backward(
E       RuntimeError: The size of tensor a (720600) must match the size of tensor b (1441200) at non-singleton dimension 0
```
Race condition in gradient reduction
Following up on the previous issue, as I couldn't reopen: https://github.com/facebookresearch/fairscale/issues/777

There seems to be a more fundamental issue here unrelated to mixed precision, and the race condition I mentioned in that issue appears real.

When `method FC` is used, enabling `maintain_forward_counter` in `checkpoint_wrapper` is not optional, as I would expect.

The forward counters seem like a trade-off. With them, you can reduce network overhead at the cost of more VRAM, by retaining the full parameters until the last backward. But in my use case, I have multiple large modules called like Sequence[A, B, C, .., A, B, C...]. I want to reshard after both forward and backward, so using the counters is undesirable.

The unit test fails to capture this error, because the weight tensors are so small that the race condition I was concerned about actually happens. If you make the weights a little larger or force synchronous gradient reduction, the test will fail, without the forward counters. See the diffs below for how to trigger it.

These errors could be avoided with stronger consistency in FSDP's internal state. In my code, I partially avoid the problem by using `method CF` with some small hacks, but the race condition is still technically present. Perhaps a better solution would use `register_full_backward_hook` or similar to detect when we have completed the backward pass for the FSDP module, and reset the appropriate state. Also have a barrier in place to ensure any previous async reduction work is done before the next pre-backward call. This was discussed in the previous issue as well.

Diff for unit test
```
diff --git a/tests/nn/data_parallel/test_fsdp_multiple_forward_checkpoint.py b/tests/nn/data_parallel/test_fsdp_multiple_forward_checkpoint.py
index b4127bc..6776e79 100644
--- a/tests/nn/data_parallel/test_fsdp_multiple_forward_checkpoint.py
+++ b/tests/nn/data_parallel/test_fsdp_multiple_forward_checkpoint.py
@@ -56,10 +56,10 @@ class Model2(nn.Module):
 
     def __init__(self):
         super().__init__()
-        self.block1 = nn.Sequential(nn.Conv2d(3, 4, kernel_size=3), nn.BatchNorm2d(4), nn.ReLU(inplace=True))
-        self.block2 = nn.Sequential(nn.Conv2d(4, 4, kernel_size=3), nn.BatchNorm2d(4), nn.ReLU(inplace=False))
-        self.block3 = nn.Sequential(nn.Conv2d(4, 8, kernel_size=3), nn.BatchNorm2d(8), nn.ReLU(inplace=True))
-        self.head = nn.Sequential(nn.AdaptiveAvgPool2d(output_size=(1, 1)), nn.Flatten(), nn.Linear(8, 10))
+        self.block1 = nn.Sequential(nn.Conv2d(3, 400, kernel_size=3), nn.BatchNorm2d(400), nn.ReLU(inplace=True))
+        self.block2 = nn.Sequential(nn.Conv2d(400, 400, kernel_size=3), nn.BatchNorm2d(400), nn.ReLU(inplace=False))
+        self.block3 = nn.Sequential(nn.Conv2d(400, 800, kernel_size=3), nn.BatchNorm2d(800), nn.ReLU(inplace=True))
+        self.head = nn.Sequential(nn.AdaptiveAvgPool2d(output_size=(1, 1)), nn.Flatten(), nn.Linear(800, 10))
 
     def forward(self, x):
         if isinstance(x, torch.Tensor):
@@ -91,9 +91,9 @@ def _create_model(
             if with_model2:
                 model.block3 = auto_wrap_bn(model.block3, single_rank_pg=False, fsdp_config=fsdp_config)
         if with_checkpoint:
-            model.block2 = checkpoint_wrapper(model.block2, maintain_forward_counter=True)
+            model.block2 = checkpoint_wrapper(model.block2, maintain_forward_counter=False)
             if with_model2:
-                model.block3 = checkpoint_wrapper(model.block3, maintain_forward_counter=True)
+                model.block3 = checkpoint_wrapper(model.block3, maintain_forward_counter=False)
         with enable_wrap(
             wrapper_cls=FSDP,
             flatten_parameters=flatten,
```

Alternatively, force synchronous reduction, and the unit test will fail with small weights.
```
diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index f9a75c8..b77bc08 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -1360,6 +1360,8 @@ class FullyShardedDataParallel(nn.Module):
                 assert self._reducer is not None
                 grad_chunks = chunk_and_pad(param.grad.data, self.world_size)
                 self._reducer.reduce_scatter_async(grad_chunks, group=self.process_group, callback_fn=callback_fn)
+                self._reducer.flush()
+                self._streams["post_backward"].synchronize()
             else:
                 # Currently the only way for _is_sharded to be False is if
                 # world_size == 1. This could be relaxed in the future, in which
```

Example error:

`NCCL_P2P_DISABLE=1 python -m pytest -s tests/nn/data_parallel/test_fsdp_multiple_forward_checkpoint.py::test_multiple_forward_checkpoint[bn-model2-no_auto_wrap_bn-flatten-full]`
```
E       torch.multiprocessing.spawn.ProcessRaisedException: 
E       
E       -- Process 1 terminated with the following error:
E       Traceback (most recent call last):
E         File "/home/darryl/miniconda3/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 59, in _wrap
E           fn(i, *args)
E         File "/home/darryl/Projects/fairscale/tests/nn/data_parallel/test_fsdp_multiple_forward_checkpoint.py", line 200, in _distributed_worker
E           loss.backward()
E         File "/home/darryl/miniconda3/lib/python3.8/site-packages/torch/_tensor.py", line 255, in backward
E           torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
E         File "/home/darryl/miniconda3/lib/python3.8/site-packages/torch/autograd/__init__.py", line 147, in backward
E           Variable._execution_engine.run_backward(
E         File "/home/darryl/miniconda3/lib/python3.8/site-packages/torch/autograd/function.py", line 87, in apply
E           return self._forward_cls.backward(self, *args)  # type: ignore[attr-defined]
E         File "/home/darryl/Projects/fairscale/fairscale/nn/checkpoint/checkpoint_activations.py", line 323, in backward
E           torch.autograd.backward(outputs_with_grad, args_with_grad)
E         File "/home/darryl/miniconda3/lib/python3.8/site-packages/torch/autograd/__init__.py", line 147, in backward
E           Variable._execution_engine.run_backward(
E       RuntimeError: The size of tensor a (720600) must match the size of tensor b (1441200) at non-singleton dimension 0
```


Discussion:
Thanks again. Sorry about the not being able to reopen. I didn't know that you don't have the ability to reopen. Opening a new issue is probably better for clarity anyway. If you have the PR, feel free to send it even if it is in draft mode. Maybe we can review and discuss better around the code.
So the "race condition" is really between the reducer bucket filling, and the rest of the backward pass. You can also replicate the error by setting `maintain_forward_counter=False` and `bucket_cap_mb=0`.

I'll have a PR tomorrow that fixes this robustly.
Thanks again. Sorry about the not being able to reopen. I didn't know that you don't have the ability to reopen. Opening a new issue is probably better for clarity anyway. If you have the PR, feel free to send it even if it is in draft mode. Maybe we can review and discuss better around the code.
So the "race condition" is really between the reducer bucket filling, and the rest of the backward pass. You can also replicate the error by setting `maintain_forward_counter=False` and `bucket_cap_mb=0`.

I'll have a PR tomorrow that fixes this robustly.

Resolving patch:
diff --git a/fairscale/nn/checkpoint/checkpoint_activations.py b/fairscale/nn/checkpoint/checkpoint_activations.py
--- a/fairscale/nn/checkpoint/checkpoint_activations.py
+++ b/fairscale/nn/checkpoint/checkpoint_activations.py
@@ -16,7 +16,7 @@
 
 from fairscale.utils.containers import pack_kwargs, split_non_tensors, unpack_kwargs, unpack_non_tensors
 
-from .checkpoint_utils import dec_counter, inc_counter, init_counter, patch_batchnorm
+from .checkpoint_utils import patch_batchnorm
 
 
 # https://docs.python.org/3/library/threading.html#thread-local-data
@@ -95,9 +95,7 @@ def forward(self, input):
     return thread_local.is_recomputing
 
 
-def checkpoint_wrapper(
-    module: nn.Module, offload_to_cpu: bool = False, maintain_forward_counter: bool = False
-) -> nn.Module:
+def checkpoint_wrapper(module: nn.Module, offload_to_cpu: bool = False,) -> nn.Module:
     """
     A friendlier wrapper for performing activation checkpointing.
 
@@ -139,10 +137,6 @@ def checkpoint_wrapper(
             The module to be wrapped
         offload_to_cpu (bool):
             Whether to offload activations to CPU.
-        maintain_forward_counter (bool):
-            If True, maintain a forward counter per inner module. The counter will first
-            increases in forward calls of outer forward pass and then decreases in the
-            forward calls of outer backward pass. It is used by FullyShardedDataParallel.
 
     Returns:
         (nn.Module):
@@ -151,9 +145,6 @@ def checkpoint_wrapper(
     # Patch the batchnorm layers in case there are any in this module.
     patch_batchnorm(module)
 
-    if maintain_forward_counter:
-        init_counter(module)
-
     # The use of weakref here is to prevent creating a ref cycle: m -> m.forward -> m.
     # When such cycle exists, gc won't collect the module when the module is freed.
     # That causes GPU memory to be leaked. See the unit test for how we catch that.
@@ -172,10 +163,6 @@ def _checkpointed_forward(
     module = weak_self()
 
     # If gradients are disabled, just use original `.forward()` method directly.
-    # Doing so also ensures the internal fwd counter is not incremented in the forward pass,
-    # which would be an issue during eval since there wouldn't be a corresponding backward pass
-    # to decrement the fwd counter.
-    # See https://github.com/facebookresearch/fairscale/pull/709.
     if not torch.is_grad_enabled() or thread_local.is_checkpointing_disabled:
         return original_forward(module, *args, **kwargs)
 
@@ -270,8 +257,6 @@ def forward(  # type: ignore
         with torch.no_grad(), enable_checkpointing():
             unpacked_args, unpacked_kwargs = unpack_kwargs(kwarg_keys, args)
             outputs = run_function(*unpacked_args, **unpacked_kwargs)
-            the_module = unpacked_args[0]
-            inc_counter(the_module)
 
         if not isinstance(outputs, torch.Tensor):
             # Autograd Functions don't like non-Tensor outputs. We can split the
@@ -304,8 +289,6 @@ def backward(ctx: Any, *args: Any) -> Tuple[Optional[Tensor], ...]:
             unpacked_args, unpacked_kwargs = unpack_kwargs(ctx.kwarg_keys, inputs)
             outputs = ctx.run_function(*unpacked_args, **unpacked_kwargs)
             tensor_outputs, _ = split_non_tensors(outputs)
-            the_module = unpacked_args[0]
-            dec_counter(the_module)
 
         # Set the states back to what it was at the start of this function.
         set_rng_state(bwd_rng_state)
diff --git a/fairscale/nn/checkpoint/checkpoint_utils.py b/fairscale/nn/checkpoint/checkpoint_utils.py
--- a/fairscale/nn/checkpoint/checkpoint_utils.py
+++ b/fairscale/nn/checkpoint/checkpoint_utils.py
@@ -48,28 +48,3 @@ def post_forward(module: _BatchNorm, input: Tensor, result: Tensor) -> None:
             post_handle = child.register_forward_hook(post_forward)
             hooks += [pre_handle, post_handle]
     return hooks
-
-
-def init_counter(module: nn.Module) -> None:
-    """Add a checkpoint forward pass counter to a module and all its child FSDP modules.
-
-       ``inc_counter`` and ``dec_counter`` are used together with this to maintain counters
-       for FSDP to use in case of multiple forward pass and checkpoint being used at the same time.
-    """
-    for mod in module.modules():
-        mod._checkpoint_fwd_counter = 0
-
-
-def _add_counter(module: nn.Module, value: int) -> None:
-    if not hasattr(module, "_checkpoint_fwd_counter"):
-        return
-    for mod in module.modules():
-        mod._checkpoint_fwd_counter += value
-
-
-def inc_counter(module: nn.Module) -> None:
-    _add_counter(module, 1)
-
-
-def dec_counter(module: nn.Module) -> None:
-    _add_counter(module, -1)
diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -369,7 +369,7 @@ def __init__(
             f"FSDP.__init__(done): total_init_time: {(init_end - init_start): .4f} num_params: {(sum(p.numel() for p in self.params))}"
         )
 
-        # Flag to guard multiple pre-backward hook being executed per iteration.
+        # Flag to guard against preparing gradients multiple times per iteration.
         # This is reset at the end of the backward pass.
         self._pre_backward_hook_has_run = False
 
@@ -1039,11 +1039,6 @@ def _set_is_root(self) -> None:
         # For children instances, if they are checkpointed, state will not be reset to
         # IDLE after each inner forward/backward.
         self.assert_state(TrainingState.IDLE)
-        # Check if the root instance is being checkpointed. It doesn't make sense to
-        # checkpoint the root instance since it won't save GPU memory.
-        assert (
-            getattr(self, "_checkpoint_fwd_counter", 0) == 0
-        ), "Is the root FSDP module wrapping an activation checkpointed module? If so, please remove that."
         # As the root, we now set all children instances to False and
         # give them a closure to try to queue a wait_for_post_backward.
         self.children_share_process_group = True
@@ -1206,16 +1201,19 @@ def _pre_backward_hook(*unused: Any) -> None:
 
             # Only run the ``self._prep_grads_for_backward`` once per iteration (i.e. in case
             # it is multiple outputs or multiple forward passes).
-            if self._pre_backward_hook_has_run:
-                return
-            self._pre_backward_hook_has_run = True
-
-            # Start of a backward pass for the first time in an iteration.
-            self.assert_state([TrainingState.IDLE, TrainingState.BACKWARD_PRE])
-            self.training_state = TrainingState.BACKWARD_PRE
-
-            # Prepare p.grad so that it is in the right shape, device, accumulated values, etc.
-            self._prep_grads_for_backward()
+            if not self._pre_backward_hook_has_run:
+                self._pre_backward_hook_has_run = True
+                # Start of a backward pass for the first time in an iteration.
+                self.assert_state([TrainingState.IDLE, TrainingState.BACKWARD_PRE])
+                # Prepare p.grad so that it is in the right shape, device, accumulated values, etc.
+                self._prep_grads_for_backward()
+
+            # Transition to BACKWARD_PRE state if currently IDLE. We can transition from BACKWARD_POST
+            # to IDLE when FSDP is within activation checkpointing and called multiple times, due to the
+            # extra forward pass for re-computation.
+            if self.training_state == TrainingState.IDLE:
+                self.training_state = TrainingState.BACKWARD_PRE
+            self.assert_state([TrainingState.BACKWARD_PRE, TrainingState.BACKWARD_POST])
 
         def _register_hook(t: torch.Tensor) -> torch.Tensor:
             if t.requires_grad:
@@ -1298,14 +1296,8 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
         the local optimizer only sees the relevant parameter shard.
         """
         # First hook callback will see PRE state. If we have multiple params,
-        # then subsequent hook callbacks will see POST state. When checkpoint
-        # fwd counter is used, IDLE is also possible since the pre-backward hook
-        # is not triggered (see ``auto_wrap_bn`` below, we have to use
-        # FSDP(checkpoint(conv, FSDP(bn), ...)), with reshard_after_forward=False).
-        if hasattr(self, "_checkpoint_fwd_counter"):
-            self.assert_state([TrainingState.BACKWARD_PRE, TrainingState.BACKWARD_POST, TrainingState.IDLE])
-        else:
-            self.assert_state([TrainingState.BACKWARD_PRE, TrainingState.BACKWARD_POST])
+        # then subsequent hook callbacks will see POST state.
+        self.assert_state([TrainingState.BACKWARD_PRE, TrainingState.BACKWARD_POST])
         self.training_state = TrainingState.BACKWARD_POST
         if param.grad is None:
             return
@@ -1313,13 +1305,6 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
         if param.grad.requires_grad:
             raise RuntimeError("FSDP only works with gradients that don't require gradients")
 
-        # If this is a checkpointed module, we check if the following
-        # counter reaches 0. If not, it is not the final backward call
-        # for this module yet. Therefore, we early return in that case.
-        if hasattr(self._fsdp_wrapped_module, "_checkpoint_fwd_counter"):
-            if self._fsdp_wrapped_module._checkpoint_fwd_counter != 0:
-                return
-
         if self._require_backward_grad_sync or self.reshard_after_forward:
             # Free full params. As a special case, we don't free the full params
             # when in a ``no_sync`` context (as inversely indicated by
@@ -1354,18 +1339,34 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                 # Average grad by world_size for consistency with PyTorch DDP.
                 param.grad.data.div_(self.gradient_predivide_factor)
 
-            callback_fn = functools.partial(self._post_reduction_hook, param)
             if param._is_sharded:
-                assert param._is_sharded
                 assert self._reducer is not None
-                grad_chunks = chunk_and_pad(param.grad.data, self.world_size)
+                # Save the unsharded grad for reduction. We will asynchronously accumulate the reduced gradient into
+                # param._saved_grad_shard. If this FSDP module was called multiple times it's possible that multiple
+                # gradient reductions will happen in an undefined order. But addition commutes, so this order doesn't
+                # matter, neglecting rounding.
+                grad = param.grad.data
+                # Clear grad on the tensor, so any repeated gradient computations do not interfere with this reduction.
+                #
+                # The effect on memory consumption is not usually significant. No extra memory is allocated if this
+                # module is called only once, reduction happens quickly, or the tensor is bucketed. If the module is
+                # called multiple times, and the backwards pass runs far enough ahead of the `post_backward` stream,
+                # then we can end up with multiple unsharded gradients allocated and queued for reduction.
+                #
+                # We could guard against this by using CUDA events (see record_event, wait_event in torch.cuda.Stream).
+                # This ensures the `default` stream will wait for the `post_backward` stream to complete the last
+                # reduction for this module, before scheduling additional reduction work. Then at most there are two
+                # unsharded gradients allocated; one for a pending reduction, and one for gradient computation.
+                param.grad = None
+                callback_fn = functools.partial(self._post_reduction_hook, param)
+                grad_chunks = chunk_and_pad(grad, self.world_size)
                 self._reducer.reduce_scatter_async(grad_chunks, group=self.process_group, callback_fn=callback_fn)
             else:
                 # Currently the only way for _is_sharded to be False is if
                 # world_size == 1. This could be relaxed in the future, in which
                 # case grads should be all-reduced here.
                 assert self.world_size == 1
-                callback_fn(param.grad.data)
+                self._post_reduction_hook(param, param.grad.data)
 
             # After _post_backward_hook returns, orig_grad_data will eventually
             # go out of scope, at which point it could otherwise be freed for
@@ -1377,33 +1378,36 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
     def _post_reduction_hook(self, param: Parameter, reduced_grad: torch.Tensor) -> None:
         """Hook to call on each param after the reduce-scatter."""
         assert torch.cuda.current_stream() == self._streams["post_backward"]
-        assert param.grad is not None
         self.assert_state(TrainingState.BACKWARD_POST)
-        param.grad.data = reduced_grad
         if self.gradient_postdivide_factor > 1:
             # Average grad by world_size for consistency with PyTorch DDP.
-            param.grad.data.div_(self.gradient_postdivide_factor)
+            reduced_grad.data.div_(self.gradient_postdivide_factor)
         # Cast grad to param's dtype (typically FP32). Note: we do this
         # before the move_grads_to_cpu step so that this entire hook remains
         # non-blocking. The downside is a bit more D2H transfer in that case.
         if self.mixed_precision:
-            orig_param_grad_data = param.grad.data
-            param.grad.data = param.grad.data.to(dtype=param.data.dtype)
+            orig_param_grad_data = reduced_grad.data
+            reduced_grad.data = reduced_grad.data.to(dtype=param.data.dtype)
             # Don't let this memory get reused until after the transfer.
             orig_param_grad_data.record_stream(torch.cuda.current_stream())
-        if hasattr(param, "_saved_grad_shard") and param._saved_grad_shard is not None:
-            assert (
-                param._saved_grad_shard.shape == param.grad.shape
-            ), f"{param._saved_grad_shard.shape} vs {param.grad.shape}"
-            param.grad.data += param._saved_grad_shard
-            delattr(param, "_saved_grad_shard")
-        # Optionally move gradients to CPU, typically used if one is running
-        # the optimizer on the CPU.
+
+        if param._is_sharded:
+            # Accumulate into the gradient shard.
+            if getattr(param, "_saved_grad_shard", None) is None:
+                param._saved_grad_shard = reduced_grad.data
+            else:
+                assert (
+                    param._saved_grad_shard.shape == reduced_grad.shape
+                ), f"{param._saved_grad_shard.shape} vs {reduced_grad.shape}"
+                param._saved_grad_shard.data += reduced_grad.data
+            reduced_grad = param._saved_grad_shard.data
+
+        # Optionally move gradients to CPU, typically used if one is running the optimizer on the CPU. Once the full
+        # backwards pass completes, we will set `.grad` to the CPU copy.
         if self.move_grads_to_cpu:
-            param._cpu_grad.copy_(param.grad.data, non_blocking=True)
+            param._cpu_grad.copy_(reduced_grad.data, non_blocking=True)
             # Don't let this memory get reused until after the transfer.
-            param.grad.data.record_stream(torch.cuda.current_stream())
-            param.grad.data = param._cpu_grad
+            reduced_grad.data.record_stream(torch.cuda.current_stream())
 
     def _queue_wait_for_post_backward(self) -> None:
         """Try to queue a `wait_for_post_backward` callback.
@@ -1446,19 +1450,38 @@ def _wait_for_post_backward(self) -> None:
         if self._reducer is not None:
             self._reducer.teardown()
 
-        def _remove_shard_bwd_hook(fsdp_module: FullyShardedDataParallel) -> None:
+        def _finalize_parameters(fsdp_module: FullyShardedDataParallel) -> None:
             """Helper used below on all fsdp modules."""
             for p in fsdp_module.params:
-                if p.requires_grad:
-                    if hasattr(p, "_shard_bwd_hook"):
-                        assert len(p._shard_bwd_hook) == 2, len(p._shard_bwd_hook)
-                        p._shard_bwd_hook[1].remove()
-                        delattr(p, "_shard_bwd_hook")
+                if not p.requires_grad:
+                    continue
+                if hasattr(p, "_shard_bwd_hook"):
+                    assert len(p._shard_bwd_hook) == 2, len(p._shard_bwd_hook)
+                    p._shard_bwd_hook[1].remove()
+                    delattr(p, "_shard_bwd_hook")
+
+                # Leave the gradient accumulation state as-is if not synchronizing this pass. This ensures p.grad
+                # remains the unsharded gradient accumulated from prior no-sync passes, and p._saved_grad_shard
+                # remains the sharded gradient from the last synchronized pass. This also allows interleaved no-sync and
+                # sync passes, if desired.
+                if not self._require_backward_grad_sync:
+                    continue
+
+                # Parameter and gradient devices must match.
+                if hasattr(p, "_cpu_grad"):
+                    assert p.device == torch.device("cpu")
+                    p.grad = p._cpu_grad
+                elif hasattr(p, "_saved_grad_shard"):
+                    assert p.device == p._saved_grad_shard.device
+                    p.grad = p._saved_grad_shard
+
+                if hasattr(p, "_saved_grad_shard"):
+                    delattr(p, "_saved_grad_shard")
 
         # Update root and nested FSDP's hooks and flags.
         for m in self.modules():  # includes self
             if isinstance(m, FullyShardedDataParallel):
-                _remove_shard_bwd_hook(m)
+                _finalize_parameters(m)
                 m._pre_backward_hook_has_run = False
                 if any(p.requires_grad for p in m.parameters()):
                     # Check if the module has params and if any of them has


Test patch:
diff --git a/tests/experimental/nn/test_sync_batchnorm.py b/tests/experimental/nn/test_sync_batchnorm.py
--- a/tests/experimental/nn/test_sync_batchnorm.py
+++ b/tests/experimental/nn/test_sync_batchnorm.py
@@ -106,7 +106,7 @@ def parity3d_checkpoint_syncbn():
     x = torch.randn(4, 3, 4, 4, 4).cuda() * rank
     torch_bn = torch.nn.SyncBatchNorm(3).cuda()
     fs_bn = SyncBatchNorm(3).cuda()
-    fs_bn = checkpoint_wrapper(fs_bn, maintain_forward_counter=True)
+    fs_bn = checkpoint_wrapper(fs_bn)
     check_parity_ddp(torch_bn, fs_bn, x)
 
 
diff --git a/tests/nn/data_parallel/test_fsdp_multiple_forward_checkpoint.py b/tests/nn/data_parallel/test_fsdp_multiple_forward_checkpoint.py
--- a/tests/nn/data_parallel/test_fsdp_multiple_forward_checkpoint.py
+++ b/tests/nn/data_parallel/test_fsdp_multiple_forward_checkpoint.py
@@ -70,7 +70,15 @@ def forward(self, x):
 
 
 def _create_model(
-    with_model2, with_sync_bn, with_fsdp, with_checkpoint, mixed_precision, flatten, wrap_bn, fp32_reduce_scatter
+    with_model2,
+    with_sync_bn,
+    with_fsdp,
+    with_checkpoint,
+    mixed_precision,
+    flatten,
+    wrap_bn,
+    fp32_reduce_scatter,
+    bucket_cap_mb,
 ):
     model = Model2() if with_model2 else Model()
     fsdp_config = None
@@ -84,33 +92,32 @@ def _create_model(
             "force_input_to_fp32": True,  # SyncBN needs this.
         }
 
+    if with_fsdp and wrap_bn:
+        model.block1 = auto_wrap_bn(model.block1, single_rank_pg=False, fsdp_config=fsdp_config)
+        model.block2 = auto_wrap_bn(model.block2, single_rank_pg=False, fsdp_config=fsdp_config)
+        if with_model2:
+            model.block3 = auto_wrap_bn(model.block3, single_rank_pg=False, fsdp_config=fsdp_config)
+
+    if with_checkpoint:
+        model.block2 = checkpoint_wrapper(model.block2)
+        if with_model2:
+            model.block3 = checkpoint_wrapper(model.block3)
+
     if with_fsdp:
-        if wrap_bn:
-            model.block1 = auto_wrap_bn(model.block1, single_rank_pg=False, fsdp_config=fsdp_config)
-            model.block2 = auto_wrap_bn(model.block2, single_rank_pg=False, fsdp_config=fsdp_config)
-            if with_model2:
-                model.block3 = auto_wrap_bn(model.block3, single_rank_pg=False, fsdp_config=fsdp_config)
-        if with_checkpoint:
-            model.block2 = checkpoint_wrapper(model.block2, maintain_forward_counter=True)
-            if with_model2:
-                model.block3 = checkpoint_wrapper(model.block3, maintain_forward_counter=True)
         with enable_wrap(
             wrapper_cls=FSDP,
             flatten_parameters=flatten,
             mixed_precision=mixed_precision,
             compute_dtype=torch.float32,
             fp32_reduce_scatter=fp32_reduce_scatter,
+            bucket_cap_mb=bucket_cap_mb,
         ):
             model.block1 = wrap(model.block1)
             model.block2 = wrap(model.block2)
             if with_model2:
                 model.block3 = wrap(model.block3)
             model.head = wrap(model.head)
-    else:
-        if with_checkpoint:
-            model.block2 = checkpoint_wrapper(model.block2, maintain_forward_counter=False)
-            if with_model2:
-                model.block3 = checkpoint_wrapper(model.block3, maintain_forward_counter=False)
+
     return model
 
 
@@ -126,6 +133,7 @@ def _distributed_worker(
     flatten,
     wrap_bn,
     fp32_reduce_scatter,
+    bucket_cap_mb,
 ):
     filename, filename_rpc = files[:2]
     filename_loss = files[2:]
@@ -155,7 +163,15 @@ def _distributed_worker(
         batch = [x.half() for x in batch]
 
     model = _create_model(
-        with_model2, with_sync_bn, with_fsdp, with_checkpoint, mixed_precision, flatten, wrap_bn, fp32_reduce_scatter
+        with_model2,
+        with_sync_bn,
+        with_fsdp,
+        with_checkpoint,
+        mixed_precision,
+        flatten,
+        wrap_bn,
+        fp32_reduce_scatter,
+        bucket_cap_mb,
     )
     model = model.cuda()
 
@@ -166,6 +182,7 @@ def _distributed_worker(
             mixed_precision=mixed_precision,
             compute_dtype=torch.float32,
             fp32_reduce_scatter=fp32_reduce_scatter,
+            bucket_cap_mb=bucket_cap_mb,
         )
         model.set_gradient_divide_factors(1.0, 2.0, True)
         no_sync_context = contextlib.suppress()
@@ -228,6 +245,7 @@ def _get_cached_results(
     flatten,
     wrap_bn,
     fp32_reduce_scatter,
+    bucket_cap_mb,
 ):
     """ Cache the training to save time. For DDP, flatten, wrap_bn etc. doesn't matter, so
         the results can be cached.
@@ -247,6 +265,7 @@ def _get_cached_results(
         flatten,
         wrap_bn,
         fp32_reduce_scatter,
+        bucket_cap_mb,
     )
     global _result_cache
     if key not in _result_cache:
@@ -265,6 +284,7 @@ def _get_cached_results(
                     flatten,
                     wrap_bn,
                     fp32_reduce_scatter,
+                    bucket_cap_mb,
                 ),
                 nprocs=world_size,
             )
@@ -310,39 +330,66 @@ def test_multiple_forward_checkpoint(precision, flatten, wrap_bn, model_type, bn
 
     world_size = 2
     expected_losses = None
-    # Ensure ddp == ddp+ckpt == fsdp == fsdp+ckpt.
+
+    # Ensure ddp == fsdp when modules are called multiple times per forward pass with/without checkpointing, forward
+    # counters and reducer bucketing.
+    #
+    # The bucketing check exists because the asynchronous gradient reduction it induces can interact with multiple
+    # forward passes in complex ways. For example, in the midst of a sharded backward pass, `parameter.grad` may only be
+    # `None` or an unsharded gradient tensor. The sharded tensor is then set at the end of the backwards pass. But a
+    # unit test with bucketing enabled might not catch violations of this invariant. For very small models, like the
+    # kind used in this unit test, bucketing will delay gradient reduction until after all the gradient computation is
+    # done. If the reduction incorrectly sets `.grad` to the _sharded_ variant, the test might not fail, since the
+    # gradient computations have already happened. Toggling bucketing helps verify that gradient reduction and
+    # computation interact correctly.
+    combinations = []
     for with_fsdp in [False, True]:
         for with_checkpoint in [False, True]:
             if not with_fsdp and with_checkpoint:
                 continue
-            final_losses = _get_cached_results(
-                world_size,
-                with_model2,
-                with_sync_bn,
-                with_fsdp,
-                with_checkpoint,
-                mixed_precision,
-                flatten,
-                wrap_bn,
-                fp32_reduce_scatter,
-            )
-            if expected_losses is None:
-                expected_losses = final_losses
-            else:
-                print(f"checking: fsdp {with_fsdp} ckpt {with_checkpoint} with ddp+no_ckpt")
-
-                def check(exp, res):
-                    assert list(exp.keys()) == list(res.keys()), f"{list(exp.keys())} vs. {list(res.keys())}"
-                    rtol = 1e-4
-                    atol = 1e-5
-                    if with_model2 and mixed_precision and torch_version() >= (1, 9, 0):
-                        # On CI, with longer model2, mixed precsion and 1.9, even ddp vs. ddp+ckpt has
-                        # larger errors.
-                        rtol = 1e-3
-                        atol = 1e-4
-                    for key in exp.keys():
-                        exp_loss = exp[key]
-                        res_loss = res[key]
-                        torch.testing.assert_allclose(exp_loss, res_loss, rtol=rtol, atol=atol)
-
-                check(expected_losses, final_losses)
+            for with_bucketing in [False, True]:
+                if not with_fsdp and with_bucketing:
+                    continue
+                combinations.append((with_fsdp, with_checkpoint, with_bucketing))
+    print("")
+    print("Testing the following configurations:")
+    for with_fsdp, with_checkpoint, with_bucketing in combinations:
+        print(f"  fsdp {with_fsdp} ckpt {with_checkpoint} bucketing {with_bucketing}")
+
+    for with_fsdp, with_checkpoint, with_bucketing in combinations:
+        if with_bucketing:
+            bucket_cap_mb = 25
+        else:
+            bucket_cap_mb = 0
+        final_losses = _get_cached_results(
+            world_size,
+            with_model2,
+            with_sync_bn,
+            with_fsdp,
+            with_checkpoint,
+            mixed_precision,
+            flatten,
+            wrap_bn,
+            fp32_reduce_scatter,
+            bucket_cap_mb,
+        )
+        if expected_losses is None:
+            expected_losses = final_losses
+        else:
+            print(f"checking: fsdp {with_fsdp} ckpt {with_checkpoint} bucketing {with_bucketing} with ddp+no_ckpt")
+
+            def check(exp, res):
+                assert list(exp.keys()) == list(res.keys()), f"{list(exp.keys())} vs. {list(res.keys())}"
+                rtol = 1e-4
+                atol = 1e-5
+                if with_model2 and mixed_precision and torch_version() >= (1, 9, 0):
+                    # On CI, with longer model2, mixed precsion and 1.9, even ddp vs. ddp+ckpt has
+                    # larger errors.
+                    rtol = 1e-3
+                    atol = 1e-4
+                for key in exp.keys():
+                    exp_loss = exp[key]
+                    res_loss = res[key]
+                    torch.testing.assert_allclose(exp_loss, res_loss, rtol=rtol, atol=atol)
+
+            check(expected_losses, final_losses)
diff --git a/tests/nn/data_parallel/test_fsdp_with_checkpoint_wrapper.py b/tests/nn/data_parallel/test_fsdp_with_checkpoint_wrapper.py
--- a/tests/nn/data_parallel/test_fsdp_with_checkpoint_wrapper.py
+++ b/tests/nn/data_parallel/test_fsdp_with_checkpoint_wrapper.py
@@ -71,12 +71,10 @@ def _train_step(model, optim, expected_param_shapes):
     # Create input and run forward pass.
     input = torch.randn(2, 3).cuda()
     loss = model(input).sum()
-    _check_fwd_counter(model, 1)
     _check_params(model, expected_param_shapes)
 
     # Run backward pass.
     loss.backward()
-    _check_fwd_counter(model, 0)
     _check_params(model, expected_param_shapes)
 
     # Finally, take a step.
@@ -90,7 +88,6 @@ def _eval_step(model, optim, expected_param_shapes):
     with torch.no_grad():
         input = torch.randn(2, 3).cuda()
         model(input).sum()
-    _check_fwd_counter(model, 0)
     _check_params(model, expected_param_shapes)
 
 
@@ -104,20 +101,11 @@ def _check_params(model, expected_param_shapes):
         ), f"Parameter {key} should have shape {expected_shape}, but found shape {current_shape}"
 
 
-def _check_fwd_counter(model, expected_value):
-    current_value = model._fpw_module.ffn[1]._fsdp_wrapped_module.module._checkpoint_fwd_counter
-    assert (
-        current_value == expected_value
-    ), f"forward counter of checkpointed submodule should be {expected_value}, but found {current_value}"
-
-
 class SimpleModuleWithCheckpointing(nn.Module):
     def __init__(self):
         super().__init__()
         self.ffn = nn.Sequential(
-            nn.Linear(3, 3),
-            FullyShardedDataParallel(checkpoint_wrapper(nn.Linear(3, 3), maintain_forward_counter=True)),
-            nn.Linear(3, 3),
+            nn.Linear(3, 3), FullyShardedDataParallel(checkpoint_wrapper(nn.Linear(3, 3))), nn.Linear(3, 3),
         )
 
     def forward(self, x):