Skip to content

[Full DTensor][Reland] Add full_dtensor flag #2013

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Nov 11, 2025
Merged

[Full DTensor][Reland] Add full_dtensor flag #2013

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2592021
Update (base update)
fegin Nov 11, 2025
7b6ff5e
Update
fegin Nov 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion torchtitan/experiments/simple_fsdp/simple_fsdp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -185,6 +185,7 @@ def __init__(
mode: str,
mp_policy: MixedPrecisionPolicy | None,
reduction_divide_factor: float | None,
full_dtensor: bool = False,
) -> None:
super().__init__()
self.device_mesh = device_mesh
Expand All @@ -201,6 +202,7 @@ def __init__(
mp_policy = mp_policy or MixedPrecisionPolicy()
self.param_dtype: torch.dtype | None = mp_policy.param_dtype
self.reduce_dtype: torch.dtype | None = mp_policy.reduce_dtype
self.full_dtensor = full_dtensor

def replicate_compute(self, x: DTensor) -> torch.Tensor:
# data parallel runtime replicate parameters and do local compute
Expand All @@ -210,6 +212,10 @@ def replicate_compute(self, x: DTensor) -> torch.Tensor:
non_dp_mesh_dims = x._spec.mesh.ndim - self.device_mesh.ndim
assert non_dp_mesh_dims <= 2, "Only DP + EP/TP/EP+TP is supported"
if non_dp_mesh_dims > 0:
if self.full_dtensor:
raise NotImplementedError(
"full_dtensor not implemented for nD parallelisms"
)
dp_mesh = self.device_mesh
# re-wrap 2D DTensor to 1D DTensor on dp_mesh for efficient FSDP all-gather
sharded_local_tensor = x.to_local()
Expand Down Expand Up @@ -245,7 +251,10 @@ def replicate_compute(self, x: DTensor) -> torch.Tensor:
placements=self.compute_placements,
forward_dtype=self.param_dtype,
backward_dtype=self.reduce_dtype,
).to_local(grad_placements=self.grad_placements)
)

if not self.full_dtensor:
output = output.to_local(grad_placements=self.grad_placements)
else:
raise AssertionError(
f"Unsupported replicate compute on placement {x._spec.placements} for DTensor {x}"
Expand Down Expand Up @@ -274,6 +283,7 @@ def data_parallel(
mp_policy: MixedPrecisionPolicy | None = None,
shard_dim: int = 0,
reduction_divide_factor: float | None = None,
full_dtensor: bool = False,
) -> nn.Module:
param_sharding: tuple[Placement, ...]
if mode == "replicate":
Expand Down Expand Up @@ -333,6 +343,7 @@ def data_parallel(
mode,
mp_policy=mp_policy,
reduction_divide_factor=reduction_divide_factor,
full_dtensor=full_dtensor,
),
)
return model