[FSDP][7/N] Support `replicate` in `fully_shard` (#91044)
This PR supports nesting `replicate` in `fully_shard`.
- The PR achieves this by treating `replicate`-annotated modules are ignored modules. This means that all submodules in the `replicate`-annotated module's subtree are ignored, including nested `fully_shard`-annotated modules, which is the desired behavior.
---
This PR reworks some tree traversal.
One end goal is for `state._handles` to follow the same order for both the wrapper and composable paths. This implies that `_get_fsdp_handles()` returns the same value for both paths.
- The helper function `_get_fully_sharded_module_to_states()` now follows a left-to-right DFS from each fully sharded module instead of a BFS. The left-to-right DFS follows `.modules()` order.
- The composable auto "wrap" initialization function `_init_param_handles_from_module()` follows the reverse left-to-right DFS order. As noted in the code comments, this initialization order is a valid reverse topological sort, but it differs from the wrapper path. This is the _only_ difference with respect to initialization order through the entire process.
```
mod: Module(
submod1: Submodule()
submod2: Submodule(
subsubmod: Subsubmodule(),
),
)
```
For left-to-right DFS, the order is `mod`, `submod1`, `submod2`, `subsubmod`. (For context, right-to-left DFS would be `mod`, `submod2`, `subsubmod`, `submod1`. In other words, the left-to-right vs. right-to-left corresponds to `.children()` vs. `reversed(.children())` respectively.) Then, reverse left-to-right DFS is `subsubmod`, `submod2`, `submod1`, `mod`, which is a valid initialization order. However, the wrapper auto wrap initialization order would be `submod1`, `subsubmod`, `submod2`, `mod` since it directly follows a left-to-right DFS and initializes as a part of the recursive DFS logic.
- At the end of `_init_param_handles_from_module()`, we reverse the newly populated `state._handles`, so this is the reverse reverse left-to-right DFS order, which is equivalent to the left-to-right DFS order. Thus, `state._handles` has the same order for both paths.
Another goal is for `_get_fsdp_states()` to not traverse into any submodule that is annotated with an API that is not compatible with `fully_shard` (e.g. `replicate`). To achieve this while preserving that `_get_fsdp_states()` follows `.modules()` order, we again use a left-to-right DFS.
The reason the DFSs may look strange is because I implemented them non-recursively, which requires a stack.
- `test_get_fully_sharded_module_to_states()` in `test_utils.py` checks the traversal order of `_get_fully_sharded_module_to_states()`.
- `test_policy()` in `test_fully_shard.py` checks the traversal order returned by `_get_fsdp_handles()`.
---
Due to a circular dependency issue, we must move the graph/tree traversal helpers to their own file `_traversal_utils.py`, and any usages must import the entire file like `import torch.distributed.fsdp._traversal_utils as traversal_utils` instead of `from torch.distributed.fsdp._traversal_utils import ...`.
The cycle comes from the fact that the traversals require `_composable()`, which requires `_get_registry()` from `composable/contract.py`, which when imported, imports `composable/fully_shard.py`, which requires the traversals.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/91044
Approved by: https://github.com/mrshenli
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)))
