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/11] 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/11] 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 1efe26a2370a99620a6291c21d4fce09a86c3cb4 Mon Sep 17 00:00:00 2001
From: Ruan Silva <ruanrms@devfair0830.h2.fair>
Date: Thu, 28 Jul 2022 09:32:56 -0700
Subject: [PATCH 03/11] format

---
 .../fully_sharded_data_parallel.py            | 24 ++++++++++---------
 1 file changed, 13 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 f454d67f4..12e316e4f 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -310,7 +310,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.default,
+        process_group_reduce_scatter: Any = ProcessGroupName.reduce_scatter,
         reshard_after_forward: bool = True,
         disable_reshard_on_root: bool = True,
         mixed_precision: bool = False,
@@ -1162,7 +1162,7 @@ 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._fsdp_forward_ordering: List[FullyShardedDataParallel] = []
         self._my_fsdp_instance_idx: Optional[int] = None
         for p in self.params:
             if hasattr(p, "_fp32_shard"):
@@ -1412,11 +1412,10 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
 
         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
+            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
-            )
+            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.
@@ -1509,9 +1508,12 @@ def _pre_backward_hook(*unused: Any) -> None:
                 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
+                    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)
+                    self._fsdp_forward_ordering[self._my_fsdp_instance_idx - 1]._rebuild_full_params(
+                        wait_for_all_gather=False
+                    )
             else:
                 self._use_full_params()
 
@@ -1875,7 +1877,9 @@ 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, wait_for_all_gather = True) -> Optional[List[Tuple[torch.Tensor, bool]]]:
+    def _rebuild_full_params(
+        self, force_full_precision: bool = False, wait_for_all_gather: bool = True
+    ) -> Optional[List[Tuple[torch.Tensor, bool]]]:
         """
         Gather all shards of params.
 
@@ -1946,8 +1950,6 @@ 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

From 418d5db70dbf4d906fbb65dc9ac11766825d6bad Mon Sep 17 00:00:00 2001
From: Ruan Silva <ruanrms@devfair0830.h2.fair>
Date: Fri, 29 Jul 2022 06:51:49 -0700
Subject: [PATCH 04/11] add check in backward

---
 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 12e316e4f..577a5e9e3 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -1511,9 +1511,9 @@ def _pre_backward_hook(*unused: Any) -> 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
-                    )
+                    t = self._fsdp_forward_ordering[self._my_fsdp_instance_idx - 1]
+                    if id(t) in self._output_pre_backward_hook_registered:
+                        t._rebuild_full_params(wait_for_all_gather=False)
             else:
                 self._use_full_params()
 

From b259a21d19ca4e4e9bd490b4f180d26be654a0f9 Mon Sep 17 00:00:00 2001
From: Ruan Silva <ruanrms@devfair0830.h2.fair>
Date: Fri, 29 Jul 2022 15:59:54 -0700
Subject: [PATCH 05/11] fix tests

---
 .../data_parallel/fully_sharded_data_parallel.py   | 14 +++++++++++++-
 .../test_fsdp_multiple_forward_checkpoint.py       |  4 ++--
 2 files changed, 15 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 577a5e9e3..96c95ae72 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -1388,6 +1388,9 @@ def _wait_for_previous_optim_step(self) -> None:
     def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
         self._lazy_init()
 
+        # if (self.rank == 0):
+        #     print("forward size = " + str(len(self._fsdp_forward_ordering)) + " idx = " + str(self._my_fsdp_instance_idx))
+
         # Start of a forward pass.
         self.training_state = TrainingState.FORWARD
 
@@ -1415,7 +1418,11 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
             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)
+            t = self._fsdp_forward_ordering[self._my_fsdp_instance_idx + 1]
+            if not t._pre_backward_hook_has_run:
+                # if (self.rank == 0):
+                #     print("fetching idx = " + str(self._my_fsdp_instance_idx + 1))
+                t._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.
@@ -1486,6 +1493,9 @@ def _pre_backward_hook(*unused: Any) -> None:
             # that final backward callback is attached to the outer most
             # backward graph task and called after all the backward
             # calls are completed.
+            # if (self.rank == 0):
+            #     print("pre-backward size = " + str(len(self._fsdp_forward_ordering)) + " idx = " + str(self._my_fsdp_instance_idx))
+
             if self._is_root:
                 self._queue_wait_for_post_backward()
 
@@ -1950,6 +1960,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
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 f7575430d..046545712 100644
--- a/tests/nn/data_parallel/test_fsdp_multiple_forward_checkpoint.py
+++ b/tests/nn/data_parallel/test_fsdp_multiple_forward_checkpoint.py
@@ -344,10 +344,10 @@ def test_multiple_forward_checkpoint(precision, flatten, wrap_bn, model_type, bn
     # computation interact correctly.
     combinations = []
     for with_fsdp in [False, True]:
-        for with_checkpoint in [False, True]:
+        for with_checkpoint in [False]:
             if not with_fsdp and with_checkpoint:
                 continue
-            for with_bucketing in [False, True]:
+            for with_bucketing in [False]:
                 if not with_fsdp and with_bucketing:
                     continue
                 combinations.append((with_fsdp, with_checkpoint, with_bucketing))

From 4ef39dcfcfaeb934330bf04e52b7c8a530b363fc Mon Sep 17 00:00:00 2001
From: Ruan Silva <ruanrms@devfair0830.h2.fair>
Date: Mon, 1 Aug 2022 15:36:02 -0700
Subject: [PATCH 06/11] refactor to single method

---
 .../fully_sharded_data_parallel.py            | 117 +++++++++++++-----
 .../test_fsdp_multiple_forward_checkpoint.py  |   4 +-
 .../data_parallel/test_fsdp_shared_weights.py |   4 +-
 .../test_fsdp_shared_weights_mevo.py          |   2 +-
 4 files changed, 89 insertions(+), 38 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 96c95ae72..f6f8baa7f 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -549,7 +549,7 @@ def module(self) -> FlattenParamsWrapper:
         assert isinstance(self._fsdp_wrapped_module, FlattenParamsWrapper)
         return self._fsdp_wrapped_module
 
-    def append_shared_param(self, p: Parameter) -> None:
+    def append_shared_param(self, p: Parameter, original_module: "FullyShardedDataParallel") -> None:
         """Add a param that's already owned by another FSDP wrapper.
 
             .. warning:: This is experimental!
@@ -576,6 +576,7 @@ def append_shared_param(self, p: Parameter) -> None:
         ), "Must have at least 1 non-shared param."
         self.params.append(p)
         self._has_shared_params = True
+        original_module._has_shared_params = True
 
     def non_shared_params(self) -> List[nn.Parameter]:
         """Return the list of non-shared parameters."""
@@ -1162,8 +1163,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[FullyShardedDataParallel] = []
-        self._my_fsdp_instance_idx: Optional[int] = None
+        self._forward_ordering: List[FullyShardedDataParallel] = []
+        self._backward_ordering: List[FullyShardedDataParallel] = []
         for p in self.params:
             if hasattr(p, "_fp32_shard"):
                 del p._fp32_shard  # reset _init_param_attributes
@@ -1336,7 +1337,8 @@ 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
+                m._forward_ordering = self._forward_ordering
+                m._backward_ordering = self._backward_ordering
 
     def _setup_streams(self) -> None:
         """Create streams to overlap data transfer and computation."""
@@ -1388,9 +1390,6 @@ def _wait_for_previous_optim_step(self) -> None:
     def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
         self._lazy_init()
 
-        # if (self.rank == 0):
-        #     print("forward size = " + str(len(self._fsdp_forward_ordering)) + " idx = " + str(self._my_fsdp_instance_idx))
-
         # Start of a forward pass.
         self.training_state = TrainingState.FORWARD
 
@@ -1399,30 +1398,39 @@ 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.
         if self.force_input_to_fp32 and not self.mixed_precision:
             args, kwargs = cast_floats_to_right_precision(False, False, *args, **kwargs)
 
+        if self not in self._forward_ordering:
+            self._forward_ordering.append(self)
+        forward_idx = self._forward_ordering.index(self)
+        if self.rank == 0:
+            print(
+                "forward size = "
+                + str(len(self._forward_ordering))
+                + " idx = "
+                + str(forward_idx)
+                + " has shared = "
+                + str(self._has_shared_params)
+                + " has full params = "
+                + str(self.has_full_params)
+            )
+
         # All-gather full parameters. This will also transfer FP32 parameters to
         # ``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
-        ):
-            t = self._fsdp_forward_ordering[self._my_fsdp_instance_idx + 1]
-            if not t._pre_backward_hook_has_run:
-                # if (self.rank == 0):
-                #     print("fetching idx = " + str(self._my_fsdp_instance_idx + 1))
-                t._rebuild_full_params(wait_for_all_gather=False)
+        # if (
+        #     forward_idx < len(self._forward_ordering) - 1
+        # ):
+        #     t = self._forward_ordering[forward_idx + 1]
+        #     if not t._pre_backward_hook_has_run:
+        #         if (self.rank == 0):
+        #             print("fetching idx = " + str(forward_idx + 1))
+        #         t._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.
@@ -1430,6 +1438,8 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
 
         outputs = self.module(*args, **kwargs)
 
+        # if(self.rank == 0):
+        #     print(str(self.reshard_after_forward) + " forward_idx = " + str(self._forward_ordering.index(self)))
         if self.reshard_after_forward:
             self._free_full_params()
             if self.mixed_precision or self.move_params_to_cpu:
@@ -1493,8 +1503,6 @@ def _pre_backward_hook(*unused: Any) -> None:
             # that final backward callback is attached to the outer most
             # backward graph task and called after all the backward
             # calls are completed.
-            # if (self.rank == 0):
-            #     print("pre-backward size = " + str(len(self._fsdp_forward_ordering)) + " idx = " + str(self._my_fsdp_instance_idx))
 
             if self._is_root:
                 self._queue_wait_for_post_backward()
@@ -1514,16 +1522,36 @@ def _pre_backward_hook(*unused: Any) -> None:
             # idempotent.  So in case they are called unnecessarily, they don't incur much
             # overhead.
             if self.reshard_after_forward:
+                if self not in self._backward_ordering:
+                    self._backward_ordering.append(self)
+                backward_idx = self._backward_ordering.index(self)
+                if self.rank == 0:
+                    print(
+                        "pre-backward size = "
+                        + str(len(self._backward_ordering))
+                        + " idx = "
+                        + str(backward_idx)
+                        + " forward_idx = "
+                        + str(self._forward_ordering.index(self))
+                        + "has shared = "
+                        + str(self._has_shared_params)
+                        + " has full params = "
+                        + str(self.has_full_params)
+                    )
+
                 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
-                ):
-                    t = self._fsdp_forward_ordering[self._my_fsdp_instance_idx - 1]
-                    if id(t) in self._output_pre_backward_hook_registered:
-                        t._rebuild_full_params(wait_for_all_gather=False)
+
+                # if (
+                #     backward_idx < len(self._backward_ordering) - 1
+                # ):
+                #     t = self._backward_ordering[backward_idx + 1]
+                #     # is_hook_registered = id(t) in self._output_pre_backward_hook_registered
+                #     # # if (self.rank == 0):
+                #     # #     print("is_hook_registered = " + str(is_hook_registered) + " id = " + str(id(t)) + " len = " + str(len(self._output_pre_backward_hook_registered)))
+                #     # if not is_hook_registered:
+                #     if (self.rank == 0):
+                #         print("fetching backward idx = " + str(backward_idx + 1) + " forward_idx = " + str(self._forward_ordering.index(t)))
+                #     t._rebuild_full_params(wait_for_all_gather=False)
             else:
                 self._use_full_params()
 
@@ -1651,6 +1679,8 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
         """
         # First hook callback will see PRE state. If we have multiple params,
         # then subsequent hook callbacks will see POST state.
+        # if (self.rank == 0):
+        #     print("post backward forward_idx = " + str(self._forward_ordering.index(self)))
         self.assert_state([TrainingState.BACKWARD_PRE, TrainingState.BACKWARD_POST])
         self.training_state = TrainingState.BACKWARD_POST
         if param.grad is None:
@@ -2093,7 +2123,28 @@ 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()
+        self._schedule_next_all_gather_and_synchronize()
+        # if(self.rank == 0):
+        #     print(str(self.training_state) + " forward_idx = " + str(self._forward_ordering.index(self)))
+
+    @torch.no_grad()
+    def _schedule_next_all_gather_and_synchronize(self) -> None:
+        self.assert_state([TrainingState.FORWARD, TrainingState.BACKWARD_POST])
+        ordering = self._forward_ordering
+        if self.training_state == TrainingState.BACKWARD_POST:
+            ordering = self._backward_ordering
+        next_idx = len(ordering)
+        if self in ordering:
+            next_idx = ordering.index(self) + 1
+        while next_idx < len(ordering):
+            next_module = ordering[next_idx]
+            if not next_module._pre_backward_hook_has_run and not next_module.has_full_params:
+                if self.rank == 0:
+                    print(str(self.training_state) + " fetching idx = " + str(next_idx))
+                next_module._rebuild_full_params(wait_for_all_gather=False)
+                break
+            next_idx = next_idx + 1
+        torch.cuda.current_stream().synchronize()
 
     def local_metadata_dict(self) -> Dict[str, Any]:
         """
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 046545712..f7575430d 100644
--- a/tests/nn/data_parallel/test_fsdp_multiple_forward_checkpoint.py
+++ b/tests/nn/data_parallel/test_fsdp_multiple_forward_checkpoint.py
@@ -344,10 +344,10 @@ def test_multiple_forward_checkpoint(precision, flatten, wrap_bn, model_type, bn
     # computation interact correctly.
     combinations = []
     for with_fsdp in [False, True]:
-        for with_checkpoint in [False]:
+        for with_checkpoint in [False, True]:
             if not with_fsdp and with_checkpoint:
                 continue
-            for with_bucketing in [False]:
+            for with_bucketing in [False, True]:
                 if not with_fsdp and with_bucketing:
                     continue
                 combinations.append((with_fsdp, with_checkpoint, with_bucketing))
diff --git a/tests/nn/data_parallel/test_fsdp_shared_weights.py b/tests/nn/data_parallel/test_fsdp_shared_weights.py
index e6711f9f2..b55f2dad5 100644
--- a/tests/nn/data_parallel/test_fsdp_shared_weights.py
+++ b/tests/nn/data_parallel/test_fsdp_shared_weights.py
@@ -46,9 +46,9 @@ def __init__(self, with_fsdp=False, inner_flat=False, sharing=None):
             self.l3 = FSDP(self.l3, flatten_parameters=False)
 
             if sharing in ["share_only_weights"]:
-                self.l3.append_shared_param(self.l1.module.weight)
+                self.l3.append_shared_param(self.l1.module.weigh, self.l1)
             if sharing in ["share_only_bias"]:
-                self.l3.append_shared_param(self.l1.module.bias)
+                self.l3.append_shared_param(self.l1.module.bias, self.l1)
 
     def forward(self, x):
         x = self.l0(x)
diff --git a/tests/nn/data_parallel/test_fsdp_shared_weights_mevo.py b/tests/nn/data_parallel/test_fsdp_shared_weights_mevo.py
index 149289ed9..108633462 100644
--- a/tests/nn/data_parallel/test_fsdp_shared_weights_mevo.py
+++ b/tests/nn/data_parallel/test_fsdp_shared_weights_mevo.py
@@ -62,7 +62,7 @@ def __init__(self, with_fsdp=False, wrap_middle="none"):
             # Shared layers must be un-flatten.
             self.l0 = FSDP(self.l0, flatten_parameters=False, mixed_precision=False, compute_dtype=torch.float16)
             self.l1 = FSDP(self.l1, flatten_parameters=False, mixed_precision=False, compute_dtype=torch.float16)
-            self.l1.append_shared_param(self.l0.module.weight)
+            self.l1.append_shared_param(self.l0.module.weight, self.l0)
             # These are for debugging.
             # print(id(self.l0), "is emb")
             # print(id(self.l1), "is out")

From 546fb062aecaa968da370bca08485cec7b36750c Mon Sep 17 00:00:00 2001
From: Ruan Silva <ruanrms@devfair0830.h2.fair>
Date: Tue, 2 Aug 2022 13:14:42 -0700
Subject: [PATCH 07/11] fix tests 2

---
 .../fully_sharded_data_parallel.py            | 88 +++++++++++--------
 .../data_parallel/test_fsdp_shared_weights.py |  4 +-
 .../test_fsdp_shared_weights_mevo.py          |  2 +-
 3 files changed, 56 insertions(+), 38 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index f6f8baa7f..c86ed1476 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -549,7 +549,7 @@ def module(self) -> FlattenParamsWrapper:
         assert isinstance(self._fsdp_wrapped_module, FlattenParamsWrapper)
         return self._fsdp_wrapped_module
 
-    def append_shared_param(self, p: Parameter, original_module: "FullyShardedDataParallel") -> None:
+    def append_shared_param(self, p: Parameter) -> None:
         """Add a param that's already owned by another FSDP wrapper.
 
             .. warning:: This is experimental!
@@ -576,7 +576,6 @@ def append_shared_param(self, p: Parameter, original_module: "FullyShardedDataPa
         ), "Must have at least 1 non-shared param."
         self.params.append(p)
         self._has_shared_params = True
-        original_module._has_shared_params = True
 
     def non_shared_params(self) -> List[nn.Parameter]:
         """Return the list of non-shared parameters."""
@@ -1404,21 +1403,6 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
         if self.force_input_to_fp32 and not self.mixed_precision:
             args, kwargs = cast_floats_to_right_precision(False, False, *args, **kwargs)
 
-        if self not in self._forward_ordering:
-            self._forward_ordering.append(self)
-        forward_idx = self._forward_ordering.index(self)
-        if self.rank == 0:
-            print(
-                "forward size = "
-                + str(len(self._forward_ordering))
-                + " idx = "
-                + str(forward_idx)
-                + " has shared = "
-                + str(self._has_shared_params)
-                + " has full params = "
-                + str(self.has_full_params)
-            )
-
         # All-gather full parameters. This will also transfer FP32 parameters to
         # ``self.compute_dtype`` (e.g., FP16 if *mixed_precision* is ``True``).
         self._rebuild_full_params()
@@ -1438,6 +1422,21 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
 
         outputs = self.module(*args, **kwargs)
 
+        if self not in self._forward_ordering:
+            self._forward_ordering.append(self)
+        forward_idx = self._forward_ordering.index(self)
+        if self.rank == 0:
+            print(
+                "forward size = "
+                + str(len(self._forward_ordering))
+                + " idx = "
+                + str(forward_idx)
+                + " has shared = "
+                + str(self._has_shared_params)
+                + " has full params = "
+                + str(self.has_full_params)
+            )
+
         # if(self.rank == 0):
         #     print(str(self.reshard_after_forward) + " forward_idx = " + str(self._forward_ordering.index(self)))
         if self.reshard_after_forward:
@@ -1988,13 +1987,21 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
             # Therefore, we update the flag accordingly here.
             self.has_full_params = not any(p._full_param_padded.storage().size() == 0 for p in self.params)
 
-        # 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
+        # Early exit if we already have full params.
+        if self.has_full_params:
+            assert (
+                force_full_precision and wait_for_all_gather
+            ) or not force_full_precision, (
+                "If you require full_precision, you need to wait for all_gather to be completed"
+            )
+            if not wait_for_all_gather:
+                return None
+            torch.cuda.current_stream().wait_stream(self._streams["all_gather"])
+
+            if not force_full_precision:
+                for p in self.params:
+                    update_p_data()
+                return output_tensors
 
         self.has_full_params = True
 
@@ -2054,8 +2061,9 @@ 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 wait_for_all_gather:
-            torch.cuda.current_stream().wait_stream(self._streams["all_gather"])
+        if not wait_for_all_gather:
+            return None
+        torch.cuda.current_stream().wait_stream(self._streams["all_gather"])
         return output_tensors
 
     @torch.no_grad()
@@ -2104,6 +2112,13 @@ def _prep_grads_for_backward(self) -> None:
     @torch.no_grad()
     def _free_full_params(self, params: Optional[List[Parameter]] = None) -> None:
         """Free up storage for full parameters."""
+        if self.rank == 0:
+            print(
+                "free full params "
+                + str(self.training_state)
+                + " fetching idx = "
+                + str(self._forward_ordering.index(self))
+            )
         if params is None:
             params = self.params
         self.has_full_params = False
@@ -2133,17 +2148,20 @@ def _schedule_next_all_gather_and_synchronize(self) -> None:
         ordering = self._forward_ordering
         if self.training_state == TrainingState.BACKWARD_POST:
             ordering = self._backward_ordering
-        next_idx = len(ordering)
         if self in ordering:
             next_idx = ordering.index(self) + 1
-        while next_idx < len(ordering):
-            next_module = ordering[next_idx]
-            if not next_module._pre_backward_hook_has_run and not next_module.has_full_params:
-                if self.rank == 0:
-                    print(str(self.training_state) + " fetching idx = " + str(next_idx))
-                next_module._rebuild_full_params(wait_for_all_gather=False)
-                break
-            next_idx = next_idx + 1
+            if next_idx < len(ordering):
+                next_module = ordering[next_idx]
+                # _pre_backward_hook_has_run prevents us from kicking off all-gather on a forward happening due to activation
+                # checkpointing. In these scenarios, forward only runs up until the module that already ran through the backward pass.
+                # If both modules have shared params, there is a potential race condition where params for current module are freed
+                # and all gather for the next module is happening, which may cause multiple all-gathers to be scheduled. So we just
+                # do not schedule all gathers if both modules have shared params.
+                if not next_module._pre_backward_hook_has_run and not next_module._has_shared_params:
+                    if self.rank == 0:
+                        print(str(self.training_state) + " fetching idx = " + str(next_idx))
+                    next_module._rebuild_full_params(wait_for_all_gather=False)
+        # Wait for computation kernels to finish running.
         torch.cuda.current_stream().synchronize()
 
     def local_metadata_dict(self) -> Dict[str, Any]:
diff --git a/tests/nn/data_parallel/test_fsdp_shared_weights.py b/tests/nn/data_parallel/test_fsdp_shared_weights.py
index b55f2dad5..e6711f9f2 100644
--- a/tests/nn/data_parallel/test_fsdp_shared_weights.py
+++ b/tests/nn/data_parallel/test_fsdp_shared_weights.py
@@ -46,9 +46,9 @@ def __init__(self, with_fsdp=False, inner_flat=False, sharing=None):
             self.l3 = FSDP(self.l3, flatten_parameters=False)
 
             if sharing in ["share_only_weights"]:
-                self.l3.append_shared_param(self.l1.module.weigh, self.l1)
+                self.l3.append_shared_param(self.l1.module.weight)
             if sharing in ["share_only_bias"]:
-                self.l3.append_shared_param(self.l1.module.bias, self.l1)
+                self.l3.append_shared_param(self.l1.module.bias)
 
     def forward(self, x):
         x = self.l0(x)
diff --git a/tests/nn/data_parallel/test_fsdp_shared_weights_mevo.py b/tests/nn/data_parallel/test_fsdp_shared_weights_mevo.py
index 108633462..149289ed9 100644
--- a/tests/nn/data_parallel/test_fsdp_shared_weights_mevo.py
+++ b/tests/nn/data_parallel/test_fsdp_shared_weights_mevo.py
@@ -62,7 +62,7 @@ def __init__(self, with_fsdp=False, wrap_middle="none"):
             # Shared layers must be un-flatten.
             self.l0 = FSDP(self.l0, flatten_parameters=False, mixed_precision=False, compute_dtype=torch.float16)
             self.l1 = FSDP(self.l1, flatten_parameters=False, mixed_precision=False, compute_dtype=torch.float16)
-            self.l1.append_shared_param(self.l0.module.weight, self.l0)
+            self.l1.append_shared_param(self.l0.module.weight)
             # These are for debugging.
             # print(id(self.l0), "is emb")
             # print(id(self.l1), "is out")

From b2de8fd3da3f8fd13db58d10980c2e90e2d79b58 Mon Sep 17 00:00:00 2001
From: Ruan Silva <ruanrms@devfair0830.h2.fair>
Date: Tue, 2 Aug 2022 15:25:27 -0700
Subject: [PATCH 08/11] remove prints

---
 .../fully_sharded_data_parallel.py            | 112 +++++-------------
 1 file changed, 31 insertions(+), 81 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index c86ed1476..629163253 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -1407,38 +1407,19 @@ 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 (
-        #     forward_idx < len(self._forward_ordering) - 1
-        # ):
-        #     t = self._forward_ordering[forward_idx + 1]
-        #     if not t._pre_backward_hook_has_run:
-        #         if (self.rank == 0):
-        #             print("fetching idx = " + str(forward_idx + 1))
-        #         t._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()
 
         outputs = self.module(*args, **kwargs)
 
+        # In the first forward pass, track the order that modules are computed.
+        # In the following passes, we assume that the order remains the same to
+        # to kick off the all-gather for the next module in the list, in case we
+        # are waiting for the computation to finish.
         if self not in self._forward_ordering:
             self._forward_ordering.append(self)
-        forward_idx = self._forward_ordering.index(self)
-        if self.rank == 0:
-            print(
-                "forward size = "
-                + str(len(self._forward_ordering))
-                + " idx = "
-                + str(forward_idx)
-                + " has shared = "
-                + str(self._has_shared_params)
-                + " has full params = "
-                + str(self.has_full_params)
-            )
 
-        # if(self.rank == 0):
-        #     print(str(self.reshard_after_forward) + " forward_idx = " + str(self._forward_ordering.index(self)))
         if self.reshard_after_forward:
             self._free_full_params()
             if self.mixed_precision or self.move_params_to_cpu:
@@ -1521,36 +1502,14 @@ def _pre_backward_hook(*unused: Any) -> None:
             # idempotent.  So in case they are called unnecessarily, they don't incur much
             # overhead.
             if self.reshard_after_forward:
-                if self not in self._backward_ordering:
-                    self._backward_ordering.append(self)
-                backward_idx = self._backward_ordering.index(self)
-                if self.rank == 0:
-                    print(
-                        "pre-backward size = "
-                        + str(len(self._backward_ordering))
-                        + " idx = "
-                        + str(backward_idx)
-                        + " forward_idx = "
-                        + str(self._forward_ordering.index(self))
-                        + "has shared = "
-                        + str(self._has_shared_params)
-                        + " has full params = "
-                        + str(self.has_full_params)
-                    )
-
                 self._rebuild_full_params()
 
-                # if (
-                #     backward_idx < len(self._backward_ordering) - 1
-                # ):
-                #     t = self._backward_ordering[backward_idx + 1]
-                #     # is_hook_registered = id(t) in self._output_pre_backward_hook_registered
-                #     # # if (self.rank == 0):
-                #     # #     print("is_hook_registered = " + str(is_hook_registered) + " id = " + str(id(t)) + " len = " + str(len(self._output_pre_backward_hook_registered)))
-                #     # if not is_hook_registered:
-                #     if (self.rank == 0):
-                #         print("fetching backward idx = " + str(backward_idx + 1) + " forward_idx = " + str(self._forward_ordering.index(t)))
-                #     t._rebuild_full_params(wait_for_all_gather=False)
+                # Similar to _forward_ordering, in the first backward pass we track the order
+                # that weights were gathered for modules in the backward pass. Then, we use
+                # this order in future passes to kick off the all-gather for the next module
+                # in case we are waiting for the computation of the current module to finish.
+                if self not in self._backward_ordering:
+                    self._backward_ordering.append(self)
             else:
                 self._use_full_params()
 
@@ -1678,8 +1637,6 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
         """
         # First hook callback will see PRE state. If we have multiple params,
         # then subsequent hook callbacks will see POST state.
-        # if (self.rank == 0):
-        #     print("post backward forward_idx = " + str(self._forward_ordering.index(self)))
         self.assert_state([TrainingState.BACKWARD_PRE, TrainingState.BACKWARD_POST])
         self.training_state = TrainingState.BACKWARD_POST
         if param.grad is None:
@@ -1987,21 +1944,14 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
             # Therefore, we update the flag accordingly here.
             self.has_full_params = not any(p._full_param_padded.storage().size() == 0 for p in self.params)
 
-        # Early exit if we already have full params.
-        if self.has_full_params:
-            assert (
-                force_full_precision and wait_for_all_gather
-            ) or not force_full_precision, (
-                "If you require full_precision, you need to wait for all_gather to be completed"
-            )
+        # 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 not wait_for_all_gather:
                 return None
             torch.cuda.current_stream().wait_stream(self._streams["all_gather"])
-
-            if not force_full_precision:
-                for p in self.params:
-                    update_p_data()
-                return output_tensors
+            for p in self.params:
+                update_p_data()
+            return output_tensors
 
         self.has_full_params = True
 
@@ -2112,13 +2062,6 @@ def _prep_grads_for_backward(self) -> None:
     @torch.no_grad()
     def _free_full_params(self, params: Optional[List[Parameter]] = None) -> None:
         """Free up storage for full parameters."""
-        if self.rank == 0:
-            print(
-                "free full params "
-                + str(self.training_state)
-                + " fetching idx = "
-                + str(self._forward_ordering.index(self))
-            )
         if params is None:
             params = self.params
         self.has_full_params = False
@@ -2138,9 +2081,16 @@ 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)
+
+        # When we are memory bound (which is the case here as we are freeing up
+        # params), we are not able to let the CPU run completely free because
+        # it will end up scheduling GPU operations required to compute all
+        # future modules. This causes significant increases in GPU reserved
+        # memory and potential thrashing.
+        # So instead, we simply schedule the all-gather for the next module
+        # to be executed and wait for the computations of the current module
+        # to finish before moving forward.
         self._schedule_next_all_gather_and_synchronize()
-        # if(self.rank == 0):
-        #     print(str(self.training_state) + " forward_idx = " + str(self._forward_ordering.index(self)))
 
     @torch.no_grad()
     def _schedule_next_all_gather_and_synchronize(self) -> None:
@@ -2148,18 +2098,18 @@ def _schedule_next_all_gather_and_synchronize(self) -> None:
         ordering = self._forward_ordering
         if self.training_state == TrainingState.BACKWARD_POST:
             ordering = self._backward_ordering
+        # With activation checkpointing, we may have modules in the backward pass that are
+        # not part of _backward_ordering. So this check is required.
         if self in ordering:
             next_idx = ordering.index(self) + 1
             if next_idx < len(ordering):
                 next_module = ordering[next_idx]
-                # _pre_backward_hook_has_run prevents us from kicking off all-gather on a forward happening due to activation
-                # checkpointing. In these scenarios, forward only runs up until the module that already ran through the backward pass.
-                # If both modules have shared params, there is a potential race condition where params for current module are freed
-                # and all gather for the next module is happening, which may cause multiple all-gathers to be scheduled. So we just
-                # do not schedule all gathers if both modules have shared params.
+                # _pre_backward_hook_has_run prevents us from kicking off all-gather on a forward pass happening due to activation
+                # checkpointing. In these scenarios, forward only runs up until the module that already went through the backward pass.
+                # In addition, if the module to be scheduled has a shared param, there is a potential race condition where params for
+                # the current module are freed and all gather for the next module is happening. So we just skip such modules.
                 if not next_module._pre_backward_hook_has_run and not next_module._has_shared_params:
-                    if self.rank == 0:
-                        print(str(self.training_state) + " fetching idx = " + str(next_idx))
+                    # Kick-off all gather for the next module without waiting.
                     next_module._rebuild_full_params(wait_for_all_gather=False)
         # Wait for computation kernels to finish running.
         torch.cuda.current_stream().synchronize()

From a9b071fc2cbb358769acc7312b3c5f67d765e702 Mon Sep 17 00:00:00 2001
From: Ruan Silva <ruanrms@devfair0830.h2.fair>
Date: Wed, 3 Aug 2022 14:08:11 -0700
Subject: [PATCH 09/11] add unit-test

---
 .../fully_sharded_data_parallel.py            |  18 +--
 tests/nn/data_parallel/test_fsdp_overlap.py   |  31 +----
 .../data_parallel/test_fsdp_prefetch_order.py | 111 ++++++++++++++++++
 3 files changed, 123 insertions(+), 37 deletions(-)
 create mode 100644 tests/nn/data_parallel/test_fsdp_prefetch_order.py

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 629163253..fb8cae4ac 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -1163,7 +1163,7 @@ def _reset_lazy_init(self) -> None:
         self._streams: Dict[str, torch.cuda.Stream] = {}
         self._reducer: Optional[ReduceScatterBucketer] = None
         self._forward_ordering: List[FullyShardedDataParallel] = []
-        self._backward_ordering: List[FullyShardedDataParallel] = []
+        self._backward_rebuild_ordering: List[FullyShardedDataParallel] = []
         for p in self.params:
             if hasattr(p, "_fp32_shard"):
                 del p._fp32_shard  # reset _init_param_attributes
@@ -1337,7 +1337,7 @@ 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._forward_ordering = self._forward_ordering
-                m._backward_ordering = self._backward_ordering
+                m._backward_rebuild_ordering = self._backward_rebuild_ordering
 
     def _setup_streams(self) -> None:
         """Create streams to overlap data transfer and computation."""
@@ -1508,8 +1508,8 @@ def _pre_backward_hook(*unused: Any) -> None:
                 # that weights were gathered for modules in the backward pass. Then, we use
                 # this order in future passes to kick off the all-gather for the next module
                 # in case we are waiting for the computation of the current module to finish.
-                if self not in self._backward_ordering:
-                    self._backward_ordering.append(self)
+                if self not in self._backward_rebuild_ordering:
+                    self._backward_rebuild_ordering.append(self)
             else:
                 self._use_full_params()
 
@@ -2067,6 +2067,9 @@ def _free_full_params(self, params: Optional[List[Parameter]] = None) -> None:
         self.has_full_params = False
         current_stream = torch.cuda.current_stream()
         for p in params:
+            # Shared params are not owned by this FSDP instance.
+            if hasattr(p, "_is_shared") and p._is_shared:
+                continue
             if not p._is_sharded:  # e.g., world_size == 1
                 if self.mixed_precision or self.move_params_to_cpu:
                     self._free_fp16_param_shard([p])
@@ -2097,15 +2100,14 @@ def _schedule_next_all_gather_and_synchronize(self) -> None:
         self.assert_state([TrainingState.FORWARD, TrainingState.BACKWARD_POST])
         ordering = self._forward_ordering
         if self.training_state == TrainingState.BACKWARD_POST:
-            ordering = self._backward_ordering
-        # With activation checkpointing, we may have modules in the backward pass that are
-        # not part of _backward_ordering. So this check is required.
+            ordering = self._backward_rebuild_ordering
+        # Not all modules require rebuilding in backward pass, so this check is required.
         if self in ordering:
             next_idx = ordering.index(self) + 1
             if next_idx < len(ordering):
                 next_module = ordering[next_idx]
                 # _pre_backward_hook_has_run prevents us from kicking off all-gather on a forward pass happening due to activation
-                # checkpointing. In these scenarios, forward only runs up until the module that already went through the backward pass.
+                # checkpointing. In these scenarios, forward only runs up until the module that already ran their backward hook.
                 # In addition, if the module to be scheduled has a shared param, there is a potential race condition where params for
                 # the current module are freed and all gather for the next module is happening. So we just skip such modules.
                 if not next_module._pre_backward_hook_has_run and not next_module._has_shared_params:
diff --git a/tests/nn/data_parallel/test_fsdp_overlap.py b/tests/nn/data_parallel/test_fsdp_overlap.py
index e1e584609..87f86cb4f 100644
--- a/tests/nn/data_parallel/test_fsdp_overlap.py
+++ b/tests/nn/data_parallel/test_fsdp_overlap.py
@@ -112,8 +112,6 @@ def run(compute_cycles, all_gather_cycles):
 
         # We run 20 iterations but only collect timing data from the minimal 10
         # data points because nondeterministic system events can disturb the timing.
-        cpu_iter = Min10()
-        cpu_wait = Min10()
         gpu_compute = Min10()
         gpu_total = Min10()
         for _ in range(20):
@@ -165,12 +163,7 @@ def _delayed_all_gather_base(*args, **kwargs):
             else:
                 for p in model.parameters():
                     p.grad = None
-
-            cpu_iter_time = time.process_time() - cpu_start
-
-            # wait for gpu
             out.item()
-            cpu_wait_for_gpu_time = time.process_time() - cpu_start - cpu_iter_time
 
             # get sum of the compute time
             times = []
@@ -182,15 +175,11 @@ def _delayed_all_gather_base(*args, **kwargs):
             # get gpu compute + all_gather time
             overall_gpu_time = e1.elapsed_time(e2)
 
-            cpu_iter.add(cpu_iter_time)
-            cpu_wait.add(cpu_wait_for_gpu_time)
             gpu_compute.add(sum(times))
             gpu_total.add(overall_gpu_time)
 
         del model
         return {
-            "cpu_iter": cpu_iter.avg(),
-            "cpu_wait": cpu_wait.avg(),
             "gpu_compute": gpu_compute.avg(),
             "gpu_total": gpu_total.avg(),
         }
@@ -204,33 +193,17 @@ def _delayed_all_gather_base(*args, **kwargs):
     debug_string = f"\nrank{rank}:\n  e1: {e1}\n  e2: {e2}\n  e3: {e3}\n  e4: {e4}"
     print(debug_string)
 
-    # Check the cpu/gpu timing. CPU should run ahead of GPU. Therefore, cpu-gpu
-    # wait should be long, except when there is no real work on GPU.
-    #
-    # If the assertions fail below, we likely have a cpu-gpu wait in the forward/backward pass.
-    short = [e1["cpu_iter"], e2["cpu_iter"], e3["cpu_iter"], e4["cpu_iter"], e1["cpu_wait"]]
-    long = [e3["cpu_wait"], e4["cpu_wait"]]
-    if world_size == 1:
-        short.append(e2["cpu_wait"])  # all gather should not be happening.
-    else:
-        long.append(e2["cpu_wait"])  # all gather should happen and prolong the cpu-gpu wait.
-    for s in short:
-        for l in long:
-            # 5X longer is a safe margin, since the GPU work timing is around 100X more
-            # of that of the CPU.
-            assert s * 5 < l, f"{s} * 5 < {l} in " + debug_string
-
     # Check the GPU timing.
     short = [e1["gpu_compute"], e1["gpu_total"], e2["gpu_compute"]]
     long = [e3["gpu_compute"], e3["gpu_total"], e4["gpu_compute"], e4["gpu_total"]]
     if world_size == 1:
         short.append(e2["gpu_total"])  # all gather should not be happening.
     else:
-        long.append(e2["gpu_total"])  # all gather should happen and prolong the cpu-gpu wait.
+        long.append(e2["gpu_total"])  # all gather should happen and prolong the gpu wait.
     for s in short:
         for l in long:
             # 10X longer is a safe margin, since the time is around 100X longer
-            # when there is work on GPU vs. no work.
+            # when there is compute work on GPU vs. no work.
             assert s * 10 < l, f"{s} * 10 < {l} in " + debug_string
 
     # Check the GPU overlapping when there is all-gather.
diff --git a/tests/nn/data_parallel/test_fsdp_prefetch_order.py b/tests/nn/data_parallel/test_fsdp_prefetch_order.py
new file mode 100644
index 000000000..631aa8107
--- /dev/null
+++ b/tests/nn/data_parallel/test_fsdp_prefetch_order.py
@@ -0,0 +1,111 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+#
+# This source code is licensed under the BSD license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pylint: disable=missing-module-docstring
+# pylint: disable=missing-class-docstring
+# pylint: disable=missing-function-docstring
+
+""" Check that the ordering used for prefetching model weights
+    matches the expected execution order for the model.
+"""
+
+import tempfile
+
+import pytest
+import torch
+import torch.multiprocessing as mp
+import torch.nn as nn
+from torch.optim import SGD
+
+from fair_dev.testing.testing import dist_init, skip_if_single_gpu, teardown
+from fairscale.internal import torch_version
+from fairscale.nn import checkpoint_wrapper
+from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
+from fairscale.nn.data_parallel import TrainingState, auto_wrap_bn
+from fairscale.nn.wrap import enable_wrap, wrap
+
+
+def _get_module_type(fsdp):
+    m = fsdp._fsdp_wrapped_module._fpw_module
+    if isinstance(m, nn.Sequential):
+        return type(m[0])
+    return type(m)
+
+
+def _test_func(rank, world_size, fsdp_config, tempfile_name, unused):
+    result = dist_init(rank, world_size, tempfile_name, unused)
+    assert result, "Dist init failed"
+
+    assert isinstance(fsdp_config, dict), str(fsdp_config)
+
+    torch.cuda.set_device(rank)
+
+    class Model(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.Linear(12, 8)
+            self.head = nn.Sequential(nn.AdaptiveAvgPool2d(output_size=(1, 1)), nn.Flatten(), nn.Linear(4, 10))
+
+        def forward(self, x):
+            return self.head(self.block3(self.block2(self.block1(x))))
+
+    model = Model()
+    # Wrapping BatchNorm as separate modules for the forward pass.
+    model.block1 = auto_wrap_bn(model.block1, fsdp_config={"reshard_after_forward": True})
+    model.block2 = auto_wrap_bn(model.block2, fsdp_config={"reshard_after_forward": True})
+
+    # Checkpoints shouldn't affect the ordering.
+    model.block2 = checkpoint_wrapper(model.block2)
+
+    with enable_wrap(
+        wrapper_cls=FSDP,
+    ):
+        model.block1 = wrap(model.block1)
+        model.block2 = wrap(model.block2)
+        model.block3 = wrap(model.block3)
+        model = wrap(model)
+
+    optim = SGD(model.parameters(), lr=0.1)
+    model = model.cuda()
+
+    # Orderings are stored in the first pass.
+    in_data = torch.randn(size=(2, 3, 16, 16)).cuda()
+    in_data.requires_grad = True
+    out = model(in_data)
+    out.sum().backward()
+    optim.step()
+
+    expected_forward_ordering = [nn.BatchNorm2d, nn.Conv2d, nn.BatchNorm2d, nn.Conv2d, nn.Linear, Model]
+    actual_forward_ordering = [_get_module_type(m) for m in model._forward_ordering]
+    assert expected_forward_ordering == actual_forward_ordering
+
+    expected_backward_ordering = [nn.Linear, nn.Conv2d, nn.BatchNorm2d, nn.BatchNorm2d, nn.Conv2d]
+    actual_backward_ordering = [_get_module_type(m) for m in model._backward_rebuild_ordering]
+    assert expected_backward_ordering == actual_backward_ordering
+
+    model.assert_state(TrainingState.IDLE)
+    teardown()
+
+
+@skip_if_single_gpu
+def test():
+    if torch_version() < (1, 6, 0):
+        pytest.skip("older pytorch doesn't support reduce_scatter")
+
+    temp_file_name = tempfile.mkstemp()[1]
+    unused = tempfile.mkstemp()[1]
+
+    fsdp_config = {}
+
+    # Using world_size > 1 to trigger all-gathers.
+    world_size = 2
+    mp.spawn(
+        _test_func,
+        args=(world_size, fsdp_config, temp_file_name, unused),
+        nprocs=world_size,
+        join=True,
+    )

From 8acbde051f48de7cbba0a4a9348f3d722274c2cf Mon Sep 17 00:00:00 2001
From: Ruan Silva <ruanrms@devfair0830.h2.fair>
Date: Wed, 3 Aug 2022 15:46:32 -0700
Subject: [PATCH 10/11] add new test to CI list

---
 tests/ci_test_list_1.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/ci_test_list_1.txt b/tests/ci_test_list_1.txt
index aa2fe4460..45678bd57 100644
--- a/tests/ci_test_list_1.txt
+++ b/tests/ci_test_list_1.txt
@@ -9,4 +9,5 @@ tests/nn/data_parallel/test_fsdp_input.py
 tests/nn/data_parallel/test_fsdp_optimizer_utils.py
 tests/nn/data_parallel/test_fsdp.py
 tests/nn/data_parallel/test_fsdp_with_checkpoint_wrapper.py
+tests/nn/data_parallel/test_fsdp_prefetch_order.py
 tests/optim/test_layerwise_gradient_scaler.py

From 7d46cba0ac2bc2d69922d75a454c08edf07bb6ce Mon Sep 17 00:00:00 2001
From: Ruan Silva <ruanrms@devfair0830.h2.fair>
Date: Thu, 4 Aug 2022 06:42:48 -0700
Subject: [PATCH 11/11] skip modules with ssd_offload

---
 .../nn/data_parallel/fully_sharded_data_parallel.py      | 9 +++++++--
 1 file changed, 7 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 acf4da7ac..c97cdeeae 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -2119,8 +2119,13 @@ def _schedule_next_all_gather_and_synchronize(self) -> None:
                 # _pre_backward_hook_has_run prevents us from kicking off all-gather on a forward pass happening due to activation
                 # checkpointing. In these scenarios, forward only runs up until the module that already ran their backward hook.
                 # In addition, if the module to be scheduled has a shared param, there is a potential race condition where params for
-                # the current module are freed and all gather for the next module is happening. So we just skip such modules.
-                if not next_module._pre_backward_hook_has_run and not next_module._has_shared_params:
+                # the current module are freed and all gather for the next module is happening. Similarly, modules with ssd_offload
+                # are not supported because ssd_offload happens before every all-gather call. So we just skip such modules.
+                if (
+                    not next_module._pre_backward_hook_has_run
+                    and not next_module._has_shared_params
+                    and not next_module.ssd_offload
+                ):
                     # Kick-off all gather for the next module without waiting.
                     next_module._rebuild_full_params(wait_for_all_gather=False)
         # Wait for computation kernels to finish running.