diff --git a/test/distributed/_composable/test_compose.py b/test/distributed/_composable/test_compose.py
index 3500a35..643530d 100644
--- a/test/distributed/_composable/test_compose.py
+++ b/test/distributed/_composable/test_compose.py
@@ -42,18 +42,28 @@
     def world_size(self) -> int:
         return 2
 
+    # TODO: Define `use_same_inputs_across_ranks` for now for BC since some
+    # test model configs do not have a simple base model to compare against. In
+    # those cases, we use the same inputs across ranks so that the averaged
+    # gradient equals the local gradient to check for parity. This means that
+    # the gradient reduction is unchecked.
     def _test_parity(
         self,
         base_model: nn.Module,
         test_model: nn.Module,
-        x: torch.Tensor,
+        inp_size: torch.Size,
+        inp_device: torch.device,
         grad_to_none: bool,
+        use_same_inputs_across_ranks: bool,
     ):
         LR = 0.01
         base_optim = torch.optim.Adam(base_model.parameters(), lr=LR)
         test_optim = torch.optim.Adam(test_model.parameters(), lr=LR)
 
         for _ in range(5):
+            if use_same_inputs_across_ranks:
+                torch.manual_seed(0)
+            x = torch.randn(inp_size, device=inp_device)
             test_loss = test_model(x).sum()
             base_loss = base_model(x).sum()
 
@@ -86,8 +96,10 @@
             {
                 "base_model": [base_model],
                 "test_model": [test_model],
-                "x": [torch.randn(2, 100, device="cuda")],
+                "inp_size": [torch.Size((2, 100))],
+                "inp_device": [torch.device("cuda")],
                 "grad_to_none": [True, False],
+                "use_same_inputs_across_ranks": [True],
             },
             self._test_parity,
         )
@@ -108,8 +120,10 @@
             {
                 "base_model": [base_model],
                 "test_model": [test_model],
-                "x": [torch.randn(2, 100, device="cuda")],
+                "inp_size": [torch.Size((2, 100))],
+                "inp_device": [torch.device("cuda")],
                 "grad_to_none": [True, False],
+                "use_same_inputs_across_ranks": [True],
             },
             self._test_parity,
         )
@@ -144,8 +158,10 @@
             {
                 "base_model": [base_model],
                 "test_model": [test_model],
-                "x": [torch.randn(2, 100, device="cuda")],
+                "inp_size": [torch.Size((2, 100))],
+                "inp_device": [torch.device("cuda")],
                 "grad_to_none": [True, False],
+                "use_same_inputs_across_ranks": [True],
             },
             self._test_parity,
         )
@@ -165,8 +181,10 @@
             {
                 "base_model": [base_model],
                 "test_model": [test_model],
-                "x": [torch.randn(2, 100, device="cuda")],
+                "inp_size": [torch.Size((2, 100))],
+                "inp_device": [torch.device("cuda")],
                 "grad_to_none": [True, False],
+                "use_same_inputs_across_ranks": [True],
             },
             self._test_parity,
         )
@@ -182,20 +200,83 @@
         replicate(model.l2)  # should not raise
 
     @skip_if_lt_x_gpu(2)
-    def test_fsdp_in_replicate(self):
-        model = CompositeModel(device=torch.device("cuda"))
-        base_model = copy.deepcopy(model)
-        test_model = copy.deepcopy(model)
-        fully_shard(test_model.l1)
-        replicate(test_model)
+    def test_fully_shard_replicate_composability(self):
+        """
+        Tests composing ``fully_shard`` and ``replicate``. To save unit test
+        time, we run the different configs in subtests.
+        """
         self.run_subtests(
             {
-                "base_model": [base_model],
-                "test_model": [test_model],
-                "x": [torch.randn(2, 100, device="cuda")],
-                "grad_to_none": [True, False],
+                "config": [
+                    "1fm,1r",
+                    "1r,1fm",
+                    "1r,1fa",
+                    "1r1fm,1fm",
+                    "1r1fa,1fm",
+                    "1fm1fm,1r1r,1fm",
+                ]
             },
-            self._test_parity,
+            self._test_replicate_in_fully_shard,
+        )
+
+    def _test_replicate_in_fully_shard(self, config: str):
+        """
+        To interpret the config, each comma delineates a level in the module
+        tree ordered bottom-up; 'r' means ``replicate``; 'f' means
+        ``fully_shard``; 'a' means auto wrap; and 'm' means manual wrap.
+        """
+        # Set the seed to ensure that all ranks initialize the same model
+        torch.manual_seed(0)
+        if config == "1fm,1r":
+            base_model = CompositeModel(device=torch.device("cuda"))
+            test_model = copy.deepcopy(base_model)
+            fully_shard(test_model.l1)
+            replicate(test_model)
+        elif config == "1r,1fm":
+            base_model = CompositeParamModel(torch.device("cuda"))
+            test_model = copy.deepcopy(base_model)
+            replicate(test_model.u1)
+            fully_shard(test_model)
+        elif config == "1r,1fa":
+            base_model = CompositeParamModel(torch.device("cuda"))
+            test_model = copy.deepcopy(base_model)
+            replicate(test_model.u1)
+            fully_shard(test_model, policy=ModuleWrapPolicy({UnitModule}))
+        elif config == "1r1fm,1fm":
+            base_model = CompositeParamModel(torch.device("cuda"))
+            test_model = copy.deepcopy(base_model)
+            replicate(test_model.u1)
+            fully_shard(test_model.u2)
+            fully_shard(test_model)
+        elif config == "1r1fa,1fm":
+            base_model = CompositeParamModel(torch.device("cuda"))
+            test_model = copy.deepcopy(base_model)
+            replicate(test_model.u1)
+            fully_shard(test_model.u2, policy=ModuleWrapPolicy({UnitModule}))
+            fully_shard(test_model)
+        elif config == "1fm1fm,1r1r,1fm":
+            base_model = CompositeParamModel(torch.device("cuda"))
+            test_model = copy.deepcopy(base_model)
+            fully_shard(test_model.u1.seq)
+            fully_shard(test_model.u2.seq)
+            replicate(test_model.u1)
+            replicate(test_model.u2)
+            fully_shard(test_model)
+        else:
+            raise ValueError(f"Unknown config: {config}")
+        # Apply data parallelism to the base model for parity since we apply
+        # data parallelism to the test model
+        replicate(base_model)
+
+        # Set the seed to ensure that ranks get different input data
+        torch.manual_seed(self.rank + 1)
+        self._test_parity(
+            base_model,
+            test_model,
+            torch.Size((2, 100)),
+            torch.device("cuda"),
+            True,
+            False,
         )
 
 
diff --git a/test/distributed/_composable/test_fully_shard.py b/test/distributed/_composable/test_fully_shard.py
index da14aac..6eb59ba 100644
--- a/test/distributed/_composable/test_fully_shard.py
+++ b/test/distributed/_composable/test_fully_shard.py
@@ -3,35 +3,38 @@
 import contextlib
 import copy
 import functools
+import itertools
 import sys
 from enum import auto, Enum
 from typing import Callable, Dict, Iterable, List, Optional, Tuple
-import itertools
 
 import torch
 import torch.distributed as dist
+import torch.distributed.fsdp._traversal_utils as traversal_utils
 import torch.nn as nn
-from torch.nn import TransformerEncoderLayer, TransformerDecoderLayer
 from torch.distributed._composable import fully_shard
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-from torch.distributed.fsdp._common_utils import (
-    _FSDPState,
-    _get_fsdp_handles,
-    _is_fsdp_flattened,
-)
-from torch.testing._internal.common_fsdp import _zero_model
+from torch.distributed.fsdp._common_utils import _FSDPState, _is_fsdp_flattened
 from torch.distributed.fsdp.api import MixedPrecision
 from torch.distributed.fsdp.flat_param import _HandlesKey, FlatParamHandle
 from torch.distributed.fsdp.wrap import _FSDPPolicy, ModuleWrapPolicy
+from torch.nn import TransformerDecoderLayer, TransformerEncoderLayer
 from torch.testing._internal.common_dist_composable import (
     CompositeParamModel,
+    NestedSequentialModel,
     UnitModule,
 )
 from torch.testing._internal.common_distributed import (
     SaveForwardInputsModel,
     skip_if_lt_x_gpu,
 )
-from torch.testing._internal.common_fsdp import FSDPTest, TransformerWithSharedParams, CUDAInitMode, FSDPInitMode
+from torch.testing._internal.common_fsdp import (
+    _zero_model,
+    CUDAInitMode,
+    FSDPInitMode,
+    FSDPTest,
+    TransformerWithSharedParams,
+)
 from torch.testing._internal.common_utils import run_tests, TEST_WITH_DEV_DBG_ASAN
 
 if not dist.is_available():
@@ -62,12 +65,25 @@
     def test_policy(self):
         """Tests passing a ``policy`` for pseudo-auto-wrapping."""
         self.run_subtests(
-            {"policy": [None, ModuleWrapPolicy({UnitModule})]},
+            {
+                "policy": [
+                    None,
+                    ModuleWrapPolicy({UnitModule}),
+                    ModuleWrapPolicy({nn.Sequential}),
+                ],
+            },
             self._test_policy,
         )
 
     def _test_policy(self, policy: Optional[_FSDPPolicy]):
-        local_model = CompositeParamModel(torch.device("cuda"))
+        use_nested_sequential_model = "Sequential" in getattr(
+            policy, "_module_classes_str", ""
+        )
+        local_model = (
+            NestedSequentialModel(torch.device("cuda"))
+            if use_nested_sequential_model
+            else CompositeParamModel(torch.device("cuda"))
+        )
         fsdp_wrapped_model = FSDP(
             copy.deepcopy(local_model),
             auto_wrap_policy=policy,
@@ -122,8 +138,8 @@
 
         # Check that the composable module has the same  `FlatParameter`
         # construction as the FSDP-wrapped model
-        composable_handles = _get_fsdp_handles(composable_module)
-        fsdp_wrapped_handles = _get_fsdp_handles(fsdp_wrapped_model)
+        composable_handles = traversal_utils._get_fsdp_handles(composable_module)
+        fsdp_wrapped_handles = traversal_utils._get_fsdp_handles(fsdp_wrapped_model)
         self.assertEqual(len(composable_handles), len(fsdp_wrapped_handles))
         for (composable_handle, fsdp_wrapped_handle) in zip(
             composable_handles, fsdp_wrapped_handles
@@ -131,6 +147,10 @@
             self.assertEqual(
                 composable_handle.flat_param.shape, fsdp_wrapped_handle.flat_param.shape
             )
+            self.assertEqual(
+                composable_handle.flat_param._fqns,
+                fsdp_wrapped_handle.flat_param._fqns,
+            )
 
         # Check that the composable module does not add any wrapper class
         local_module_classes = set()
@@ -384,8 +404,8 @@
         ) = self._init_models_and_optims(device, fsdp_wrap_mode)
         # Before checking the unshard/reshard order, sanity check that the
         # assumption about wrapper FQN being a suffix of composable FQN holds
-        all_composable_handles = _get_fsdp_handles(composable_module)
-        all_wrapped_handles = _get_fsdp_handles(fsdp_wrapped_model)
+        all_composable_handles = traversal_utils._get_fsdp_handles(composable_module)
+        all_wrapped_handles = traversal_utils._get_fsdp_handles(fsdp_wrapped_model)
         self._check_same_param_handles(all_composable_handles, all_wrapped_handles)
         num_handles = len(all_composable_handles)
 
@@ -564,8 +584,7 @@
 
         # Validate load
         load_composable = fully_shard(
-            copy.deepcopy(local_model),
-            policy=ModuleWrapPolicy({UnitModule})
+            copy.deepcopy(local_model), policy=ModuleWrapPolicy({UnitModule})
         )
         _zero_model(load_composable, summon_full=False)
         for p in load_composable.parameters():
@@ -634,10 +653,12 @@
                 save_model = copy.deepcopy(local_model)
                 save_model = fully_shard(
                     save_model,
-                    policy=ModuleWrapPolicy({TransformerEncoderLayer, TransformerDecoderLayer}),
+                    policy=ModuleWrapPolicy(
+                        {TransformerEncoderLayer, TransformerDecoderLayer}
+                    ),
                     ignored_modules=(
                         save_model.get_ignored_modules() if ignore_modules else []
-                    )
+                    ),
                 )
 
                 # TODO: test state_dict_type after https://github.com/pytorch/pytorch/issues/90954 is resolved
@@ -654,10 +675,12 @@
                 _zero_model(load_model, zero_buffers=True, summon_full=False)
                 fully_shard(
                     load_model,
-                    policy=ModuleWrapPolicy({TransformerDecoderLayer, TransformerEncoderLayer}),
+                    policy=ModuleWrapPolicy(
+                        {TransformerDecoderLayer, TransformerEncoderLayer}
+                    ),
                     ignored_modules=(
                         load_model.get_ignored_modules() if ignore_modules else []
-                    )
+                    ),
                 )
                 load_model.load_state_dict(state_dict)
                 self._check_model_parity(load_model, save_model)
@@ -670,7 +693,9 @@
         for k in composable_sd.keys():
             v1 = composable_sd[k]
             v2 = local_sd[k]
-            self.assertEqual(v1.shape, v2.shape, f"Shape mismatch for {k} {v1.shape} vs {v2.shape}")
+            self.assertEqual(
+                v1.shape, v2.shape, f"Shape mismatch for {k} {v1.shape} vs {v2.shape}"
+            )
 
         # Check actual values
         for k in composable_sd.keys():
@@ -682,9 +707,7 @@
         """
         Checks that m1 and m2 have equivalent named_parameters.
         """
-        for (n1, p1), (n2, p2) in zip(
-            m1.named_parameters(), m2.named_parameters()
-        ):
+        for (n1, p1), (n2, p2) in zip(m1.named_parameters(), m2.named_parameters()):
             self.assertEqual(n1, n2)
             self.assertEqual(p1, p2)
 
diff --git a/test/distributed/fsdp/test_fsdp_hybrid_shard.py b/test/distributed/fsdp/test_fsdp_hybrid_shard.py
index 2f7d3e2..aa2e226 100644
--- a/test/distributed/fsdp/test_fsdp_hybrid_shard.py
+++ b/test/distributed/fsdp/test_fsdp_hybrid_shard.py
@@ -8,6 +8,7 @@
 
 import torch
 import torch.distributed as dist
+import torch.distributed.fsdp._traversal_utils as traversal_utils
 import torch.nn as nn
 
 from torch.distributed.distributed_c10d import _rank_not_in_group
@@ -256,12 +257,14 @@
             loss.backward()
 
         if sharding_strategy_mode == ShardingStrategyMode.ALL_HYBRID_SHARD:
-            num_flat_params = len(list(FSDP._fsdp_handles(fsdp_model)))
+            num_flat_params = len(list(traversal_utils._get_fsdp_handles(fsdp_model)))
             self.assertEqual(num_flat_params, cntr[orig_ar])
             self.assertEqual(num_flat_params, cntr[orig_rs])
         elif sharding_strategy_mode == ShardingStrategyMode.MIXED_HYBRID_FULL_SHARD:
-            num_hsdp_flat_params = len(list(FSDP._fsdp_handles(fsdp_model.transformer)))
-            num_flat_params = len(list(FSDP._fsdp_handles(fsdp_model)))
+            num_hsdp_flat_params = len(
+                list(traversal_utils._get_fsdp_handles(fsdp_model.transformer))
+            )
+            num_flat_params = len(list(traversal_utils._get_fsdp_handles(fsdp_model)))
             self.assertEqual(num_hsdp_flat_params, cntr[orig_ar])
             self.assertEqual(num_flat_params, cntr[orig_rs])
 
diff --git a/test/distributed/fsdp/test_fsdp_overlap.py b/test/distributed/fsdp/test_fsdp_overlap.py
index 8bd5354..1eb3f7a 100644
--- a/test/distributed/fsdp/test_fsdp_overlap.py
+++ b/test/distributed/fsdp/test_fsdp_overlap.py
@@ -93,9 +93,9 @@
     def _dist_train(self):
         rank = self.rank
         world_size = self.world_size
-        # Save the original torch.distributed._all_gather_base function since we will
+        # Save the original torch.distributed.all_gather_into_tensor function since we will
         # patch it to include an artificial delay.
-        orig_all_gather = torch.distributed._all_gather_base
+        orig_all_gather = torch.distributed.all_gather_into_tensor
 
         def run(compute_cycles, all_gather_cycles):
             has_params = all_gather_cycles > 0
@@ -139,7 +139,7 @@
                 # Even though both e1 & e2 are on the compute stream, since
                 # compute depends on all_gather, e2-e1 includes all_gather time.
                 e1.record()
-                with patch("torch.distributed._all_gather_base", _delayed_all_gather):
+                with patch("torch.distributed.all_gather_into_tensor", _delayed_all_gather):
                     out = model(batch)
                     if has_params and world_size > 1:
                         self.assertTrue(all_gather_called)
diff --git a/test/distributed/fsdp/test_utils.py b/test/distributed/fsdp/test_utils.py
index b5ad65a..b9f9677 100644
--- a/test/distributed/fsdp/test_utils.py
+++ b/test/distributed/fsdp/test_utils.py
@@ -160,20 +160,20 @@
         def forward(self, x: torch.Tensor) -> torch.Tensor:
             return self.seq2(self.lin(self.seq1(x)))  # equivalent to one matmul
 
-    def test_module_wrap_policy(self):
+    def test_get_fully_sharded_module_to_states(self):
         """
-        Tests the module wrap policy on a nested model with buffers and a
-        shared parameter.
+        Tests the helper function ``_get_fully_sharded_module_states()`` that
+        performs the pseudo-auto-wrapping for the non-wrapper path.
 
         NOTE: This test is hard coded against ``Model``.
         """
         model = self.Model(TestGetSubmoduleToStates.SharedParameterMode.PARENT_CHILD)
 
-        # Compute the mapping from submodule to states according to a logical
-        # module wrap policy
+        # Compute the mapping from fully sharded module to states according to
+        # a logical module wrap policy
         module_classes = (nn.Sequential,)
         auto_wrap_policy = ModuleWrapPolicy(set(module_classes))
-        submodule_to_states = _get_fully_sharded_module_to_states(
+        fully_sharded_module_to_states = _get_fully_sharded_module_to_states(
             model, auto_wrap_policy, set(), set()
         )
         # Check the number of submodules with states in the mapping
@@ -183,41 +183,46 @@
         if not isinstance(model, module_classes):
             num_submodules_with_states += 1  # always include the root
         assert num_submodules_with_states == 4, f"{num_submodules_with_states}"
-        self.assertEqual(len(submodule_to_states), num_submodules_with_states)
+        self.assertEqual(len(fully_sharded_module_to_states), num_submodules_with_states)
 
-        # Check the mapping, i.e. that the dict order follows a post-order
-        # traversal and that the contents are expected
-        submodules = list(submodule_to_states.keys())
+        # Check the mapping, i.e. that the dict order follows `model.modules()`
+        # order and that the contents are expected
+        fully_sharded_modules = list(fully_sharded_module_to_states.keys())
+        expected_fully_sharded_modules = [
+            module for module in model.modules()
+            if isinstance(module, nn.Sequential) or module is model
+        ]
+        self.assertEqual(expected_fully_sharded_modules, fully_sharded_modules)
         # - Root module `model`
-        self.assertEqual(submodules[0], model)
-        root_states = submodule_to_states[submodules[0]]
+        self.assertEqual(fully_sharded_modules[0], model)
+        root_states = fully_sharded_module_to_states[fully_sharded_modules[0]]
         self.assertEqual(root_states.params, [model.lin.weight])
         self.assertEqual(root_states.param_names, ["lin.weight"])
         self.assertEqual(root_states.buffers, [])
         self.assertEqual(root_states.buffer_names, [])
-        # # - `seq2`
-        self.assertEqual(submodules[1], model.seq2)
-        seq2_states = submodule_to_states[submodules[1]]
-        self.assertEqual(seq2_states.params, [model.seq2[1].weight])
-        self.assertEqual(seq2_states.param_names, ["1.weight"])
-        self.assertEqual(seq2_states.buffers, [model.seq2[1].seq2_1_buffer])
-        self.assertEqual(seq2_states.buffer_names, ["1.seq2_1_buffer"])
-        # - `seq2[0]`
-        self.assertEqual(submodules[2], model.seq2[0])
-        seq2_0_states = submodule_to_states[submodules[2]]
-        self.assertEqual(seq2_0_states.params, [])  # shared parameter
-        self.assertEqual(seq2_0_states.param_names, [])
-        self.assertEqual(seq2_0_states.buffers, [])
-        self.assertEqual(seq2_0_states.buffer_names, [])
         # - `seq1`
-        self.assertEqual(submodules[3], model.seq1)
-        seq1_states = submodule_to_states[submodules[3]]
+        self.assertEqual(fully_sharded_modules[1], model.seq1)
+        seq1_states = fully_sharded_module_to_states[fully_sharded_modules[1]]
         self.assertEqual(
             seq1_states.params, [model.seq1[0].weight, model.seq1[1].weight]
         )
         self.assertEqual(seq1_states.param_names, ["0.weight", "1.weight"])
         self.assertEqual(seq1_states.buffers, [model.seq1.seq1_buffer])
         self.assertEqual(seq1_states.buffer_names, ["seq1_buffer"])
+        # - `seq2`
+        self.assertEqual(fully_sharded_modules[2], model.seq2)
+        seq2_states = fully_sharded_module_to_states[fully_sharded_modules[2]]
+        self.assertEqual(seq2_states.params, [model.seq2[1].weight])
+        self.assertEqual(seq2_states.param_names, ["1.weight"])
+        self.assertEqual(seq2_states.buffers, [model.seq2[1].seq2_1_buffer])
+        self.assertEqual(seq2_states.buffer_names, ["1.seq2_1_buffer"])
+        # - `seq2[0]`
+        self.assertEqual(fully_sharded_modules[3], model.seq2[0])
+        seq2_0_states = fully_sharded_module_to_states[fully_sharded_modules[3]]
+        self.assertEqual(seq2_0_states.params, [])  # shared parameter
+        self.assertEqual(seq2_0_states.param_names, [])
+        self.assertEqual(seq2_0_states.buffers, [])
+        self.assertEqual(seq2_0_states.buffer_names, [])
 
 
 instantiate_parametrized_tests(TestUtils)
diff --git a/torch/distributed/_composable/fully_shard.py b/torch/distributed/_composable/fully_shard.py
index f71b8ee..5065761 100644
--- a/torch/distributed/_composable/fully_shard.py
+++ b/torch/distributed/_composable/fully_shard.py
@@ -6,6 +6,7 @@
 from torch.distributed._composable.contract import contract
 from torch.distributed._composable_state import _get_module_state, _insert_module_state
 from torch.distributed.fsdp._common_utils import _FSDPState
+
 from torch.distributed.fsdp._init_utils import (
     _init_buffer_state,
     _init_core_state,
diff --git a/torch/distributed/fsdp/_common_utils.py b/torch/distributed/fsdp/_common_utils.py
index 322c1fe..52125de 100644
--- a/torch/distributed/fsdp/_common_utils.py
+++ b/torch/distributed/fsdp/_common_utils.py
@@ -29,24 +29,6 @@
 FSDP_PREFIX = FSDP_WRAPPED_MODULE + "."
 FSDP_FLATTENED = "_fsdp_flattened"
 
-"""
-[Note: FSDP State Traversal]
-For the wrapper code path, ``_FSDPState`` is the ``FullyShardedDataParallel``
-module wrapping a fully sharded module, and for the non-wrapper code path,
-``_FSDPState`` is an object that gets embedded on a fully sharded module.
-See [Note: Fully Sharded Module] for the definition.
-
-There are three common traversal idioms: Given a root module,
-- ``_get_fsdp_states()`` returns all ``_FSDPState`` s in the tree.
-- ``get_fsdp_root_states()`` returns all local root ``_FSDPState`` s in the
-tree (i.e. those with ``_is_root == True``).
-- ``_get_fsdp_handles()``returns all ``FlatParamHandle`` s in the tree.
-
-All of these methods must take in the root module (i.e. an ``nn.Module``) and
-not a general ``_FSDPState`` because ``_FSDPState`` does not support a graph
-traversal, whereas ``nn.Module`` has ``nn.Module.modules()`` for traversal.
-"""
-
 
 class _FSDPState(_State):
     def __init__(self) -> None:
@@ -58,6 +40,7 @@
         self._state_dict_config: StateDictConfig = FullStateDictConfig()
         self._is_root: Optional[bool] = None
         self._handles: List[flat_param_file.FlatParamHandle] = []
+        self._ignored_modules: Set[nn.Module] = set()
         self.rank: int = -1
 
 
@@ -68,40 +51,6 @@
     return state
 
 
-def _get_fsdp_states(module: nn.Module) -> List[_FSDPState]:
-    """
-    Returns all ``_FSDPState`` instances in the module tree rooted at
-    ``module`` without any duplicates and following the ``module.modules()``
-    traversal order.
-
-    For the wrapper code path, this returns all ``FullyShardedDataParallel``
-    instances. For the non-wrapper code path, this returns composable state
-    instances.
-
-    NOTE: For now, we must pass an ``nn.Module`` as the argument because
-    ``_FSDPState`` does not support graph traversal.
-    """
-    fsdp_states: List[_FSDPState] = []
-    visited_fsdp_states: Set[_FSDPState] = set()
-    for submodule in module.modules():
-        optional_state = _get_module_fsdp_state(submodule)
-        if optional_state is not None and optional_state not in visited_fsdp_states:
-            visited_fsdp_states.add(optional_state)
-            fsdp_states.append(optional_state)
-    return fsdp_states
-
-
-def _get_fsdp_handles(module: nn.Module) -> List:
-    """
-    Returns all ``FlatParamHandle`` s in the module tree rooted at ``module``.
-    """
-    return [
-        handle
-        for fsdp_state in _get_fsdp_states(module)
-        for handle in fsdp_state._handles
-    ]
-
-
 class TrainingState(Enum):
     """
     An enum that indicates the state of a ``FullyShardedDataParallel` instance.
diff --git a/torch/distributed/fsdp/_exec_order_utils.py b/torch/distributed/fsdp/_exec_order_utils.py
index 73b63d4..dd8a2b2 100644
--- a/torch/distributed/fsdp/_exec_order_utils.py
+++ b/torch/distributed/fsdp/_exec_order_utils.py
@@ -5,12 +5,9 @@
 
 import torch
 import torch.distributed as dist
+import torch.distributed.fsdp._traversal_utils as traversal_utils
 import torch.nn as nn
-from torch.distributed.fsdp._common_utils import (
-    _FSDPState,
-    _get_fsdp_handles,
-    _get_param_to_fqns,
-)
+from torch.distributed.fsdp._common_utils import _FSDPState, _get_param_to_fqns
 from torch.distributed.fsdp.flat_param import FlatParamHandle
 
 _HandlesKey = Tuple[FlatParamHandle, ...]
@@ -87,7 +84,7 @@
         self.rank = process_group.rank()
         self.world_size = process_group.size()
         # Fix an order over the handles, which should be the same across ranks
-        for handle in _get_fsdp_handles(root_module):
+        for handle in traversal_utils._get_fsdp_handles(root_module):
             index = len(self.all_handles)
             self.all_handles.append(handle)
             self.handle_to_handle_index[handle] = index
diff --git a/torch/distributed/fsdp/_init_utils.py b/torch/distributed/fsdp/_init_utils.py
index 83aea59..2d0107d 100644
--- a/torch/distributed/fsdp/_init_utils.py
+++ b/torch/distributed/fsdp/_init_utils.py
@@ -18,18 +18,20 @@
 
 import torch
 import torch.distributed as dist
+import torch.distributed.fsdp._exec_order_utils as exec_order_utils
+import torch.distributed.fsdp._traversal_utils as traversal_utils
 import torch.distributed.fsdp.fully_sharded_data_parallel as fsdp_file
 import torch.nn as nn
 from torch.distributed.algorithms._comm_hooks import default_hooks
 from torch.distributed.distributed_c10d import _get_default_group
 from torch.distributed.fsdp._common_utils import (
     _FSDPState,
+    _get_module_fsdp_state,
     _get_param_to_fqns,
     _is_fsdp_flattened,
     clean_tensor_name,
     TrainingState,
 )
-from torch.distributed.fsdp._exec_order_utils import _ExecOrderData
 from torch.distributed.fsdp._limiter_utils import _FreeEventQueue
 from torch.distributed.fsdp._wrap_utils import _get_fully_sharded_module_to_states
 from torch.distributed.fsdp.api import (
@@ -303,7 +305,7 @@
     state._stream_to_name = _stream_to_name
     state._free_event_queue = _FreeEventQueue()
     state._debug_level = dist.get_debug_level()
-    state._exec_order_data = _ExecOrderData(
+    state._exec_order_data = exec_order_utils._ExecOrderData(
         state._debug_level,
         backward_prefetch_limit,
         forward_prefetch_limit,
@@ -438,8 +440,13 @@
     )
     _check_single_device_module(root_module, state._ignored_params)
     device_from_device_id = _get_device_from_device_id(device_id, state.rank)
-    # Initialize and shard `FlatParamHandle`s one by one following bottom-up
-    # order (hence the `reversed`) to avoid increasing peak GPU memory usage
+    # Initialize and shard `FlatParamHandle`s one by one following reverse
+    # depth-first order (i.e. reverse `.modules()` order), which represents a
+    # reverse topological sort order. This avoids increasing peak GPU memory
+    # usage when the unsharded model exists on CPU or meta device.
+    # NOTE: This order differs from that followed by the wrapper path when
+    # using auto wrapping, which also represents a valid reverse toplogical
+    # sort order, but the difference does not matter.
     materialized_module = False
     for fully_sharded_module, (params, buffers, param_names, buffer_names) in reversed(
         fully_sharded_module_to_states.items()
@@ -475,7 +482,9 @@
         # Pass `root_module` to have internal FQN metadata prefix starting from
         # it instead of `submodule`
         _init_param_handle_from_params(state, params, fully_sharded_module)
-    # Reverse to preserve top-down order like `_fsdp_handles()`
+    # Reverse `_handles` to preserve depth-first `.modules()` order for
+    # consistency with the wrapper path (namely, so that `_get_fsdp_handles()`
+    # returns the same ordering for both paths).
     state._handles.reverse()
     return state
 
@@ -527,19 +536,31 @@
     any FSDP instances, and returns the modules contained in their module
     subtrees as a :class:`set`. Nested FSDP instances are excluded, but their
     already-computed ignored modules are included.
+
+    ``_ignored_modules`` represents the argument passed by the user to FSDP.
     """
-    if _ignored_modules is None:
-        return set()
     msg_prefix = "`ignored_modules` should be an iterable of `torch.nn.Module`s "
     try:
-        ignored_root_modules = set(_ignored_modules)
+        ignored_root_modules = (
+            set(_ignored_modules) if _ignored_modules is not None else set()
+        )
     except TypeError as e:
         raise TypeError(msg_prefix + f"but got {type(_ignored_modules)}") from e
     for module in ignored_root_modules:
         if not isinstance(module, torch.nn.Module):
             raise TypeError(msg_prefix + f"but got an iterable with {type(module)}")
         if isinstance(module, fsdp_file.FullyShardedDataParallel):
+            # TODO: We may relax this by taking the FSDP instance's wrapped
+            # module to provide more flexibility to the user.
             raise ValueError("`ignored_modules` should not include FSDP modules")
+    # Treat modules that cannot compose with `fully_shard` as ignored modules,
+    # meaning that their subtrees are ignored
+    for module in root_module.modules():
+        if not traversal_utils._composable(module):
+            ignored_root_modules.add(module)
+    # NOTE: Even if `ignored_root_modules` is empty, do not return early so
+    # that this FSDP instance can get any ignored modules from its children.
+
     # Include child modules and exclude nested FSDP modules themselves
     ignored_modules = set(
         child
@@ -555,9 +576,10 @@
         )
     # Include nested FSDP modules' ignored modules
     for submodule in root_module.modules():
-        if isinstance(submodule, fsdp_file.FullyShardedDataParallel):
-            assert hasattr(submodule, "_ignored_modules")
-            ignored_modules.update(submodule._ignored_modules)
+        optional_fsdp_state = _get_module_fsdp_state(submodule)
+        if optional_fsdp_state is not None:
+            assert hasattr(optional_fsdp_state, "_ignored_modules")
+            ignored_modules.update(optional_fsdp_state._ignored_modules)
     return ignored_modules
 
 
diff --git a/torch/distributed/fsdp/_optim_utils.py b/torch/distributed/fsdp/_optim_utils.py
index 2b6e191..6b2c635 100644
--- a/torch/distributed/fsdp/_optim_utils.py
+++ b/torch/distributed/fsdp/_optim_utils.py
@@ -17,6 +17,7 @@
 
 import torch
 import torch.distributed as dist
+import torch.distributed.fsdp._traversal_utils as traversal_utils
 
 # Import the entire FSDP file to avoid circular imports
 import torch.distributed.fsdp.fully_sharded_data_parallel as fsdp_file
@@ -131,7 +132,7 @@
         otherwise. The final optimizer state dict will need to map these
         entries using the proper unflattened parameter IDs.
     """
-    _clear_grads_if_needed(fsdp_module._fsdp_handles(fsdp_module))
+    _clear_grads_if_needed(traversal_utils._get_fsdp_handles(fsdp_module))
     consolidated_state = _communicate_optim_state(
         flat_param,
         flat_param_state,
diff --git a/torch/distributed/fsdp/_runtime_utils.py b/torch/distributed/fsdp/_runtime_utils.py
index bac37a8..b66e4aa 100644
--- a/torch/distributed/fsdp/_runtime_utils.py
+++ b/torch/distributed/fsdp/_runtime_utils.py
@@ -12,6 +12,7 @@
 )
 
 import torch
+import torch.distributed.fsdp._traversal_utils as traversal_utils
 import torch.nn as nn
 import torch.nn.functional as F
 from torch.autograd import Variable
@@ -19,8 +20,6 @@
 from torch.distributed.fsdp._common_utils import (
     _assert_in_training_states,
     _FSDPState,
-    _get_fsdp_handles,
-    _get_fsdp_states,
     _get_module_fsdp_state,
     _get_sharding_strategy,
     _is_composable,
@@ -108,7 +107,7 @@
     intra_node_pgs = set()
     inter_node_pgs = set()
     inter_node_states = set()
-    for fsdp_module in _get_fsdp_states(root_module):
+    for fsdp_module in traversal_utils._get_fsdp_states(root_module):
         # TODO: Change this to handle's sharding strategy if we deprecate
         # `ShardingStrategy` internally.
         # https://github.com/pytorch/pytorch/issues/90857
@@ -181,7 +180,7 @@
     attr_name_to_values: Dict[str, Set[Any]] = {}
     for attr_name in HOMOGENEOUS_ATTR_NAMES:
         attr_name_to_values[attr_name] = set()
-    for fsdp_state in _get_fsdp_states(root_module):
+    for fsdp_state in traversal_utils._get_fsdp_states(root_module):
         for attr_name in HOMOGENEOUS_ATTR_NAMES:
             p_assert(
                 hasattr(fsdp_state, attr_name),
@@ -491,7 +490,7 @@
             # TODO: This assumes singleton handles keys.
             handles_keys = [tuple(handle) for handle in state._handles]
         else:
-            for fsdp_module in _get_fsdp_states(state):
+            for fsdp_module in traversal_utils._get_fsdp_states(state):
                 handles_key = tuple(fsdp_module._handles)
                 handles_keys.append(handles_key)
         for handles_key in handles_keys:
@@ -501,7 +500,7 @@
         state._streams["unshard"],
         state._streams["pre_unshard"],
     )
-    _clear_grads_if_needed(_get_fsdp_handles(module))
+    _clear_grads_if_needed(traversal_utils._get_fsdp_handles(module))
 
     # Prepares the forward inputs by moving them to ``compute_device``
     # TODO: Do not use the side stream for tensor copies for now; investigate
@@ -573,7 +572,7 @@
         # after all backward calls complete
         if state._is_root and not state._post_backward_callback_queued:
             _register_post_backward_final_callback(state, module)
-            _clear_grads_if_needed(_get_fsdp_handles(module))
+            _clear_grads_if_needed(traversal_utils._get_fsdp_handles(module))
         elif _handles_key:
             allowed_states = [TrainingState.IDLE]
             if _is_composable(state):
@@ -868,7 +867,7 @@
             torch.cuda.current_stream().synchronize()
     root_state._exec_order_data.next_iter()
 
-    for fsdp_state in _get_fsdp_states(module):
+    for fsdp_state in traversal_utils._get_fsdp_states(module):
         _catch_all_reshard(fsdp_state)
         _finalize_params(fsdp_state)
         fsdp_state._ran_pre_backward_hook.clear()
@@ -1280,7 +1279,7 @@
         visited_buffers = set()
         # Traverse the FSDP instances bottom-up so that we prefer the owning
         # FSDP instance's mixed precision setting for each buffer
-        for fsdp_module in reversed(_get_fsdp_states(root_module)):
+        for fsdp_module in reversed(traversal_utils._get_fsdp_states(root_module)):
             for buffer in fsdp_module.buffers():
                 if buffer in visited_buffers:
                     continue
diff --git a/torch/distributed/fsdp/_state_dict_utils.py b/torch/distributed/fsdp/_state_dict_utils.py
index ac9d440..81a8ea4 100644
--- a/torch/distributed/fsdp/_state_dict_utils.py
+++ b/torch/distributed/fsdp/_state_dict_utils.py
@@ -6,6 +6,7 @@
 import torch
 import torch.distributed as dist
 import torch.distributed.algorithms._checkpoint.checkpoint_wrapper as checkpoint_wrapper
+import torch.distributed.fsdp._traversal_utils as traversal_utils
 
 # Import the entire FSDP file to avoid circular imports
 import torch.nn as nn
@@ -18,7 +19,6 @@
 )
 from torch.distributed.fsdp._common_utils import (
     _FSDPState,
-    _get_fsdp_handles,
     _has_fsdp_params,
     _is_composable,
     _module_handles,
@@ -130,7 +130,7 @@
     _lazy_init(fsdp_state, module)
     # TODO: change to this call after pre_state_dict_hook is in `nn.Module`.
     if fsdp_state._is_root:
-        _clear_grads_if_needed(_get_fsdp_handles(module))
+        _clear_grads_if_needed(traversal_utils._get_fsdp_handles(module))
 
 
 def _common_unshard_pre_state_dict_hook(
diff --git a/torch/distributed/fsdp/_traversal_utils.py b/torch/distributed/fsdp/_traversal_utils.py
new file mode 100644
index 0000000..8607323
--- /dev/null
+++ b/torch/distributed/fsdp/_traversal_utils.py
@@ -0,0 +1,94 @@
+"""
+NOTE: This file must be imported like
+``import torch.distributed.fsdp._traversal_utils`` and not like
+``from torch.distirbuted.fsdp._traversal_utils import ...`` to avoid circular
+imports. For brevity, we may import the file as ``traversal_utils``.
+"""
+
+import collections
+from typing import Deque, List, Set
+import torch.nn as nn
+from torch.distributed._composable.contract import _get_registry
+from torch.distributed.fsdp._common_utils import _FSDPState, _get_module_fsdp_state
+
+
+"""
+[Note: FSDP State Traversal]
+For the wrapper code path, ``_FSDPState`` is the ``FullyShardedDataParallel``
+module wrapping a fully sharded module, and for the non-wrapper code path,
+``_FSDPState`` is an object that gets embedded on a fully sharded module.
+See [Note: Fully Sharded Module] for the definition.
+
+There are three common traversal idioms: Given a root module,
+- ``_get_fsdp_states()`` returns all ``_FSDPState`` s in the tree.
+- ``get_fsdp_root_states()`` returns all local root ``_FSDPState`` s in the
+tree (i.e. those with ``_is_root == True``).
+- ``_get_fsdp_handles()``returns all ``FlatParamHandle`` s in the tree.
+
+All of these methods must take in the root module (i.e. an ``nn.Module``) and
+not a general ``_FSDPState`` because ``_FSDPState`` does not support a graph
+traversal, whereas ``nn.Module`` has ``nn.Module.modules()`` for traversal.
+"""
+
+
+def _composable(module: nn.Module) -> bool:
+    """
+    Returns if ``module`` can compose with ``fully_shard``.
+    """
+    # TODO: Add any other composable APIs that are mutually exclusive.
+    return "replicate" not in _get_registry(module)
+
+
+def _get_fsdp_states(module: nn.Module) -> List[_FSDPState]:
+    """
+    Returns all ``_FSDPState`` instances in the module tree rooted at
+    ``module`` without any duplicates and following the ``module.modules()``
+    traversal order (which is assumed to remain as depth-first). However, the
+    traversal does not proceed into any module annotated by an incompatible
+    API (e.g. ``replicate``).
+
+    For the wrapper code path, this returns all ``FullyShardedDataParallel``
+    instances. For the non-wrapper code path, this returns composable state
+    instances.
+
+    NOTE: For now, we must pass an ``nn.Module`` as the argument because
+    ``_FSDPState`` does not support graph traversal.
+    """
+    fsdp_states: List[_FSDPState] = []
+    # Track the visited FSDP states since multiple modules may share the same
+    # one and we want to return a de-duplicated list
+    visited_fsdp_states: Set[_FSDPState] = set()
+    # Track the visited modules in case of shared modules, which implies the
+    # module graph is no longer a tree
+    visited_modules: Set[nn.Module] = set()
+
+    # Perform depth-first search from `module` to ensure that we do not
+    # traverse into an incompatible API's subtree (use DFS instead of BFS to
+    # match `.modules()` order)
+    deque: Deque[nn.Module] = collections.deque()
+    deque.append(module)
+    while deque:
+        submodule = deque.popleft()
+        visited_modules.add(submodule)
+        if not _composable(submodule):
+            continue
+        for child_module in reversed(list(submodule.children())):
+            if child_module not in visited_modules:
+                deque.appendleft(child_module)
+        optional_state = _get_module_fsdp_state(submodule)
+        if optional_state is not None and optional_state not in visited_fsdp_states:
+            visited_fsdp_states.add(optional_state)
+            fsdp_states.append(optional_state)
+    return fsdp_states
+
+
+def _get_fsdp_handles(module: nn.Module) -> List:
+    """
+    Returns all ``FlatParamHandle`` s in the module tree rooted at ``module``
+    following the rules in :func:`_get_fsdp_state`.
+    """
+    return [
+        handle
+        for fsdp_state in _get_fsdp_states(module)
+        for handle in fsdp_state._handles
+    ]
diff --git a/torch/distributed/fsdp/_wrap_utils.py b/torch/distributed/fsdp/_wrap_utils.py
index 9533187..f9b5f89 100644
--- a/torch/distributed/fsdp/_wrap_utils.py
+++ b/torch/distributed/fsdp/_wrap_utils.py
@@ -102,8 +102,8 @@
     instead.
     """
     # Record the modules to wrap without actually wrapping
-    wrapped_modules: List[nn.Module] = []  # these are only logically wrapped
-    wrapper_cls = functools.partial(_record_module_wrapper_cls, wrapped_modules)
+    wrapped_modules_set: Set[nn.Module] = set()  # these are only logically wrapped
+    wrapper_cls = functools.partial(_record_module_wrapper_cls, wrapped_modules_set)
     if auto_wrap_policy is not None:
         _recursive_wrap(
             root_module,
@@ -114,33 +114,41 @@
             only_wrap_children=False,
         )
     # Always include the root module even if not wrapped by the given policy
-    if root_module not in wrapped_modules:
-        wrapped_modules.append(root_module)
+    wrapped_modules_set.add(root_module)
 
-    submodule_to_states = collections.OrderedDict()
+    fully_sharded_module_to_states = collections.OrderedDict()
     visited_params = set()
     for ignored_param in ignored_params:
         visited_params.add(ignored_param)
     visited_buffers = set()
-    # Constructing `wrapped_modules` with `_recursive_wrap()` follows a
-    # post-order traversal. We record state in `submodule_to_states` using a
-    # reverse post-ordering since that is a topological sort. This assigns
-    # parent-child shared parameters to the parent submodule.
-    # TODO: To handle sibling shared parameters, we need to pre-compute the
-    # shared parameters and assign them to the LCA submodule manually.
-    wrapped_modules.reverse()
-    wrapped_modules_set = set(wrapped_modules)
+    # Construct `wrapped_modules` to follow `.modules()` order to ensure that
+    # downstream data structures (`._handles`) match those of the wrapper path.
+    # NOTE: Since `.modules()` follows a depth-first order, which is a
+    # topological sort, and we iterate over `wrapped_modules` following that
+    # order, parent-child shared parameters are assigned to the parent module.
+    wrapped_modules: List[nn.Module] = []
+    for module in root_module.modules():
+        if module in wrapped_modules_set:
+            wrapped_modules.append(module)
     for submodule in wrapped_modules:
-        # Perform a BFS from `submodule` and record all unvisited state that is
-        # not already associated with another module in `wrapped_modules`.
-        queue: Deque[Tuple[nn.Module, str]] = collections.deque()
-        queue.append((submodule, ""))
+        # Perform a DFS from `submodule` and record all unvisited state that is
+        # not already associated with another module in `wrapped_modules`. We
+        # use DFS to follow the `.modules()` order.
+        deque: Deque[Tuple[nn.Module, str]] = collections.deque()
+        deque.append((submodule, ""))
         params: List[nn.Parameter] = []
         param_names: List[str] = []
         buffers: List[torch.Tensor] = []
         buffer_names: List[str] = []
-        while len(queue) > 0:
-            module, prefix = queue.popleft()
+        while len(deque) > 0:
+            module, prefix = deque.popleft()
+            # Reverse `named_children()`, use `appendleft()`, and add to the
+            # deque before processing to perform non-recursive DFS
+            for child_module_name, child_module in reversed(
+                list(module.named_children())
+            ):
+                if child_module not in wrapped_modules_set:
+                    deque.appendleft((child_module, prefix + child_module_name + "."))
             for param_name, param in module.named_parameters(recurse=False):
                 if param not in visited_params and not _is_fsdp_flattened(param):
                     params.append(param)
@@ -151,24 +159,21 @@
                     buffers.append(buffer)
                     visited_buffers.add(buffer)
                     buffer_names.append(prefix + buffer_name)
-            for child_module_name, child_module in module.named_children():
-                if child_module not in wrapped_modules_set:
-                    queue.append((child_module, prefix + child_module_name + "."))
-        submodule_to_states[submodule] = FullyShardedModuleState(
+        fully_sharded_module_to_states[submodule] = FullyShardedModuleState(
             params, buffers, param_names, buffer_names
         )
-    return submodule_to_states
+    return fully_sharded_module_to_states
 
 
 def _record_module_wrapper_cls(
-    wrapped_modules: List[nn.Module],
+    wrapped_modules_set: Set[nn.Module],
     module: nn.Module,
     **kwargs,
 ) -> nn.Module:
     """
     This defines a pseudo-wrapper class to be passed to ``_recursive_wrap()``
-    that records the wrapped module to the input ``wrapped_modules`` without
-    actually wrapping with a class.
+    that records the wrapped module to the input ``wrapped_modules_set``
+    without actually wrapping with a class.
     """
-    wrapped_modules.append(module)
+    wrapped_modules_set.add(module)
     return module
diff --git a/torch/distributed/fsdp/fully_sharded_data_parallel.py b/torch/distributed/fsdp/fully_sharded_data_parallel.py
index 586abc4..3a98ffb 100644
--- a/torch/distributed/fsdp/fully_sharded_data_parallel.py
+++ b/torch/distributed/fsdp/fully_sharded_data_parallel.py
@@ -21,6 +21,7 @@
 
 import torch
 import torch.distributed as dist
+import torch.distributed.fsdp._traversal_utils as traversal_utils
 import torch.nn as nn
 from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
     _CHECKPOINT_WRAPPED_MODULE,
@@ -29,8 +30,6 @@
 from torch.distributed.algorithms._comm_hooks import LOW_PRECISION_HOOKS
 from torch.distributed.fsdp._common_utils import (
     _FSDPState,
-    _get_fsdp_handles,
-    _get_fsdp_states,
     _get_param_to_fqns,
     FSDP_PREFIX,
     FSDP_WRAPPED_MODULE,
@@ -94,7 +93,7 @@
     _unshard_params,
 )
 from ._utils import p_assert
-from .flat_param import FlatParameter, FlatParamHandle
+from .flat_param import FlatParameter
 from .wrap import _FSDPPolicy
 
 
@@ -470,15 +469,7 @@
         """
         if root_only:
             return _get_fsdp_root_states(module)
-        return _get_fsdp_states(module)
-
-    @staticmethod
-    def _fsdp_handles(module: nn.Module) -> List[FlatParamHandle]:
-        """
-        Returns all nested FSDP instances' handles in the module hierarchy
-        rooted at ``module``.
-        """
-        return _get_fsdp_handles(module)
+        return traversal_utils._get_fsdp_states(module)
 
     def apply(self, fn: Callable[[nn.Module], None]) -> "FullyShardedDataParallel":
         r"""Applies ``fn`` recursively to every submodule (as returned by ``.children()``)
@@ -503,7 +494,7 @@
         # Reset lazy init that might be called by _summon_full_params, since
         # it could have set is_root incorrectly for non-root FSDP instances.
         if uninitialized and self._is_root:
-            for module in _get_fsdp_states(self):
+            for module in traversal_utils._get_fsdp_states(self):
                 module._reset_lazy_init()
 
         return ret
@@ -583,7 +574,7 @@
         # Use the default config if a state_dict config is not set.
         if state_dict_config is None:
             state_dict_config = _state_dict_type_to_config[state_dict_type]()
-        for submodule in _get_fsdp_states(module):
+        for submodule in traversal_utils._get_fsdp_states(module):
             if prev_state_dict_type is None:
                 prev_state_dict_type = submodule._state_dict_type
             if prev_state_dict_config is None:
@@ -682,9 +673,7 @@
                     f"{self.compute_device} but got {handle.flat_param.device}",
                 )
             output = self._fsdp_wrapped_module(*args, **kwargs)
-            return _post_forward(
-                self, self._handles, reshard_fn, self, unused, output
-            )
+            return _post_forward(self, self._handles, reshard_fn, self, unused, output)
 
     @staticmethod
     @contextlib.contextmanager
@@ -809,7 +798,7 @@
 
         if recurse:
             with contextlib.ExitStack() as stack:
-                for module in _get_fsdp_states(self):
+                for module in traversal_utils._get_fsdp_states(self):
                     stack.enter_context(
                         module._summon_full_params(
                             recurse=False,
@@ -850,12 +839,12 @@
             "`_deregister_orig_params_ctx()` should only be called when "
             "`_use_orig_params=True`",
         )
-        for fsdp_module in _get_fsdp_states(self):
+        for fsdp_module in traversal_utils._get_fsdp_states(self):
             _deregister_orig_params(fsdp_module, fsdp_module)
         try:
             yield
         finally:
-            for fsdp_module in _get_fsdp_states(self):
+            for fsdp_module in traversal_utils._get_fsdp_states(self):
                 _register_orig_params(fsdp_module, fsdp_module)
 
     def _apply(self, *args, **kwargs):
@@ -1024,7 +1013,7 @@
         # the normal `nn.utils` one targeting local gradients
         all_no_shard = all(
             not handle.uses_sharded_strategy
-            for handle in FullyShardedDataParallel._fsdp_handles(self)
+            for handle in traversal_utils._get_fsdp_handles(self)
         )
         if all_no_shard:
             return torch.nn.utils.clip_grad_norm_(
@@ -1037,7 +1026,7 @@
         sharded_params = set()
         nonsharded_params = set()  # `NO_SHARD` or not FSDP-managed
         grads: List[torch.Tensor] = []
-        for handle in FullyShardedDataParallel._fsdp_handles(self):
+        for handle in traversal_utils._get_fsdp_handles(self):
             target_set = (
                 sharded_params if handle.uses_sharded_strategy else nonsharded_params
             )
@@ -1121,7 +1110,10 @@
     ):
         if full_optim and not rank0_only:
             return
-        if any(fsdp_module._use_orig_params for fsdp_module in _get_fsdp_states(model)):
+        if any(
+            fsdp_module._use_orig_params
+            for fsdp_module in traversal_utils._get_fsdp_states(model)
+        ):
             raise NotImplementedError(
                 "Optimizer state checkpointing is not supported yet for `use_orig_params=True`"
             )
@@ -1194,7 +1186,7 @@
             optim,
         )
         use_orig_params: bool = False
-        for module in _get_fsdp_states(model):
+        for module in traversal_utils._get_fsdp_states(model):
             use_orig_params = module._use_orig_params
             break
         return _optim_state_dict(
@@ -1313,7 +1305,7 @@
             optim,
         )
         use_orig_params: bool = False
-        for module in _get_fsdp_states(model):
+        for module in traversal_utils._get_fsdp_states(model):
             use_orig_params = module._use_orig_params
             break
         sharded_osd = _flatten_optim_state_dict(
@@ -1737,7 +1729,7 @@
             raise AssertionError(
                 "register_comm_hook can only be called on a root instance."
             )
-        for submodule in _get_fsdp_states(self):
+        for submodule in traversal_utils._get_fsdp_states(self):
             assert (
                 not submodule._hook_registered
             ), "communication hook can be only registered once"
diff --git a/torch/testing/_internal/common_dist_composable.py b/torch/testing/_internal/common_dist_composable.py
index 42029f1..9ec92c6 100644
--- a/torch/testing/_internal/common_dist_composable.py
+++ b/torch/testing/_internal/common_dist_composable.py
@@ -1,5 +1,7 @@
 # Owner(s): ["oncall: distributed"]
 
+from typing import Tuple
+
 import torch
 import torch.nn as nn
 
@@ -58,3 +60,47 @@
         a = self.u2(self.u1(self.l(x)))
         b = self.p
         return torch.mm(a, b)
+
+
+class FakeSequential(nn.Module):
+    # Define this class to achieve a desired nested wrapping using the module
+    # wrap policy with `nn.Sequential`
+    def __init__(self, *modules: Tuple[nn.Module, ...]) -> None:
+        super().__init__()
+        self._module_sequence = list(modules)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for module in self._module_sequence:
+            x = module(x)
+        return x
+
+
+class NestedSequentialModel(nn.Module):
+    def __init__(self, device: torch.device) -> None:
+        super().__init__()
+        # This nested structure exercises traversal order to catch differences
+        # between valid traversals (e.g. BFS and DFS variations).
+        self.seq1 = nn.Sequential(
+            nn.Linear(1, 1, device=device),
+            FakeSequential(
+                nn.Linear(1, 1, device=device),
+                nn.ReLU(),
+                FakeSequential(
+                    nn.Linear(1, 1, device=device),
+                ),
+                nn.ReLU(),
+            ),
+            nn.Linear(1, 2, device=device),
+        )
+        self.lin = nn.Linear(2, 2, device=device)
+        self.seq2 = nn.Sequential(
+            nn.ReLU(),
+            nn.Linear(2, 3, device=device),
+            FakeSequential(
+                nn.Linear(3, 2, bias=False, device=device),
+                nn.Linear(2, 4, bias=False, device=device),
+            ),
+        )
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return self.seq2(self.lin(self.seq1(x)))