1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
|
import itertools
from typing import List, Optional, Set, Tuple, Union
import torch
import torch.distributed as dist
import torch.nn as nn
from torch.distributed.device_mesh import _get_device_handle
from torch.distributed.tensor import DeviceMesh, DTensor, init_device_mesh
from torch.utils._python_dispatch import is_traceable_wrapper_subclass
from ._fsdp_common import _is_composable_with_fsdp, FSDPMeshInfo, HSDPMeshInfo
from ._fsdp_state import _get_module_fsdp_state
def _get_post_forward_mesh_info(
reshard_after_forward: Union[bool, int], mesh_info: FSDPMeshInfo
) -> Optional[FSDPMeshInfo]:
shard_mesh_size = mesh_info.shard_mesh_size
if not isinstance(reshard_after_forward, (bool, int)):
raise ValueError(
"reshard_after_forward should be a bool or an int representing the "
f"group size to reshard to, not {reshard_after_forward}"
)
# NOTE: `isinstance(False, int)` returns `True`.
if not isinstance(reshard_after_forward, bool) and isinstance(
reshard_after_forward, int
):
if (
reshard_after_forward < 1
or reshard_after_forward > shard_mesh_size
or shard_mesh_size % reshard_after_forward != 0
):
raise ValueError(
"If passing reshard_after_forward as an int, it should be a "
f"factor of {shard_mesh_size}, not {reshard_after_forward}"
)
elif reshard_after_forward == 1:
reshard_after_forward = False
elif reshard_after_forward == shard_mesh_size:
reshard_after_forward = True
post_forward_mesh_info = None
if reshard_after_forward is True:
post_forward_mesh_info = mesh_info
elif reshard_after_forward is not False: # int case
# For HSDP, we can flatten the two replicate dims into the 0th dim
post_forward_mesh_tensor = mesh_info.mesh.mesh.view(-1, reshard_after_forward)
post_forward_mesh = DeviceMesh(
mesh_info.mesh.device_type, post_forward_mesh_tensor
)
post_forward_mesh_info = HSDPMeshInfo(
post_forward_mesh, shard_mesh_dim=1, replicate_mesh_dim=0
)
return post_forward_mesh_info
def _init_default_fully_shard_mesh() -> DeviceMesh:
"""Default to global CUDA mesh if possible else global CPU mesh."""
if not dist.distributed_c10d.is_initialized():
dist.distributed_c10d.init_process_group()
default_pg = dist.distributed_c10d._get_default_group()
device = torch._C._get_accelerator()
mesh = init_device_mesh(device.type, mesh_shape=(default_pg.size(),))
return mesh
def _get_device_from_mesh(mesh: DeviceMesh) -> torch.device:
if mesh.device_type == "cpu":
return torch.device("cpu")
device_handle = _get_device_handle(mesh.device_type)
return torch.device(mesh.device_type, device_handle.current_device())
def _get_managed_modules(root_modules: Tuple[nn.Module, ...]) -> List[nn.Module]:
modules: List[nn.Module] = []
root_modules_set = set(root_modules)
# Track visisted modules to avoid visiting shared modules multiple times
visited_modules: Set[nn.Module] = set()
def dfs(module: nn.Module) -> None:
"""
Runs a DFS to collect managed modules, not recursing into modules with
a non-composable API or ``fully_shard`` already applied.
"""
if not _is_composable_with_fsdp(module):
return
elif (
module not in root_modules_set
and _get_module_fsdp_state(module) is not None
):
return # nested `fully_shard` module
visited_modules.add(module)
for submodule in module.children():
if submodule not in visited_modules:
dfs(submodule)
modules.append(module)
for root_module in root_modules:
dfs(root_module)
return modules
def _verify_managed_param(name: str, param: nn.Parameter) -> None:
"""
Verify if the parameter is accepted by fully_shard. The only restriction now
is that the parameter cannot be a scalar tensor (param.numel == 0) since we
need at least one dim to shard.
"""
if len(param.shape) == 0:
raise ValueError(
"fully_shard doesn't support salar parameters. "
f"Change {name} to a 1D tensor with numel equal to 1."
)
def _get_managed_states(
modules: List[nn.Module],
) -> Tuple[List[nn.Parameter], List[torch.Tensor]]:
params: List[nn.Parameter] = []
buffers: List[torch.Tensor] = []
# Track visited parameters/buffers to avoid visiting shared parameters and
# buffers multiple times
visited_params: Set[nn.Parameter] = set()
visited_buffers: Set[torch.Tensor] = set()
for module in modules:
for name, param in module.named_parameters(recurse=False):
if param not in visited_params:
_verify_managed_param(name, param)
params.append(param)
visited_params.add(param)
for buffer in module.buffers(recurse=False):
if buffer not in visited_buffers:
buffers.append(buffer)
visited_buffers.add(buffer)
return params, buffers
def _move_states_to_device(
params: List[nn.Parameter],
buffers: List[torch.Tensor],
device: torch.device,
) -> None:
"""
We have FSDP move states to device for simpler and faster initialization
since FSDP almost always uses CUDA for training. We move parameters/buffers
rather than modules since modules to support ignoring parameters/buffers in
the future.
"""
# Follow the logic in `nn.Module._apply`
for tensor in itertools.chain(params, buffers):
if tensor.device == device or tensor.device.type == "meta":
# Keep meta-device tensors on meta device for deferred init
continue
if isinstance(tensor, DTensor):
if (dtensor_mesh_type := tensor.device_mesh.device_type) != device.type:
raise ValueError(
"Requires DTensor to have mesh of the same type as the FSDP mesh "
f"but got {dtensor_mesh_type} for DTensor and {device.type} for FSDP"
)
raise AssertionError(
f"Expects DTensor to be moved to {dtensor_mesh_type} but got {tensor.device}"
)
tensor_ = tensor
if is_traceable_wrapper_subclass(tensor_):
with torch.no_grad(): # avoid autograd increasing C++ refcount by 1
tensor_on_device = nn.Parameter(tensor.to(device))
torch.utils.swap_tensors(tensor, tensor_on_device)
else:
tensor.data = tensor.to(device)
|
