From ac5e71269d14c06496e9543aa5a6d1ccded8e4df Mon Sep 17 00:00:00 2001
From: {{Naman Goyal}} <{{naman@fb.com}}>
Date: Mon, 11 Apr 2022 11:05:29 -0700
Subject: [PATCH 01/15] fix for high gpu reserved memory

---
 .../fully_sharded_data_parallel.py            | 30 +++++++++++++++++--
 1 file changed, 27 insertions(+), 3 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index fef71dce4..952031351 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -1156,6 +1156,8 @@ def _reset_lazy_init(self) -> None:
         self._is_root: Optional[bool] = None
         self._streams: Dict[str, torch.cuda.Stream] = {}
         self._reducer: Optional[ReduceScatterBucketer] = None
+        self._fsdp_forward_ordering: List[nn.Module] = []
+        self._my_fsdp_instance_idx: Optional[int] = None
         for p in self.params:
             if hasattr(p, "_fp32_shard"):
                 del p._fp32_shard  # reset _init_param_attributes
@@ -1327,6 +1329,7 @@ def _set_is_root(self) -> None:
                 m.no_broadcast_optim_state = m.no_broadcast_optim_state or (
                     (m.world_size == 1) and (m.world_size < self.world_size) and (m.process_group != self.process_group)
                 )
+                m._fsdp_forward_ordering = self._fsdp_forward_ordering
 
     def _setup_streams(self) -> None:
         """Create streams to overlap data transfer and computation."""
@@ -1386,6 +1389,10 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
         if self._is_root and self.mixed_precision:
             args, kwargs = cast_floats_to_right_precision(True, True, *args, **kwargs)
 
+        if self not in self._fsdp_forward_ordering:
+            self._my_fsdp_instance_idx = len(self._fsdp_forward_ordering)
+            self._fsdp_forward_ordering.append(self)
+
         # If enabled, convert the input to FP32 if we are in full precision.
         # no_grad is not used because the input might be for a non-root instance,
         # which mean autograd needs to go through the conversion.
@@ -1396,6 +1403,14 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
         # ``self.compute_dtype`` (e.g., FP16 if *mixed_precision* is ``True``).
         self._rebuild_full_params()
 
+        if (
+            self._fsdp_forward_ordering is not None
+            and self._my_fsdp_instance_idx is not None and self._my_fsdp_instance_idx < len(self._fsdp_forward_ordering) - 1
+        ):
+            self._fsdp_forward_ordering[self._my_fsdp_instance_idx + 1]._rebuild_full_params(
+                wait_for_all_gather=False
+            )
+
         # Register backward hooks to reshard params and reduce-scatter grads.
         # These need to be re-registered every forward pass.
         self._register_post_backward_hooks()
@@ -1484,6 +1499,12 @@ def _pre_backward_hook(*unused: Any) -> None:
             # overhead.
             if self.reshard_after_forward:
                 self._rebuild_full_params()
+                if (
+                    self.reshard_after_forward
+                    and self._fsdp_forward_ordering is not None
+                    and self._my_fsdp_instance_idx is not None and self._my_fsdp_instance_idx > 0
+                ):
+                    self._fsdp_forward_ordering[self._my_fsdp_instance_idx - 1]._rebuild_full_params(wait_for_all_gather=False)
             else:
                 self._use_full_params()
 
@@ -1847,7 +1868,7 @@ def _finalize_parameters(fsdp_module: FullyShardedDataParallel) -> None:
                     self._output_pre_backward_hook_registered.clear()
 
     @torch.no_grad()
-    def _rebuild_full_params(self, force_full_precision: bool = False) -> Optional[List[Tuple[torch.Tensor, bool]]]:
+    def _rebuild_full_params(self, force_full_precision: bool = False, wait_for_all_gather = True) -> Optional[List[Tuple[torch.Tensor, bool]]]:
         """
         Gather all shards of params.
 
@@ -1916,6 +1937,8 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
 
         # Early exit if we already have full params and don't need full precision.
         if self.has_full_params and not force_full_precision:
+            if wait_for_all_gather:
+                torch.cuda.current_stream().wait_stream(self._streams["all_gather"])
             for p in self.params:
                 update_p_data()
             return output_tensors
@@ -1978,8 +2001,8 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
 
                     if self.move_params_to_cpu and (self.params[0].dtype == self.compute_dtype):
                         self._free_fp16_param_shard([p])
-
-        torch.cuda.current_stream().wait_stream(self._streams["all_gather"])
+        if wait_for_all_gather:
+            torch.cuda.current_stream().wait_stream(self._streams["all_gather"])
         return output_tensors
 
     @torch.no_grad()
@@ -2047,6 +2070,7 @@ def _free_full_params(self, params: Optional[List[Parameter]] = None) -> None:
             # Storage object and unshard it in-place. For now, just resize
             # the Storage to 0 to save memory.
             free_storage_(p._full_param_padded)
+            torch.cuda.current_stream().synchronize()
 
     def local_metadata_dict(self) -> Dict[str, Any]:
         """

From 191553190d73a5ef4a48687c889d4b1d94532135 Mon Sep 17 00:00:00 2001
From: Stephen Roller <roller@fb.com>
Date: Wed, 11 May 2022 15:24:29 +0000
Subject: [PATCH 02/15] Get rid of warning

---
 fairscale/nn/data_parallel/fully_sharded_data_parallel.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 952031351..be874e125 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -311,7 +311,7 @@ def __init__(
         module: nn.Module,
         process_group: Optional[ProcessGroup] = None,
         # The type for the process_group_reduce_scatter only can be either ProcessGroup or ProcessGroupName
-        process_group_reduce_scatter: Any = ProcessGroupName.reduce_scatter,
+        process_group_reduce_scatter: Any = ProcessGroupName.default,
         reshard_after_forward: bool = True,
         disable_reshard_on_root: bool = True,
         mixed_precision: bool = False,

From ba38cf3235378607035791b98d318ffaf6f6b2af Mon Sep 17 00:00:00 2001
From: namangoyal <naman.goyal21@gmail.com>
Date: Tue, 28 Jun 2022 02:12:31 +0000
Subject: [PATCH 03/15] bf16 changes and attribute for cpu activations

---
 .../data_parallel/fully_sharded_data_parallel.py   | 14 +++++++++-----
 fairscale/nn/misc/flatten_params_wrapper.py        |  2 ++
 tests/nn/data_parallel/test_fsdp.py                | 11 +++++++++++
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index be874e125..bf3e6c018 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -1386,8 +1386,9 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
 
         # For root and mixed precision, we convert the input to FP16 (no_grad is needed for
         # the conversion).
+        is_bf16 = self.compute_dtype == torch.bfloat16
         if self._is_root and self.mixed_precision:
-            args, kwargs = cast_floats_to_right_precision(True, True, *args, **kwargs)
+            args, kwargs = cast_floats_to_right_precision(True, True, is_bf16, *args, **kwargs)
 
         if self not in self._fsdp_forward_ordering:
             self._my_fsdp_instance_idx = len(self._fsdp_forward_ordering)
@@ -1397,7 +1398,7 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
         # no_grad is not used because the input might be for a non-root instance,
         # which mean autograd needs to go through the conversion.
         if self.force_input_to_fp32 and not self.mixed_precision:
-            args, kwargs = cast_floats_to_right_precision(False, False, *args, **kwargs)
+            args, kwargs = cast_floats_to_right_precision(False, False, is_bf16, *args, **kwargs)
 
         # All-gather full parameters. This will also transfer FP32 parameters to
         # ``self.compute_dtype`` (e.g., FP16 if *mixed_precision* is ``True``).
@@ -2054,6 +2055,7 @@ def _free_full_params(self, params: Optional[List[Parameter]] = None) -> None:
         if params is None:
             params = self.params
         self.has_full_params = False
+
         current_stream = torch.cuda.current_stream()
         for p in params:
             if not p._is_sharded:  # e.g., world_size == 1
@@ -2182,7 +2184,6 @@ def consolidate_shard_weights(
                 for n, t, s in zip(names, full_param.split(numels), shapes):
                     out_state_dict_key = ".".join([fsdp_path, n]) if fsdp_path else n
                     consolidated_weights[out_state_dict_key] = t.view(s)
-
         # copy shared parameters
         for src_path, dest_path in metadata["shared_param_info"]:
             consolidated_weights[dest_path] = consolidated_weights[src_path]
@@ -2462,7 +2463,7 @@ def _get_default_cuda_device(module: nn.Module) -> torch.device:
     return torch.device("cuda")
 
 
-def cast_floats_to_right_precision(to_fp16: bool, no_grad: bool, *args: Any, **kwargs: Any) -> Tuple[Any, Any]:
+def cast_floats_to_right_precision(to_fp16: bool, no_grad: bool, is_bf16: bool, *args: Any, **kwargs: Any) -> Tuple[Any, Any]:
     """
     Cast floating point Tensors in *args or **kwargs to FP16 or FP32 if they are not.
     We also retain the requires_grad flag so that casting doesn't affect the autograd graph.
@@ -2470,7 +2471,10 @@ def cast_floats_to_right_precision(to_fp16: bool, no_grad: bool, *args: Any, **k
 
     def fn_fp16(x: torch.Tensor) -> torch.Tensor:
         if x.dtype is torch.float32:
-            y = x.half()
+            if is_bf16:
+                y = x.bfloat16()
+            else:
+                y = x.half()
             if x.is_leaf:
                 y.requires_grad = x.requires_grad
             return y
diff --git a/fairscale/nn/misc/flatten_params_wrapper.py b/fairscale/nn/misc/flatten_params_wrapper.py
index 0cbdce2bb..38265dd2b 100644
--- a/fairscale/nn/misc/flatten_params_wrapper.py
+++ b/fairscale/nn/misc/flatten_params_wrapper.py
@@ -372,6 +372,7 @@ def _unflatten_params_as_views(self) -> None:
         ps = self.get_param_views()
         param_views = []
         for (_, m, n), p in zip(self._param_infos, ps):
+            setattr(p, '_fsdp_weight', True)
             setattr(m, n, p)  # This will set as plain attr
             param_views.append(p)
 
@@ -382,6 +383,7 @@ def _unflatten_params_as_views(self) -> None:
         for (_, _, m, n, shared_m, shared_n) in self._shared_param_infos:
             setattr(m, n, getattr(shared_m, shared_n))
 
+
     @contextmanager
     def unflatten_params(self, flat_params: Optional[List[Tensor]] = None) -> Generator:
         """
diff --git a/tests/nn/data_parallel/test_fsdp.py b/tests/nn/data_parallel/test_fsdp.py
index 7313bf262..06f96a9db 100644
--- a/tests/nn/data_parallel/test_fsdp.py
+++ b/tests/nn/data_parallel/test_fsdp.py
@@ -211,6 +211,17 @@ def test_mixed_precision_autocast_fp32_compute(self):
             expected_buffer_type=torch.float32,
         )
 
+    def test_mixed_precision_bfloat16(self):
+        self._spawn_test_case(
+            {"mixed_precision": True, "compute_dtype": torch.bfloat16},
+            True,  # autocast enabled
+            torch.bfloat16,  # expected_input_dtype
+            torch.bfloat16,  # expected_param_dtype
+            torch.float32,  # expected_loss_dtype
+            torch.bfloat16,  # expected_reduce_dtype
+            expected_buffer_type=torch.float32,
+        )
+
     def test_fp32_reduce_scatter(self):
         self._spawn_test_case(
             {"mixed_precision": True, "fp32_reduce_scatter": True},

From 8e5c41625af432f2e0864db57445ac9165eae050 Mon Sep 17 00:00:00 2001
From: Naman Goyal <naman@fb.com>
Date: Fri, 22 Sep 2023 12:28:09 -0700
Subject: [PATCH 04/15] changes for pp

---
 .../fully_sharded_data_parallel.py            | 102 ++++++++++++++++--
 fairscale/utils/reduce_scatter_bucketer.py    |  12 ++-
 2 files changed, 107 insertions(+), 7 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index bf3e6c018..855a2b8bc 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -5,6 +5,7 @@
 
 import contextlib
 import copy
+import collections
 from dataclasses import dataclass
 from enum import Enum, auto
 import functools
@@ -30,6 +31,7 @@
     Tuple,
     Union,
     cast,
+    Deque,
 )
 
 import torch
@@ -75,6 +77,38 @@
     pass
 
 
+from logging import getLogger
+logger = getLogger()
+
+class _FreeEventQueue:
+    """
+    This tracks all pending frees corresponding to inflight all-gathers. The
+    queueing pattern is iterative enqueues with a single dequeue per iteration
+    once the limit ``_max_num_inflight_all_gathers`` is reached.
+    """
+
+    def __init__(self) -> None:
+        self._queue: Deque[torch.cuda.Event] = collections.deque()
+        self._max_num_inflight_all_gathers = 0  # empirically chosen
+
+    def enqueue(self, free_event: torch.cuda.Event) -> None:
+        """Enqueues a free event."""
+        self._queue.append(free_event)
+
+    def dequeue_if_needed(self) -> Optional[torch.cuda.Event]:
+        """Dequeues a single event if the limit is reached."""
+        if len(self._queue) >= self._max_num_inflight_all_gathers:
+            return self._dequeue()
+        return None
+
+    def _dequeue(self) -> Optional[torch.cuda.Event]:
+        """Dequeues a free event if possible."""
+        if self._queue:
+            event = self._queue.popleft()
+            return event
+        return None
+
+
 class TrainingState(Enum):
     """
     Simple enum to indicate what state FSDP is in. Used for asserting
@@ -332,6 +366,8 @@ def __init__(
         offload_config: Optional[OffloadConfig] = None,
         state_dict_on_rank_0_only: bool = False,
         gradient_predivide_factor: Optional[float] = None,
+        limit_all_gather_events: bool = False,
+        limit_reduce_scatter_events: bool = False,
     ):
         try:
             import torch._C
@@ -518,6 +554,10 @@ def __init__(
                 if isinstance(m, FullyShardedDataParallel):
                     m._free_ssd_offload()
 
+        self.dont_wait_current_stream_for_post_all_gather = False
+        self._all_gather_free_event_queue = _FreeEventQueue() if limit_all_gather_events else None
+        self._reduce_scatter_free_event_queue = _FreeEventQueue() if limit_reduce_scatter_events else None
+
     def _get_gradient_predivide_factor(self, world_size: int) -> float:
         factor: int = 1
         while world_size % factor == 0 and world_size / factor > factor:
@@ -1330,6 +1370,8 @@ def _set_is_root(self) -> None:
                     (m.world_size == 1) and (m.world_size < self.world_size) and (m.process_group != self.process_group)
                 )
                 m._fsdp_forward_ordering = self._fsdp_forward_ordering
+                m._reduce_scatter_free_event_queue = self._reduce_scatter_free_event_queue
+                m._all_gather_free_event_queue = self._all_gather_free_event_queue
 
     def _setup_streams(self) -> None:
         """Create streams to overlap data transfer and computation."""
@@ -1346,7 +1388,7 @@ def _setup_streams(self) -> None:
 
         # Helper for bucketing reduce-scatter ops. This is also shared with
         # children instances to improve bucket utilization.
-        self._reducer = ReduceScatterBucketer(self.bucket_cap_mb)
+        self._reducer = ReduceScatterBucketer(self.bucket_cap_mb, self._reduce_scatter_free_event_queue)
         # We share streams with all children instances, which allows them to
         # overlap transfers across the forward pass without synchronizing with
         # the default stream.
@@ -1600,8 +1642,8 @@ def _register_post_backward_hooks(self) -> None:
             return  # don't register grad hooks if grad isn't enabled
         for p in self.params:
             if p.requires_grad:
-                if hasattr(p, "_shard_bwd_hook"):
-                    continue
+                # if hasattr(p, "_shard_bwd_hook"):
+                #     continue
                 # Register a hook on the first call, empirically, autograd
                 # fires it at the end for this param, which makes sense.
                 p_tmp = p.expand_as(p)  # Get a grad_fn on p_tmp.
@@ -1723,6 +1765,11 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
             # are underway in the post_backward stream. See:
             # github.com/NVIDIA/apex/blob/master/apex/parallel/distributed.py
             orig_grad_data.record_stream(self._streams["post_backward"])
+            if self._reduce_scatter_free_event_queue is not None:
+                release_full_grad_event = torch.cuda.Event()
+                release_full_grad_event.record()
+                self._reduce_scatter_free_event_queue.enqueue(release_full_grad_event)
+
 
     def _post_reduction_hook(self, param: Parameter, reduced_grad: torch.Tensor) -> None:
         """Hook to call on each param after the reduce-scatter."""
@@ -1784,6 +1831,9 @@ def _wait_for_post_backward(self) -> None:
         else:
             self.assert_state(TrainingState.BACKWARD_PRE)
 
+        if self.dont_wait_current_stream_for_post_all_gather:
+            return
+
         if self._require_backward_grad_sync:
             # Flush any unreduced buckets in the post_backward stream.
             with torch.cuda.stream(self._streams["post_backward"]):
@@ -1808,8 +1858,8 @@ def _finalize_parameters(fsdp_module: FullyShardedDataParallel) -> None:
                     continue
                 if hasattr(p, "_shard_bwd_hook"):
                     p_assert(len(p._shard_bwd_hook) == 2, f"WFPB: incorrect hook num: {len(p._shard_bwd_hook)}")
-                    p._shard_bwd_hook[1].remove()
-                    delattr(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
@@ -1817,7 +1867,6 @@ def _finalize_parameters(fsdp_module: FullyShardedDataParallel) -> None:
                 # sync passes, if desired.
                 if not self._require_backward_grad_sync:
                     continue
-
                 # Parameter and gradient devices must match.
                 if hasattr(p, "_cpu_grad"):
                     p_assert(p.device == torch.device("cpu"), f"WFPB: incorrect cpu_grad device {p.device}")
@@ -1868,6 +1917,32 @@ def _finalize_parameters(fsdp_module: FullyShardedDataParallel) -> None:
                     assert self._output_pre_backward_hook_registered is not None  # make mypy happy
                     self._output_pre_backward_hook_registered.clear()
 
+
+    @torch.no_grad()
+    def _rebuild_full_params_recursive(self):
+        # if recurse:
+        #     with contextlib.ExitStack() as stack:
+        #         # Summon all params for any nested FSDP instances.
+        #         for module in self.modules():
+        #             if isinstance(module, FullyShardedDataParallel):
+        #                 stack.enter_context(module.summon_full_params(recurse=False, volatile=volatile))
+        #         # Yield to the caller, with full params in all nested instances.
+        #         yield
+        #     # Exiting from the ExitStack will re-shard params.
+        #     return
+        # else:
+        # torch.cuda.synchronize()
+        # self.assert_state(TrainingState.IDLE)
+        # Set the state so that we assert when trying to go into
+        # forward/backward.
+        # full_tensors = self._rebuild_full_params(force_full_precision=True)
+        for module in self.modules():
+            if isinstance(module, FullyShardedDataParallel):
+                module._lazy_init()
+                module._rebuild_full_params(wait_for_all_gather=False)
+
+
+
     @torch.no_grad()
     def _rebuild_full_params(self, force_full_precision: bool = False, wait_for_all_gather = True) -> Optional[List[Tuple[torch.Tensor, bool]]]:
         """
@@ -1946,6 +2021,14 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
 
         self.has_full_params = True
 
+        if self._all_gather_free_event_queue is not None:
+            while self._all_gather_free_event_queue._queue and self._all_gather_free_event_queue._queue[0].query():
+                self._all_gather_free_event_queue._dequeue()
+            event = self._all_gather_free_event_queue.dequeue_if_needed()
+            if event:
+                logger.warning("synching cpu thread for AG limit")
+                event.synchronize()
+
         with torch.cuda.stream(self._streams["all_gather"]):
             if (self.mixed_precision or self.move_params_to_cpu) and not force_full_precision:
                 self._cast_fp32_param_shards_to_fp16()
@@ -2002,8 +2085,15 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
 
                     if self.move_params_to_cpu and (self.params[0].dtype == self.compute_dtype):
                         self._free_fp16_param_shard([p])
+
+            if self._all_gather_free_event_queue is not None:
+                release_params_event = torch.cuda.Event()
+                release_params_event.record()
+                self._all_gather_free_event_queue.enqueue(release_params_event)
+
         if wait_for_all_gather:
             torch.cuda.current_stream().wait_stream(self._streams["all_gather"])
+
         return output_tensors
 
     @torch.no_grad()
diff --git a/fairscale/utils/reduce_scatter_bucketer.py b/fairscale/utils/reduce_scatter_bucketer.py
index a2ef5a6d2..a23ec314b 100644
--- a/fairscale/utils/reduce_scatter_bucketer.py
+++ b/fairscale/utils/reduce_scatter_bucketer.py
@@ -6,11 +6,13 @@
 import functools
 import os
 from typing import Callable, Dict, List, Optional, Tuple
+from logging import getLogger
 
 import torch
 from torch import Tensor
 import torch.distributed as dist
 from torch.distributed import ProcessGroup
+logger = getLogger()
 
 # TODO: Remove the toggle-enable_nccl_base_collectives when github open issue #801 is resolved.
 if os.getenv("ENABLE_NCCL_BASE_COLLECTIVES", "1") == "0":
@@ -97,9 +99,10 @@ class ReduceScatterBucketer:
             are sub-divided based on world_size. Values <= 0 disable bucketing.
     """
 
-    def __init__(self, bucket_cap_mb: int = 25):
+    def __init__(self, bucket_cap_mb: int = 25, _reduce_scatter_free_event_queue = None):
         self.bucket_cap_mb = bucket_cap_mb
         self.buckets: Dict[Tuple[torch.dtype, torch.device, ProcessGroup], Bucket] = {}
+        self._reduce_scatter_free_event_queue = _reduce_scatter_free_event_queue
 
     @torch.no_grad()
     def reduce_scatter_async(
@@ -137,6 +140,13 @@ def reduce_scatter_async(
 
         bucket_shard_size = self._get_shard_size(first_input.element_size(), world_size)
         if first_input_size > bucket_shard_size:
+            if self._reduce_scatter_free_event_queue is not None:
+                while self._reduce_scatter_free_event_queue._queue and self._reduce_scatter_free_event_queue._queue[0].query():
+                    self._reduce_scatter_free_event_queue._dequeue()
+                event = self._reduce_scatter_free_event_queue.dequeue_if_needed()
+                if event:
+                    event.synchronize()
+
             # TODO: investigate how to avoid using torch.cat (because it seems to be slow for CPU tensors)
             # input is too big to fit in the bucket, reduce-scatter directly
             output = torch.zeros_like(input_list[0])

From e9103207cdde05771747cdbab0af784d644a900e Mon Sep 17 00:00:00 2001
From: vedanuj <vedanuj@fb.com>
Date: Thu, 28 Sep 2023 14:28:32 -0700
Subject: [PATCH 05/15] changes for fp8

---
 .../nn/data_parallel/fully_sharded_data_parallel.py    | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 855a2b8bc..4e8e62180 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -444,8 +444,8 @@ def __init__(
         self.numel_padded_per_param: List[int] = []
         self._tstart = time.time()
 
-        if self.fp32_reduce_scatter and not self.mixed_precision:
-            raise ValueError("fp32_reduce_scatter requires mixed_precision=True")
+        # if self.fp32_reduce_scatter and not self.mixed_precision:
+        #     raise ValueError("fp32_reduce_scatter requires mixed_precision=True")
 
         if self.ssd_offload and not self.flatten_parameters:
             raise ValueError(f"offload type: '{offload_config.offload_type}' requires flatten_parameters=True")
@@ -1720,9 +1720,9 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
         with torch.cuda.stream(self._streams["post_backward"]):
             orig_grad_data = param.grad.data
 
-            if self.mixed_precision and self.fp32_reduce_scatter:
+            if self.fp32_reduce_scatter:
                 # Cast grad to FP32.
-                param.grad.data = param.grad.data.to(param.dtype)
+                param.grad.data = param.grad.data.float()
 
             if self.gradient_predivide_factor > 1:
                 # Average grad by world_size for consistency with PyTorch DDP.
@@ -1781,7 +1781,7 @@ def _post_reduction_hook(self, param: Parameter, reduced_grad: torch.Tensor) ->
         # 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:
+        if self.fp32_reduce_scatter:
             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.

From 0db6e62302d2484488cebd847d3b16a01aaafba9 Mon Sep 17 00:00:00 2001
From: Jianyu Huang <hjyahead@gmail.com>
Date: Sun, 1 Oct 2023 19:56:43 -0700
Subject: [PATCH 06/15] Fix fsdp+pp+te WPS decreasing issue (#1139)

* Fix fsdp+pp+te WPS decreasing issue

* Address comment; remove unused stuff

* split into wps fix P841842878 only and main_grad fix
---
 .../nn/data_parallel/fully_sharded_data_parallel.py | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 4e8e62180..759b9f445 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -1650,7 +1650,10 @@ def _register_post_backward_hooks(self) -> None:
                 assert p_tmp.grad_fn is not None
                 grad_acc = p_tmp.grad_fn.next_functions[0][0]  # Gets its GradAccumulation object.
                 handle = grad_acc.register_hook(functools.partial(self._post_backward_hook, p))
-                p._shard_bwd_hook = (grad_acc, handle)
+                if not hasattr(p, "_shard_bwd_hooks"):
+                    p._shard_bwd_hooks = []
+                p._shard_bwd_hooks.append((grad_acc, handle))
+                # p._shard_bwd_hook = (grad_acc, handle)
 
     @torch.no_grad()
     def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
@@ -1860,6 +1863,9 @@ def _finalize_parameters(fsdp_module: FullyShardedDataParallel) -> None:
                     p_assert(len(p._shard_bwd_hook) == 2, f"WFPB: incorrect hook num: {len(p._shard_bwd_hook)}")
                     # p._shard_bwd_hook[1].remove()
                     # delattr(p, "_shard_bwd_hook")
+                if hasattr(p, "_shard_bwd_hooks") and self._require_backward_grad_sync:
+                    for _, handle in p._shard_bwd_hooks:
+                        handle.remove()
 
                 # 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
@@ -1876,7 +1882,10 @@ def _finalize_parameters(fsdp_module: FullyShardedDataParallel) -> None:
                         p.device == p._saved_grad_shard.device,
                         f"WFPB: incorrect saved_grad_shard device {p.device} vs {p._saved_grad_shard.device}",
                     )
-                    p.grad = p._saved_grad_shard
+                    if p._saved_grad_shard.dtype != p.dtype:
+                        p.grad = p._saved_grad_shard.to(p.dtype)
+                    else:
+                        p.grad = p._saved_grad_shard
 
                 if hasattr(p, "_saved_grad_shard"):
                     delattr(p, "_saved_grad_shard")

From 71aeffecc294455b423356646fc50c5c8535828c Mon Sep 17 00:00:00 2001
From: Andrew Gu <31054793+awgu@users.noreply.github.com>
Date: Mon, 2 Oct 2023 22:06:22 -0400
Subject: [PATCH 07/15] Removed extra `cat` before reduce-scatter (#1141)

---
 .../data_parallel/fully_sharded_data_parallel.py |  7 +++++--
 fairscale/utils/reduce_scatter_bucketer.py       | 16 +++++++++-------
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 759b9f445..74947a9ac 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -1751,9 +1751,8 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                 # 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.process_group_reduce_scatter.size())
                 self._reducer.reduce_scatter_async(
-                    grad_chunks, group=self.process_group_reduce_scatter, callback_fn=callback_fn
+                    grad, group=self.process_group_reduce_scatter, callback_fn=callback_fn
                 )
             else:
                 # Currently the only way for _is_sharded to be False is if
@@ -1882,6 +1881,10 @@ def _finalize_parameters(fsdp_module: FullyShardedDataParallel) -> None:
                         p.device == p._saved_grad_shard.device,
                         f"WFPB: incorrect saved_grad_shard device {p.device} vs {p._saved_grad_shard.device}",
                     )
+                    # Reshard in case the parameter was inadvertently gathered
+                    # again after post-backward
+                    if p.shape != p._saved_grad_shard.shape:
+                        self._use_fp32_param_shard([p])
                     if p._saved_grad_shard.dtype != p.dtype:
                         p.grad = p._saved_grad_shard.to(p.dtype)
                     else:
diff --git a/fairscale/utils/reduce_scatter_bucketer.py b/fairscale/utils/reduce_scatter_bucketer.py
index a23ec314b..ea96968ff 100644
--- a/fairscale/utils/reduce_scatter_bucketer.py
+++ b/fairscale/utils/reduce_scatter_bucketer.py
@@ -12,6 +12,7 @@
 from torch import Tensor
 import torch.distributed as dist
 from torch.distributed import ProcessGroup
+from fairscale.utils.parallel import chunk_and_pad
 logger = getLogger()
 
 # TODO: Remove the toggle-enable_nccl_base_collectives when github open issue #801 is resolved.
@@ -107,7 +108,7 @@ def __init__(self, bucket_cap_mb: int = 25, _reduce_scatter_free_event_queue = N
     @torch.no_grad()
     def reduce_scatter_async(
         self,
-        input_list: List[Tensor],
+        grad: Tensor,
         group: ProcessGroup,
         callback_fn: Optional[Callable] = None,
     ) -> None:
@@ -121,16 +122,15 @@ def reduce_scatter_async(
         may also flush the relevant bucket to make room for ``input_list``.
 
         Args:
-            input_list (List[Tensor]): list of tensors to reduce-scatter. List
-                should contain ``group.size()`` tensors and each tensor should
-                have identical shape, dtype and device.
+            grad (Tensor): full gradient to reduce-scatter.
             group (ProcessGroup): process group for reduction
             callback_fn (Callable, Optional): callback function to call after
                 the reduction executes. Function will be called with a single
                 argument corresponding to the reduced result.
         """
         world_size = group.size()
-
+        needs_padding = grad.numel() % world_size != 0
+        input_list = chunk_and_pad(grad, world_size)  # copies last chunk if needs padding
         assert (
             len(input_list) == world_size
         ), f"reduce_scatter received {len(input_list)} inputs, expected group.size() ({world_size})"
@@ -151,8 +151,10 @@ def reduce_scatter_async(
             # input is too big to fit in the bucket, reduce-scatter directly
             output = torch.zeros_like(input_list[0])
             if hasattr(dist, "_reduce_scatter_base") and enable_nccl_base_collectives:
-                input_flattened = torch.cat(input_list)
-                dist._reduce_scatter_base(output, input_flattened, group=group)
+                # For the no-padding + `flatten_parameters=True` case, this
+                # avoids an unnecessary cat
+                input = grad if not needs_padding and grad.ndim == 1 else torch.cat(input_list)
+                dist._reduce_scatter_base(output, input, group=group)
             else:
                 # fallback
                 dist.reduce_scatter(output, input_list, group=group)

From 17ecf4acfd5c9eb62f405ac08b38c147421ec8a8 Mon Sep 17 00:00:00 2001
From: Andrew Gu <31054793+awgu@users.noreply.github.com>
Date: Mon, 9 Oct 2023 21:25:27 -0400
Subject: [PATCH 08/15] Cleared backward hooks to avoid accumulating over
 iterations (#1143)

---
 fairscale/nn/data_parallel/fully_sharded_data_parallel.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 74947a9ac..dee596e9a 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -1865,6 +1865,7 @@ def _finalize_parameters(fsdp_module: FullyShardedDataParallel) -> None:
                 if hasattr(p, "_shard_bwd_hooks") and self._require_backward_grad_sync:
                     for _, handle in p._shard_bwd_hooks:
                         handle.remove()
+                    p._shard_bwd_hooks.clear()
 
                 # 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

From a8189f06ea49fec1d24bc844ec63512209fe23e9 Mon Sep 17 00:00:00 2001
From: Artem <artkorenev@users.noreply.github.com>
Date: Tue, 28 Nov 2023 14:55:29 +0000
Subject: [PATCH 09/15] fix no shard case (#1150)

Co-authored-by: Artem Korenev <artk@fb.com>
---
 fairscale/nn/data_parallel/fully_sharded_data_parallel.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index dee596e9a..f5c024d5b 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -1759,7 +1759,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                 # world_size == 1. This could be relaxed in the future, in which
                 # case grads should be all-reduced here.
                 assert self.world_size == 1
-                self._post_reduction_hook(param, param.grad.data)
+                self._post_reduction_hook(param, param.grad)
 
             # After _post_backward_hook returns, orig_grad_data will eventually
             # go out of scope, at which point it could otherwise be freed for

From 3b7cc249ad48dbd703ea38afb1b7b662afa88f8c Mon Sep 17 00:00:00 2001
From: ngoyal2707 <ngoyal2707@users.noreply.github.com>
Date: Wed, 6 Dec 2023 20:28:27 -0500
Subject: [PATCH 10/15] [not to be merged yet] added temp changes for fp32 main
 grad, might not work for TE (#1151)

* added temp changes for fp32 main grad, might not work for TE

* post rebase

* changes to keep reduced grad in fp32 (#1152)

* fix .grad=None issue when param is not sharded (#1153)

* fixed broken clipping (#1154)

Co-authored-by: Naman Goyal <naman@fb.com>

---------

Co-authored-by: Naman Goyal <naman@fb.com>
Co-authored-by: Vedanuj Goswami <vedanuj@users.noreply.github.com>
Co-authored-by: Jiecao Yu <jiecaoyu@umich.edu>
---
 .../fully_sharded_data_parallel.py            | 40 ++++++++++++++-----
 1 file changed, 31 insertions(+), 9 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index f5c024d5b..09648894c 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -687,7 +687,7 @@ def _cast_buffers(
     @property
     def params_with_grad(self) -> List[Parameter]:
         """[p for p in self.parameters() if p.grad is not None]"""
-        return [p for p in self.parameters() if p.grad is not None]
+        return [p for p in self.parameters() if (p.grad is not None or p.main_grad is not None)]
 
     @torch.no_grad()
     def clip_grad_norm_(
@@ -1714,6 +1714,14 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
         # Switch to FP32 shard after backward.
         self._use_fp32_param_shard([param])
 
+        if self.fp32_reduce_scatter:
+            if getattr(param, "unsharded_main_grad", None) is None:
+                param.unsharded_main_grad = param.grad.to(torch.float32)
+            else:
+                param.unsharded_main_grad.add_(param.grad.data)
+
+            param.grad = None
+
         if not self._require_backward_grad_sync:
             return
 
@@ -1721,15 +1729,19 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
         # reductions in post_backward stream.
         self._streams["post_backward"].wait_stream(torch.cuda.current_stream())
         with torch.cuda.stream(self._streams["post_backward"]):
-            orig_grad_data = param.grad.data
 
             if self.fp32_reduce_scatter:
                 # Cast grad to FP32.
-                param.grad.data = param.grad.data.float()
+                orig_grad_data = param.unsharded_main_grad.data
+            else:
+                orig_grad_data = param.grad.data
 
             if self.gradient_predivide_factor > 1:
                 # Average grad by world_size for consistency with PyTorch DDP.
-                param.grad.data.div_(self.gradient_predivide_factor)
+                if getattr(param, "unsharded_main_grad", None) is not None:
+                    param.unsharded_main_grad.data.div_(self.gradient_predivide_factor)
+                else:
+                    param.grad.data.div_(self.gradient_predivide_factor)
 
             if param._is_sharded:
                 assert self._reducer is not None
@@ -1737,7 +1749,13 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                 # 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
+                if getattr(param, "unsharded_main_grad", None) is not None:
+                    grad = param.unsharded_main_grad.data
+                    param.unsharded_main_grad = None
+                else:
+                    grad = param.grad.data
+                    param.grad = None
+
                 # 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
@@ -1749,7 +1767,6 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                 # 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)
                 self._reducer.reduce_scatter_async(
                     grad, group=self.process_group_reduce_scatter, callback_fn=callback_fn
@@ -1759,7 +1776,10 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                 # world_size == 1. This could be relaxed in the future, in which
                 # case grads should be all-reduced here.
                 assert self.world_size == 1
-                self._post_reduction_hook(param, param.grad)
+                if getattr(param, "unsharded_main_grad", None) is not None:
+                    self._post_reduction_hook(param, param.unsharded_main_grad)
+                else:
+                    self._post_reduction_hook(param, param.grad)
 
             # After _post_backward_hook returns, orig_grad_data will eventually
             # go out of scope, at which point it could otherwise be freed for
@@ -1785,7 +1805,7 @@ def _post_reduction_hook(self, param: Parameter, reduced_grad: torch.Tensor) ->
         # non-blocking. The downside is a bit more D2H transfer in that case.
         if self.fp32_reduce_scatter:
             orig_param_grad_data = reduced_grad.data
-            reduced_grad.data = reduced_grad.data.to(dtype=param.data.dtype)
+            # 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())
 
@@ -1799,6 +1819,8 @@ def _post_reduction_hook(self, param: Parameter, reduced_grad: torch.Tensor) ->
                 ), 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
+        elif (param.grad is None) and self.fp32_reduce_scatter:
+            param.main_grad = reduced_grad.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.
@@ -1887,7 +1909,7 @@ def _finalize_parameters(fsdp_module: FullyShardedDataParallel) -> None:
                     if p.shape != p._saved_grad_shard.shape:
                         self._use_fp32_param_shard([p])
                     if p._saved_grad_shard.dtype != p.dtype:
-                        p.grad = p._saved_grad_shard.to(p.dtype)
+                        p.main_grad = p._saved_grad_shard
                     else:
                         p.grad = p._saved_grad_shard
 

From a4f02efd50a92f57cb72986a906d249568b8fe26 Mon Sep 17 00:00:00 2001
From: Andrew Gu <andgu@fb.com>
Date: Fri, 12 Jan 2024 11:19:17 -0800
Subject: [PATCH 11/15] Added reshard hook for frozen params in backward

---
 .../fully_sharded_data_parallel.py            | 74 +++++++++++---
 .../test_fsdp_freezing_weights.py             | 96 +++++++++++++++++++
 2 files changed, 159 insertions(+), 11 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 09648894c..d9b20fca7 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -9,6 +9,7 @@
 from dataclasses import dataclass
 from enum import Enum, auto
 import functools
+import itertools
 import logging
 from math import inf
 import os
@@ -47,7 +48,6 @@
 from fairscale.utils.containers import apply_to_tensors
 from fairscale.utils.parallel import (
     ProcessGroupName,
-    chunk_and_pad,
     enable_pytorch_sync_bn,
     get_process_group_cached,
     validate_process_group,
@@ -1457,6 +1457,7 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
         # Register backward hooks to reshard params and reduce-scatter grads.
         # These need to be re-registered every forward pass.
         self._register_post_backward_hooks()
+        self._register_post_backward_reshard_hooks(args, kwargs)
 
         outputs = self.module(*args, **kwargs)
 
@@ -1655,6 +1656,37 @@ def _register_post_backward_hooks(self) -> None:
                 p._shard_bwd_hooks.append((grad_acc, handle))
                 # p._shard_bwd_hook = (grad_acc, handle)
 
+    def _register_post_backward_reshard_hooks(
+        self, args: Tuple[Any, ...], kwargs: Dict[str, Any]
+    ) -> None:
+        if not hasattr(torch.autograd.graph, "register_multi_grad_hook"):
+            return  # unsupported
+        if not torch.is_grad_enabled():
+            return
+        from torch.utils._pytree import tree_flatten
+        from torch.autograd.graph import register_multi_grad_hook
+        # Construct `inp_tensors` lazily to avoid CPU overhead in typical case
+        # where each parameter requires gradient
+        inp_tensors: Optional[List[torch.Tensor]] = None
+        for param in self.params:
+            # Only register for parameters that do not require gradient
+            if param.requires_grad:
+                continue
+            if inp_tensors is None:
+                args_list, _ = tree_flatten(args)
+                kwargs_list, _ = tree_flatten(kwargs)
+                inp_tensors = [
+                    obj
+                    for obj in itertools.chain(args_list, kwargs_list)
+                    if torch.is_tensor(obj) and obj.requires_grad
+                ]
+            hook_handle = register_multi_grad_hook(
+                inp_tensors, functools.partial(self._post_backward_reshard_hook, param)
+            )
+            if not hasattr(param, "_shard_bwd_hooks"):
+                param._shard_bwd_hooks = []
+            param._shard_bwd_hooks.append((hook_handle,))
+
     @torch.no_grad()
     def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
         """
@@ -1697,12 +1729,8 @@ 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 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
-            # ``self._require_backward_grad_sync``), since the params will not
-            # get updated before the next forward. This saves networking
-            # bandwidth but uses more GPU memory.
+        if self._should_free_in_backward():
+            # Free full params.
             self._free_full_params([param])
 
         if self.mixed_precision:
@@ -1829,6 +1857,22 @@ def _post_reduction_hook(self, param: Parameter, reduced_grad: torch.Tensor) ->
             # Don't let this memory get reused until after the transfer.
             reduced_grad.data.record_stream(torch.cuda.current_stream())
 
+    @torch.no_grad()
+    def _post_backward_reshard_hook(self, param: Parameter, *unused: Any) -> None:
+        if self._should_free_in_backward():
+            self._free_full_params([param])
+        if self.mixed_precision:
+            self._free_fp16_param_shard([param])
+        self._use_fp32_param_shard([param])
+
+    def _should_free_in_backward(self):
+        # As a special case, we don't free the full params
+        # when in a ``no_sync`` context (as inversely indicated by
+        # ``self._require_backward_grad_sync``), since the params will not
+        # get updated before the next forward. This saves networking
+        # bandwidth but uses more GPU memory.
+        return self._require_backward_grad_sync or self.reshard_after_forward
+
     def _queue_wait_for_post_backward(self) -> None:
         """Try to queue a `wait_for_post_backward` callback.
 
@@ -1878,16 +1922,24 @@ def _wait_for_post_backward(self) -> None:
         def _finalize_parameters(fsdp_module: FullyShardedDataParallel) -> None:
             """Helper used below on all fsdp modules."""
             for p in fsdp_module.params:
-                if not p.requires_grad:
-                    continue
                 if hasattr(p, "_shard_bwd_hook"):
                     p_assert(len(p._shard_bwd_hook) == 2, f"WFPB: incorrect hook num: {len(p._shard_bwd_hook)}")
                     # p._shard_bwd_hook[1].remove()
                     # delattr(p, "_shard_bwd_hook")
                 if hasattr(p, "_shard_bwd_hooks") and self._require_backward_grad_sync:
-                    for _, handle in p._shard_bwd_hooks:
-                        handle.remove()
+                    for hook_state in p._shard_bwd_hooks:
+                        if len(hook_state) == 1:
+                            hook_state[0].remove()
+                        elif len(hook_state) == 2:
+                            hook_state[1].remove()
                     p._shard_bwd_hooks.clear()
+                if not p.requires_grad:
+                    # For the 1st layer, if the forward inputs did not require
+                    # gradient, then we cannot run a reshard hook for it, and
+                    # we instead free here.
+                    if p._full_param_padded.untyped_storage().size() > 0:
+                        fsdp_module._post_backward_reshard_hook(p)
+                    continue
 
                 # 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
diff --git a/tests/nn/data_parallel/test_fsdp_freezing_weights.py b/tests/nn/data_parallel/test_fsdp_freezing_weights.py
index c6ad364f7..7baadc5d9 100644
--- a/tests/nn/data_parallel/test_fsdp_freezing_weights.py
+++ b/tests/nn/data_parallel/test_fsdp_freezing_weights.py
@@ -12,6 +12,8 @@
 
 from enum import Enum
 from itertools import product
+from unittest import mock
+import copy
 import tempfile
 
 import pytest
@@ -275,3 +277,97 @@ def test_freezing_weights(temp_files, nested_trunk):
             nprocs=world_size,
         )
         temp_file_idx += 3
+
+
+@skip_if_single_gpu
+def test_reshard_frozen_weights():
+    world_size = 2
+    for flatten_parameters, reshard_after_forward, inp_requires_grad in product(
+        [False, True], [False, True], [False, True]
+    ):
+        print(
+            "Testing FSDP reshard frozen weights with "
+            f"flatten_parameters={flatten_parameters}, "
+            f"reshard_after_forward={reshard_after_forward}, "
+            f"inp_requires_grad={inp_requires_grad}"
+        )
+        mp.spawn(
+            _distributed_worker_reshard,
+            (world_size, flatten_parameters, reshard_after_forward, inp_requires_grad),
+            nprocs=world_size,
+        )
+
+
+def _distributed_worker_reshard(
+    rank: int,
+    world_size: int,
+    flatten_parameters: bool,
+    reshard_after_forward: bool,
+    inp_requires_grad: bool,
+):
+    import os
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = "12355"
+    torch.cuda.set_device(rank)
+    torch.distributed.init_process_group(backend="nccl", rank=rank, world_size=world_size)
+
+    torch.manual_seed(0)
+
+    num_linears = 6
+    modules = []
+    for _ in range(num_linears):
+        modules += [nn.Linear(5, 5, device="cuda"), nn.ReLU()]
+    model = nn.Sequential(*modules)
+    # Freeze every other linear
+    for i in range(num_linears):
+        if i % 2 == 0:
+            for param in model[i * 2].parameters(recurse=False):
+                param.requires_grad = False
+    num_frozen_linears = num_linears // 2
+
+    ref_model = DistributedDataParallel(copy.deepcopy(model), device_ids=[rank])
+    ref_optim = torch.optim.AdamW(ref_model.parameters(), lr=1e-2)
+
+    for i, module in enumerate(model):
+        if isinstance(module, nn.Linear):
+            model[i] = FSDP(
+                module,
+                flatten_parameters=flatten_parameters,
+                reshard_after_forward=reshard_after_forward,
+            )
+    fsdp_model = FSDP(
+        model,
+        flatten_parameters=flatten_parameters,
+        reshard_after_forward=reshard_after_forward,
+    )
+    fsdp_optim = torch.optim.AdamW(fsdp_model.parameters(), lr=1e-2)
+
+    orig_post_backward_reshard_hook = FSDP._post_backward_reshard_hook
+    reshard_hook_count = 0
+
+    def post_backward_reshard_hook_with_count(*args, **kwargs):
+        nonlocal reshard_hook_count
+        reshard_hook_count += 1
+        return orig_post_backward_reshard_hook(*args, **kwargs)
+
+    with mock.patch(
+        "fairscale.nn.data_parallel.FullyShardedDataParallel._post_backward_reshard_hook",
+        post_backward_reshard_hook_with_count,
+    ):
+        inp = torch.randn((8, 5), device="cuda", requires_grad=inp_requires_grad)
+        for i in range(6):
+            losses = []
+            for model, optim in ((fsdp_model, fsdp_optim), (ref_model, ref_optim)):
+                optim.zero_grad()
+                loss = model(inp).sum()
+                losses.append(loss)
+                loss.backward()
+                optim.step()
+            expected_reshard_hook_count = num_frozen_linears
+            if not flatten_parameters:
+                expected_reshard_hook_count *= 2  # weight and bias per linear
+            assert (
+                reshard_hook_count == expected_reshard_hook_count
+            ), f"Expected {expected_reshard_hook_count} but got {reshard_hook_count}"
+            assert losses[0].eq(losses[1]).all().item(), f"Expected {losses[1]} but got {losses[0]}"
+            reshard_hook_count = 0

From f2bb56f930c6f5fbb4d6cb0c3e8e09225b04f122 Mon Sep 17 00:00:00 2001
From: Jiecao Yu <jiecaoyu@umich.edu>
Date: Wed, 21 Feb 2024 03:38:59 -0800
Subject: [PATCH 12/15] Avoid calling _free_fp16_param_shard() too early with
 PR 1159

---
 fairscale/nn/data_parallel/fully_sharded_data_parallel.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index d9b20fca7..7d7fce9a4 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -1733,7 +1733,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
             # Free full params.
             self._free_full_params([param])
 
-        if self.mixed_precision:
+        if self.mixed_precision and (self._require_backward_grad_sync or self.reshard_after_forward):
             # This is a no-op if reshard_after_forward is True, since we already
             # free the param shard when rebuilding the full params in the
             # pre_backward_hook.
@@ -1861,7 +1861,7 @@ def _post_reduction_hook(self, param: Parameter, reduced_grad: torch.Tensor) ->
     def _post_backward_reshard_hook(self, param: Parameter, *unused: Any) -> None:
         if self._should_free_in_backward():
             self._free_full_params([param])
-        if self.mixed_precision:
+        if self.mixed_precision and (self._require_backward_grad_sync or self.reshard_after_forward):
             self._free_fp16_param_shard([param])
         self._use_fp32_param_shard([param])
 
@@ -1937,7 +1937,7 @@ def _finalize_parameters(fsdp_module: FullyShardedDataParallel) -> None:
                     # For the 1st layer, if the forward inputs did not require
                     # gradient, then we cannot run a reshard hook for it, and
                     # we instead free here.
-                    if p._full_param_padded.untyped_storage().size() > 0:
+                    if p._is_sharded and p._full_param_padded.untyped_storage().size() > 0:
                         fsdp_module._post_backward_reshard_hook(p)
                     continue
 

From 1307b1dc4b53ad0f4f23e64e25a5f1db6ee2e559 Mon Sep 17 00:00:00 2001
From: Jie Wang <holawj@gmail.com>
Date: Mon, 25 Mar 2024 11:56:10 -0700
Subject: [PATCH 13/15] Added requires_grad check for params_with_grad method
 (#1171)

Co-authored-by: Jie Wang <jiewang@meta.com>
---
 fairscale/nn/data_parallel/fully_sharded_data_parallel.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 7d7fce9a4..437227e7a 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -687,7 +687,7 @@ def _cast_buffers(
     @property
     def params_with_grad(self) -> List[Parameter]:
         """[p for p in self.parameters() if p.grad is not None]"""
-        return [p for p in self.parameters() if (p.grad is not None or p.main_grad is not None)]
+        return [p for p in self.parameters() if (p.requires_grad and (p.grad is not None or p.main_grad is not None))]
 
     @torch.no_grad()
     def clip_grad_norm_(

From 5faca97f32793cc08a4f24d9ada0243bb3f2360d Mon Sep 17 00:00:00 2001
From: Andrew Gu <31054793+awgu@users.noreply.github.com>
Date: Mon, 1 Apr 2024 14:08:57 -0400
Subject: [PATCH 14/15] Changed to only run reshard hook if all gradients
 computed (#1166)

* Changed to only run reshard hook if all gradients computed

* Fix decreasing it/s with multi-grad hook
---
 .../fully_sharded_data_parallel.py            | 76 ++++++++++++++++++-
 1 file changed, 73 insertions(+), 3 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 437227e7a..71b07af35 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -28,6 +28,7 @@
     Mapping,
     NamedTuple,
     Optional,
+    Sequence,
     Set,
     Tuple,
     Union,
@@ -42,6 +43,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from torch.nn.parameter import Parameter
+from torch.utils.hooks import RemovableHandle
 
 from fairscale.nn.misc import FlattenParamsWrapper
 from fairscale.nn.wrap import auto_wrap, config_auto_wrap_policy, enable_wrap
@@ -1659,12 +1661,9 @@ def _register_post_backward_hooks(self) -> None:
     def _register_post_backward_reshard_hooks(
         self, args: Tuple[Any, ...], kwargs: Dict[str, Any]
     ) -> None:
-        if not hasattr(torch.autograd.graph, "register_multi_grad_hook"):
-            return  # unsupported
         if not torch.is_grad_enabled():
             return
         from torch.utils._pytree import tree_flatten
-        from torch.autograd.graph import register_multi_grad_hook
         # Construct `inp_tensors` lazily to avoid CPU overhead in typical case
         # where each parameter requires gradient
         inp_tensors: Optional[List[torch.Tensor]] = None
@@ -2823,3 +2822,74 @@ def auto_wrap_bn(
         enable_wrap(config_auto_wrap_policy, wrapper_cls=FullyShardedDataParallel) if wrap_it else contextlib.suppress()
     ):
         return auto_wrap(module)
+
+
+class Handle(RemovableHandle):
+    handles: Tuple[RemovableHandle, ...]
+
+    def __init__(self, handles: Tuple[RemovableHandle, ...]):
+        self.handles = handles
+
+    def remove(self):
+        for handle in self.handles:
+            handle.remove()
+
+    def __getstate__(self):
+        return self.handles
+
+    def __setstate__(self, state):
+        self.handles = state
+
+
+def register_multi_grad_hook(
+    tensors: Sequence[torch.Tensor],
+    fn: Callable[[Sequence[Optional[torch.Tensor]]], None]
+):
+    count: Dict[int, int] = dict()
+    nb_calls = None
+    buffer: Dict[int, List[Optional[torch.Tensor]]] = dict()
+
+    grad_fns = list(map(_get_grad_fn_or_grad_acc, tensors))
+    len_tensors = len(tensors)
+
+    def get_inner_hook(idx):
+        def inner_hook(grad: torch.Tensor):
+            nonlocal count, nb_calls, buffer, fn
+            id = torch._C._current_graph_task_id()
+            assert (
+                id != -1
+            ), "expected this hook to be called inside a backward call"
+            count[id] = count.get(id, 0)
+            buffer[id] = buffer.get(id, [None] * len_tensors)
+
+            if count[id] == 0:
+                # On the first call, compute the actual nb_calls and buffer
+                # nb_calls = sum(torch._C._will_engine_execute_node(g) for g in grad_fns)  # type: ignore[attr-defined]
+
+                # NOTE: To avoid resharding too early when microbatches share
+                # some same module inputs, let us require all gradients to be
+                # computed in this backward for the hook to run.
+                nb_calls = len(grad_fns)
+
+            buffer[id][idx] = grad
+            count[id] += 1
+
+            if count[id] == nb_calls:
+                fn = cast(Callable[[Sequence[Optional[torch.Tensor]]], None], fn)
+                fn(buffer[id])
+                del count[id]
+                del buffer[id]
+
+        return inner_hook
+
+    handles: Tuple[RemovableHandle, ...] = tuple(
+        t.register_hook(get_inner_hook(i)) for i, t in enumerate(tensors)
+    )
+    return Handle(handles)
+
+
+def _get_grad_fn_or_grad_acc(t):
+    if t.requires_grad and t.grad_fn is None:
+        return t.view_as(t).grad_fn.next_functions[0][0]
+    else:
+        return t.grad_fn

From 7bcbc805777ffd1696934d282a11821c02dd997a Mon Sep 17 00:00:00 2001
From: Jie Wang <holawj@gmail.com>
Date: Fri, 5 Apr 2024 12:45:35 -0700
Subject: [PATCH 15/15] Add cast input argument (#1175)

Co-authored-by: Jie Wang <jiewang@meta.com>
---
 fairscale/nn/data_parallel/fully_sharded_data_parallel.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 71b07af35..bfa02ae77 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -370,6 +370,7 @@ def __init__(
         gradient_predivide_factor: Optional[float] = None,
         limit_all_gather_events: bool = False,
         limit_reduce_scatter_events: bool = False,
+        cast_input: bool = True,
     ):
         try:
             import torch._C
@@ -420,6 +421,7 @@ def __init__(
         self.reshard_after_forward = self._orig_reshard_after_forward = reshard_after_forward
         self.disable_reshard_on_root = disable_reshard_on_root
         self.mixed_precision = mixed_precision
+        self.cast_input = cast_input
         self.fp32_reduce_scatter = fp32_reduce_scatter
         self.flatten_parameters = flatten_parameters
         self.move_params_to_cpu = move_params_to_cpu or cpu_offload
@@ -1431,7 +1433,7 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
         # For root and mixed precision, we convert the input to FP16 (no_grad is needed for
         # the conversion).
         is_bf16 = self.compute_dtype == torch.bfloat16
-        if self._is_root and self.mixed_precision:
+        if self._is_root and self.mixed_precision and self.cast_input:
             args, kwargs = cast_floats_to_right_precision(True, True, is_bf16, *args, **kwargs)
 
         if self not in self._fsdp_forward_ordering: