From 844c6250248678da780d74ed791e9139bf4ce103 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Fri, 22 May 2026 06:41:27 +0000 Subject: [PATCH 01/23] WIP: add experimental minimal FSDP path Signed-off-by: Jingyue Wu --- .../megatron_fsdp/experimental/__init__.py | 13 +- .../megatron_fsdp/experimental/fully_shard.py | 556 ++++++++++++++++++ .../test_experimental_fully_shard.py | 319 ++++++++++ 3 files changed, 887 insertions(+), 1 deletion(-) create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py create mode 100644 tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py index 1bd55b7d995..1cb67e172da 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py @@ -15,6 +15,17 @@ """Experimental Megatron-FSDP implementation.""" from .dbuffer import DBuffer +from .fully_shard import FsdpModule, ParameterGroup, Placements, fully_shard from .placement import Flat, Partial, Placement, Replicate -__all__ = ["DBuffer", "Flat", "Partial", "Placement", "Replicate"] +__all__ = [ + "DBuffer", + "Flat", + "FsdpModule", + "ParameterGroup", + "Partial", + "Placement", + "Placements", + "Replicate", + "fully_shard", +] diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py new file mode 100644 index 00000000000..ba664e82387 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py @@ -0,0 +1,556 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Minimal experimental per-module Megatron-FSDP implementation.""" + +import dataclasses +from collections.abc import Callable, Sequence + +import torch +from torch import nn +from torch.distributed import DeviceMesh + +from ..mixed_precision import MixedPrecisionPolicy +from .dbuffer import DBuffer, MeshAxis, Partial, Placement, Replicate + +_CONTAINING_PARAMETER_GROUP_ATTR = "_mfsdp_parameter_group" + + +@dataclasses.dataclass(frozen=True) +class Placements: + """Per-mesh-axis placements for parameter, gradient, and optimizer buffers.""" + + dp_axes: list[MeshAxis] + parameter: list[Placement] + gradient: list[Placement] + optimizer: list[Placement] + + def __post_init__(self) -> None: + """Validate placement list lengths.""" + axis_count = len(self.dp_axes) + for name, placements in ( + ("parameter", self.parameter), + ("gradient", self.gradient), + ("optimizer", self.optimizer), + ): + if len(placements) != axis_count: + raise ValueError(f"Expected {axis_count} {name} placements, got {len(placements)}.") + + +class ParameterGroup: + """A dtype and requires-grad homogeneous group of FSDP-owned parameters.""" + + module: nn.Module + parameters: dict[str, nn.Parameter] + _unsharded_parameters: dict[str, nn.Parameter] + mesh: DeviceMesh + dtype: torch.dtype + requires_grad: bool + main_weight: DBuffer + model_weight: DBuffer + main_grad: DBuffer | None + _full_weight: DBuffer | None + _full_weight_allocated: bool + + def __init__( + self, + module: nn.Module, + parameters: dict[str, nn.Parameter], + mesh: DeviceMesh, + model_weight_placements: Sequence[Placement], + main_grad_placements: Sequence[Placement], + main_weight_placements: Sequence[Placement], + mixed_precision_policy: MixedPrecisionPolicy, + ) -> None: + """Create persistent sharded buffers for a group of parameters. + + Args: + module: Closest FSDP root module that owns this parameter group. + parameters: Root-module-relative FQNs and their parameters. + mesh: Device mesh used by the buffers. + model_weight_placements: Placements for compute-weight storage. + main_grad_placements: Placements for persistent main gradients. + main_weight_placements: Placements for optimizer-owned main weights. + mixed_precision_policy: Precision policy for main weights and gradients. + """ + if not parameters: + raise ValueError("ParameterGroup requires at least one parameter.") + + # Python dicts preserve insertion order, so values() defines the stable + # tensor order used by each DBuffer built from this group. + original_parameters = parameters + self.module = module + self.parameters = {} + self._unsharded_parameters = {} + self.mesh = mesh + first_parameter = next(iter(original_parameters.values())) + self.dtype = first_parameter.dtype + self.requires_grad = first_parameter.requires_grad + for name, parameter in original_parameters.items(): + if parameter.dtype != self.dtype: + raise ValueError( + f"Expected parameter {name!r} to have dtype {self.dtype}, got {parameter.dtype}." + ) + if parameter.requires_grad != self.requires_grad: + raise ValueError( + f"Expected parameter {name!r} to have requires_grad={self.requires_grad}, " + f"got {parameter.requires_grad}." + ) + main_params_dtype = mixed_precision_policy.main_params_dtype + if main_params_dtype is None: + raise ValueError( + "experimental FSDP requires main_params_dtype to be specified explicitly." + ) + main_grads_dtype = mixed_precision_policy.main_grads_dtype + if main_grads_dtype is None: + main_grads_dtype = self.dtype + self.main_weight = DBuffer( + mesh=self.mesh, + placements=main_weight_placements, + tensor_shapes=[parameter.shape for parameter in original_parameters.values()], + dtype=main_params_dtype, + device=_mesh_device(self.mesh), + ) + for index, parameter in enumerate(original_parameters.values()): + if parameter.is_meta: + continue + self._copy_full_tensor_to_buffer( + self.main_weight, + index, + parameter.detach().to( + dtype=main_params_dtype, device=self.main_weight.local_buffer.device + ), + ) + self.model_weight = self._make_model_weight(tuple(model_weight_placements)) + self.main_grad = ( + DBuffer( + mesh=self.mesh, + placements=main_grad_placements, + tensor_shapes=self.main_weight.layout.tensor_shapes, + dtype=main_grads_dtype, + device=self.main_weight.local_buffer.device, + ) + if self.requires_grad + else None + ) + self._full_weight: DBuffer | None = None + self._full_weight_allocated = False + + for index, (name, original_parameter) in enumerate(original_parameters.items()): + sharded_parameter = nn.Parameter( + self.main_weight.get_dtensor(index), requires_grad=original_parameter.requires_grad + ) + setattr(sharded_parameter, _CONTAINING_PARAMETER_GROUP_ATTR, self) + self.parameters[name] = sharded_parameter + self._set_module_parameter(name, sharded_parameter) + + def _make_model_weight(self, placements: Sequence[Placement]) -> DBuffer: + model_weight = self.main_weight.redistribute(placements) + if model_weight.local_buffer.dtype != self.dtype: + converted = DBuffer( + mesh=model_weight.mesh, + placements=model_weight.placements, + tensor_shapes=model_weight.layout.tensor_shapes, + dtype=self.dtype, + device=model_weight.local_buffer.device, + ) + converted.local_buffer.copy_(model_weight.local_buffer.to(dtype=self.dtype)) + return converted + return model_weight + + def _copy_full_tensor_to_buffer( + self, buffer: DBuffer, index: int, tensor: torch.Tensor + ) -> None: + """Copy this rank's overlapping slice from a full tensor into a DBuffer.""" + tensor = tensor.contiguous().view(-1) + shape = buffer.layout.tensor_shapes[index] + tensor_start = buffer.layout.tensor_to_offset[index] + tensor_end = tensor_start + shape.numel() + buffer_start = buffer.offset + buffer_end = buffer.offset + buffer.local_buffer.numel() + overlap_start = max(tensor_start, buffer_start) + overlap_end = min(tensor_end, buffer_end) + if overlap_end <= overlap_start: + return + local_numel = overlap_end - overlap_start + buffer.local_buffer.narrow(0, overlap_start - buffer_start, local_numel).copy_( + tensor.narrow(0, overlap_start - tensor_start, local_numel) + ) + + def _set_module_parameter(self, name: str, parameter: nn.Parameter) -> None: + module, parameter_name = _get_parameter_owner(self.module, name) + module._parameters[parameter_name] = parameter + + def refresh_model_weight(self) -> None: + """Refresh compute-weight storage from the current main weights.""" + self.model_weight = self._make_model_weight(self.model_weight.placements) + + def unshard_parameters(self) -> None: + """Install full parameters for local compute.""" + self.refresh_model_weight() + self._allocate_full_weight() + assert self._full_weight is not None + self._ensure_unsharded_parameters() + for name, parameter in self._unsharded_parameters.items(): + parameter.grad = None + self._set_module_parameter(name, parameter) + + def reshard_parameters(self) -> None: + """Install sharded DTensor parameters on the owning modules.""" + for name, parameter in self.parameters.items(): + self._set_module_parameter(name, parameter) + self._release_full_weight_storage() + + def _ensure_unsharded_parameters(self) -> None: + """Create stable full-size autograd leaf parameters.""" + if self._unsharded_parameters: + return + assert self._full_weight is not None + for index, (name, sharded_parameter) in enumerate(self.parameters.items()): + self._unsharded_parameters[name] = nn.Parameter( + self._full_weight.get_tensor(index), requires_grad=sharded_parameter.requires_grad + ) + + def _allocate_full_weight(self) -> None: + """Ensure the stable full-weight storage is allocated and current.""" + if self._full_weight is None: + self._full_weight = self.model_weight.redistribute([Replicate()] * self.mesh.ndim) + self._full_weight_allocated = True + return + if self._full_weight_allocated: + return + + self._resize_full_weight_storage(self._full_weight.local_buffer.numel()) + refreshed = self.model_weight.redistribute([Replicate()] * self.mesh.ndim) + with torch.autograd._unsafe_preserve_version_counter(self._full_weight.local_buffer): + self._full_weight.local_buffer.copy_(refreshed.local_buffer) + self._full_weight_allocated = True + + def _release_full_weight_storage(self) -> None: + """Free full-weight storage without replacing the Storage object.""" + if self._full_weight is None or not self._full_weight_allocated: + return + # Non-leaf parameter views saved by autograd keep their Storage object. + # Retaining that object and resizing it back before backward avoids + # stale-storage failures while still releasing the allocation. + self._resize_full_weight_storage(0) + self._full_weight_allocated = False + + def _resize_full_weight_storage(self, numel: int) -> None: + assert self._full_weight is not None + with torch.autograd._unsafe_preserve_version_counter(self._full_weight.local_buffer): + self._full_weight.local_buffer.untyped_storage().resize_( + numel * self._full_weight.local_buffer.element_size() + ) + + def reduce_gradients(self, average: bool = True) -> None: + """Reduce full local gradients into the persistent sharded gradient buffer.""" + if not self.requires_grad: + return + assert self.main_grad is not None + + accumulate = self.has_sharded_grad() + full_grads: list[torch.Tensor] = [] + for name, parameter in self._unsharded_parameters.items(): + if parameter.grad is None: + raise RuntimeError(f"Missing gradient for FSDP parameter {name!r}.") + full_grads.append( + parameter.grad.detach().to(dtype=self.main_grad.local_buffer.dtype).contiguous() + ) + + partial_grad = DBuffer.distribute_tensors( + full_grads, mesh=self.mesh, placements=[Partial()] * self.mesh.ndim + ) + reduced_grad = partial_grad.redistribute(self.main_grad.placements) + if average: + reduced_grad.local_buffer.div_(self.mesh.size()) + + if accumulate: + self.main_grad.local_buffer.add_(reduced_grad.local_buffer) + else: + self.main_grad.local_buffer.copy_(reduced_grad.local_buffer) + + for parameter in self._unsharded_parameters.values(): + parameter.grad = None + self.install_sharded_gradients() + + def has_sharded_grad(self) -> bool: + """Return whether persistent sharded gradients are currently materialized.""" + return any(parameter.grad is not None for parameter in self.parameters.values()) + + def install_sharded_gradients(self) -> None: + """Install sharded DTensor gradients backed by main_grad.""" + if self.main_grad is None: + return + for index, parameter in enumerate(self.parameters.values()): + parameter.grad = self.main_grad.get_dtensor(index) + + def zero_grad(self, set_to_none: bool = True) -> None: + """Clear persistent and parameter gradients.""" + if self.main_grad is not None: + self.main_grad.local_buffer.zero_() + if set_to_none: + for parameter in self.parameters.values(): + parameter.grad = None + else: + self.install_sharded_gradients() + for parameter in self._unsharded_parameters.values(): + parameter.grad = None + + def unsharded_parameters(self) -> tuple[nn.Parameter, ...]: + """Return temporary full-size parameters used for autograd.""" + return tuple(self._unsharded_parameters.values()) + + +class FsdpModule: + """Mixin attached to modules managed by the minimal experimental FSDP path.""" + + _parameter_groups: tuple[ParameterGroup, ...] + _ready_grad_params: set[nn.Parameter] + _registered_grad_param_ids: set[int] + _trainable_param_count: int + + def __init__(self, parameter_groups: Sequence[ParameterGroup]) -> None: + """Initialize mixin runtime state on an already-constructed module.""" + self._parameter_groups = tuple(parameter_groups) + self._ready_grad_params: set[nn.Parameter] = set() + self._registered_grad_param_ids: set[int] = set() + self._trainable_param_count = sum( + len(group.parameters) for group in self._parameter_groups if group.requires_grad + ) + self._register_hooks() + + def _register_hooks(self) -> None: + self.register_forward_pre_hook(lambda _module, _args: self.pre_forward()) + self.register_forward_hook(lambda _module, _args, _output: self.post_forward()) + self.register_full_backward_pre_hook(lambda _module, _grad_output: self.pre_backward()) + + def _register_grad_hooks(self) -> None: + """Register post-accumulate hooks on full-size autograd leaf parameters.""" + for group in self._parameter_groups: + if not group.requires_grad: + continue + for parameter in group.unsharded_parameters(): + parameter_id = id(parameter) + if parameter_id in self._registered_grad_param_ids: + continue + parameter.register_post_accumulate_grad_hook(self._make_grad_hook(parameter)) + self._registered_grad_param_ids.add(parameter_id) + + def _make_grad_hook(self, parameter: nn.Parameter) -> Callable[[nn.Parameter], None]: + def grad_hook(_parameter: nn.Parameter) -> None: + self._ready_grad_params.add(parameter) + if len(self._ready_grad_params) == self._trainable_param_count: + self.post_backward() + + return grad_hook + + def pre_forward(self) -> None: + """Prepare full parameters for forward compute.""" + self._ready_grad_params.clear() + for group in self._parameter_groups: + group.unshard_parameters() + self._register_grad_hooks() + + def post_forward(self) -> None: + """Return parameters to their sharded resting state after forward compute.""" + for group in self._parameter_groups: + group.reshard_parameters() + + def pre_backward(self) -> None: + """Prepare full parameters for backward compute.""" + for group in self._parameter_groups: + group.unshard_parameters() + + def post_backward(self) -> None: + """Reduce gradients and return parameters to their sharded resting state.""" + for group in self._parameter_groups: + group.reduce_gradients() + group.reshard_parameters() + group.install_sharded_gradients() + self._ready_grad_params.clear() + + def parameter_groups(self) -> tuple[ParameterGroup, ...]: + """Return parameter groups owned by this FSDP unit.""" + return self._parameter_groups + + +def fully_shard( + module: nn.Module, + mesh: DeviceMesh, + placements: Placements, + mixed_precision_policy: MixedPrecisionPolicy | None = None, + init_model_with_meta_device: bool = False, +) -> None: + """Shard one module as an experimental per-module FSDP unit. + + Args: + module: Module whose currently unowned parameters become this FSDP unit. + mesh: Device mesh used for sharding. + placements: Parameter, gradient, and optimizer placements. + mixed_precision_policy: Optional precision policy. Defaults to FP32 main weights + and parameter-dtype main gradients. + init_model_with_meta_device: If true, initialize owned meta parameters by + calling their direct module's reset_parameters() or _reset_parameters(). + """ + if isinstance(module, FsdpModule): + raise ValueError("This module is already managed by experimental FSDP.") + + mixed_precision_policy = mixed_precision_policy or MixedPrecisionPolicy() + model_weight_placements = _placements_in_mesh_order( + mesh, placements.dp_axes, placements.parameter + ) + main_grad_placements = _placements_in_mesh_order(mesh, placements.dp_axes, placements.gradient) + main_weight_placements = _placements_in_mesh_order( + mesh, placements.dp_axes, placements.optimizer + ) + owned_parameters = _collect_owned_parameters(module) + if ( + any(parameter.is_meta for parameter in owned_parameters.values()) + and not init_model_with_meta_device + ): + raise ValueError( + "experimental FSDP found meta parameters. Pass init_model_with_meta_device=True " + "to initialize them with reset_parameters()/_reset_parameters()." + ) + grouped_parameters = _group_parameters(owned_parameters) + parameter_groups = [ + ParameterGroup( + module, + group_parameters, + mesh=mesh, + model_weight_placements=model_weight_placements, + main_grad_placements=main_grad_placements, + main_weight_placements=main_weight_placements, + mixed_precision_policy=mixed_precision_policy, + ) + for group_parameters in grouped_parameters + ] + if init_model_with_meta_device: + _reset_owned_meta_modules(module, owned_parameters) + for group in parameter_groups: + group.refresh_model_weight() + + _attach_mixin(module) + assert isinstance(module, FsdpModule) + FsdpModule.__init__(module, parameter_groups=parameter_groups) + + +def _axis_index(mesh: DeviceMesh, axis: MeshAxis) -> int: + if isinstance(axis, int): + if axis < 0: + axis += mesh.ndim + if axis < 0 or axis >= mesh.ndim: + raise ValueError(f"Mesh axis {axis} is out of bounds for mesh ndim {mesh.ndim}.") + return axis + + dim_names = mesh.mesh_dim_names + if dim_names is None or axis not in dim_names: + raise ValueError(f"Mesh axis {axis!r} is not present in mesh dim names {dim_names}.") + return dim_names.index(axis) + + +def _placements_in_mesh_order( + mesh: DeviceMesh, dp_axes: Sequence[MeshAxis], placements: Sequence[Placement] +) -> tuple[Placement, ...]: + if len(dp_axes) != mesh.ndim: + raise ValueError( + "experimental fully_shard currently requires placements for every mesh axis: " + f"mesh ndim is {mesh.ndim}, got {len(dp_axes)} axes." + ) + result: list[Placement | None] = [None] * mesh.ndim + for axis, placement in zip(dp_axes, placements, strict=True): + axis_index = _axis_index(mesh, axis) + if result[axis_index] is not None: + raise ValueError(f"Duplicate placement for mesh axis {axis!r}.") + result[axis_index] = placement + if any(placement is None for placement in result): + raise ValueError("Missing placement for at least one mesh axis.") + return tuple(placement for placement in result if placement is not None) + + +def _mesh_device(mesh: DeviceMesh) -> torch.device: + if mesh.device_type == "cuda": + return torch.device("cuda", torch.cuda.current_device()) + return torch.device(mesh.device_type) + + +def _collect_owned_parameters(module: nn.Module) -> dict[str, nn.Parameter]: + child_prefixes = [ + f"{name}." + for name, child in module.named_modules() + if name and isinstance(child, FsdpModule) + ] + parameters: dict[str, nn.Parameter] = {} + for module_name, child in module.named_modules(): + prefix = f"{module_name}." if module_name else "" + if any(prefix.startswith(child_prefix) for child_prefix in child_prefixes): + continue + for parameter_name, parameter in child.named_parameters(recurse=False): + name = f"{prefix}{parameter_name}" + if hasattr(parameter, _CONTAINING_PARAMETER_GROUP_ATTR): + raise ValueError( + f"Parameter {name!r} is already owned by an experimental FSDP unit." + ) + parameters[name] = parameter + if not parameters: + raise ValueError("experimental fully_shard requires at least one unowned parameter.") + return parameters + + +def _group_parameters(parameters: dict[str, nn.Parameter]) -> list[dict[str, nn.Parameter]]: + grouped: dict[tuple[torch.dtype, bool], dict[str, nn.Parameter]] = {} + for name, parameter in parameters.items(): + key = (parameter.dtype, parameter.requires_grad) + grouped.setdefault(key, {})[name] = parameter + return [grouped[key] for key in grouped] + + +def _get_parameter_owner(module: nn.Module, name: str) -> tuple[nn.Module, str]: + """Resolve a root-module-relative parameter FQN to its direct owner.""" + module_name, separator, parameter_name = name.rpartition(".") + owner = module.get_submodule(module_name) if separator else module + return owner, parameter_name + + +def _reset_owned_meta_modules( + module: nn.Module, original_parameters: dict[str, nn.Parameter] +) -> None: + """Initialize modules that originally owned meta parameters.""" + modules_to_reset: dict[int, tuple[str, nn.Module]] = {} + for name, parameter in original_parameters.items(): + if not parameter.is_meta: + continue + owner_module, _ = _get_parameter_owner(module, name) + module_name = name.rsplit(".", 1)[0] if "." in name else "" + modules_to_reset.setdefault(id(owner_module), (module_name, owner_module)) + + for module_name, module in modules_to_reset.values(): + if hasattr(module, "reset_parameters"): + module.reset_parameters() + elif hasattr(module, "_reset_parameters"): + module._reset_parameters() + else: + raise ValueError( + f"[init_model_with_meta_device=True] Module {module_name!r} does not have " + "reset_parameters or _reset_parameters." + ) + + +def _attach_mixin(module: nn.Module) -> None: + if isinstance(module, FsdpModule): + return + module_cls = module.__class__ + fsdp_cls = type(f"ExperimentalFsdp{module_cls.__name__}", (FsdpModule, module_cls), {}) + module.__class__ = fsdp_cls diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py new file mode 100644 index 00000000000..9f7f6926773 --- /dev/null +++ b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py @@ -0,0 +1,319 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for the experimental minimal Megatron-FSDP path.""" + +import gc +import os +from collections.abc import Iterator +from dataclasses import dataclass + +import pytest +import torch +import torch.distributed as dist +from torch import nn +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor import DTensor + +from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental import ( + Flat, + FsdpModule, + Placements, + fully_shard, +) + + +@dataclass(frozen=True) +class DistributedSetup: + """Per-rank distributed test setup.""" + + rank: int + world_size: int + device: torch.device + + +@pytest.fixture(scope="module") +def setup() -> Iterator[DistributedSetup]: + """Read torchrun rank state and set up this rank's local device.""" + if "RANK" not in os.environ or "WORLD_SIZE" not in os.environ: + pytest.skip("Not running under torchrun.") + + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ.get("LOCAL_RANK", rank)) + + if torch.cuda.is_available(): + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + else: + device = torch.device("cpu") + + yield DistributedSetup(rank=rank, world_size=world_size, device=device) + + if dist.is_initialized(): + dist.destroy_process_group() + + +class TinyModel(nn.Module): + """Small model with two separately shardable units.""" + + def __init__(self) -> None: + super().__init__() + self.fc1 = nn.Linear(8, 16) + self.relu = nn.ReLU() + self.fc2 = nn.Linear(16, 4) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Run the tiny model.""" + return self.fc2(self.relu(self.fc1(x))) + + +class NestedModel(nn.Module): + """Model with direct and child-owned parameters.""" + + def __init__(self) -> None: + super().__init__() + self.bias = nn.Parameter(torch.ones(4)) + self.inner = nn.Linear(4, 4, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Run the nested model.""" + return self.inner(x) + self.bias + + +class SaveNonLeafWeightView(torch.autograd.Function): + """Autograd function that saves a non-leaf parameter view for backward.""" + + @staticmethod + def forward(ctx, x: torch.Tensor, weight_view: torch.Tensor) -> torch.Tensor: + """Save the non-leaf weight view and run a simple elementwise op.""" + ctx.save_for_backward(x, weight_view) + return x * weight_view + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Use the saved non-leaf weight view during backward.""" + x, weight_view = ctx.saved_tensors + return grad_output * weight_view, grad_output * x + + +class NonLeafViewModel(nn.Module): + """Model that saves a non-leaf parameter view across forward and backward.""" + + def __init__(self) -> None: + super().__init__() + self.weight = nn.Parameter(torch.randn(8)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Run using a non-leaf view of the parameter.""" + return SaveNonLeafWeightView.apply(x, self.weight.view_as(self.weight)) + + +def _flat_placements() -> Placements: + return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()]) + + +def _full_named_parameters(module: nn.Module) -> dict[str, torch.Tensor]: + result = {} + for module_name, child in module.named_modules(): + if not isinstance(child, FsdpModule): + continue + prefix = f"{module_name}." if module_name else "" + for group in child.parameter_groups(): + full_weight = group.main_weight.redistribute([Flat()]).allgather(0) + for index, name in enumerate(group.parameters): + result[f"{prefix}{name}"] = full_weight.get_tensor(index).detach().clone() + return result + + +@pytest.mark.distributed +def test_experimental_fully_shard_train_step_matches_baseline(setup: DistributedSetup): + """A minimal per-module FSDP train step should match single-rank SGD.""" + if setup.world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + + mesh = init_device_mesh(setup.device.type, (setup.world_size,)) + torch.manual_seed(1234) + baseline = TinyModel().to(setup.device) + model = TinyModel().to(setup.device) + model.load_state_dict(baseline.state_dict()) + + fully_shard(model.fc1, mesh=mesh, placements=_flat_placements()) + fully_shard(model.fc2, mesh=mesh, placements=_flat_placements()) + optimizer = torch.optim.SGD(model.parameters(), lr=0.05) + + x = torch.randn(3, 8, device=setup.device) + target = torch.randn(3, 4, device=setup.device) + + baseline_loss = torch.nn.functional.mse_loss(baseline(x), target) + baseline_loss.backward() + with torch.no_grad(): + for parameter in baseline.parameters(): + parameter.add_(parameter.grad, alpha=-0.05) + + loss = torch.nn.functional.mse_loss(model(x), target) + loss.backward() + optimizer.step() + + sharded_params = _full_named_parameters(model) + for name, expected in baseline.named_parameters(): + torch.testing.assert_close(sharded_params[name], expected, rtol=1e-5, atol=1e-6) + + +@pytest.mark.distributed +def test_nested_fully_shard_excludes_child_owned_parameters(setup: DistributedSetup): + """An outer FSDP unit owns direct parameters but not nested child-unit parameters.""" + if setup.world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + + mesh = init_device_mesh(setup.device.type, (setup.world_size,)) + model = NestedModel().to(setup.device) + + fully_shard(model.inner, mesh=mesh, placements=_flat_placements()) + fully_shard(model, mesh=mesh, placements=_flat_placements()) + + inner_names = [name for group in model.inner.parameter_groups() for name in group.parameters] + outer_names = [name for group in model.parameter_groups() for name in group.parameters] + + assert inner_names == ["weight"] + assert outer_names == ["bias"] + + +@pytest.mark.distributed +def test_frozen_parameter_group_does_not_allocate_main_grad(setup: DistributedSetup): + """A non-trainable parameter group should not allocate persistent main gradients.""" + if setup.world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + + mesh = init_device_mesh(setup.device.type, (setup.world_size,)) + model = nn.Linear(4, 4, bias=False).to(setup.device) + model.weight.requires_grad_(False) + + fully_shard(model, mesh=mesh, placements=_flat_placements()) + + (group,) = model.parameter_groups() + assert not group.requires_grad + assert group.main_grad is None + + +@pytest.mark.distributed +def test_default_main_buffer_dtypes_follow_policy_contract(setup: DistributedSetup): + """Default main weights are FP32 while default main gradients match parameter dtype.""" + if setup.world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + + mesh = init_device_mesh(setup.device.type, (setup.world_size,)) + model = nn.Linear(4, 4, bias=False, dtype=torch.float64).to(setup.device) + + fully_shard(model, mesh=mesh, placements=_flat_placements()) + + (group,) = model.parameter_groups() + assert group.dtype is torch.float64 + assert group.main_weight.local_buffer.dtype is torch.float32 + assert group.main_grad is not None + assert group.main_grad.local_buffer.dtype is torch.float64 + + +@pytest.mark.distributed +def test_sharded_parameter_contract_uses_dtensors(setup: DistributedSetup): + """Resting sharded parameters and gradients should be DTensors.""" + if setup.world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + + mesh = init_device_mesh(setup.device.type, (setup.world_size,)) + model = nn.Linear(4, 4, bias=False).to(setup.device) + + fully_shard(model, mesh=mesh, placements=_flat_placements()) + + assert isinstance(model.weight, DTensor) + assert isinstance(model.weight.data, DTensor) + model(torch.randn(2, 4, device=setup.device)).sum().backward() + assert isinstance(model.weight, DTensor) + assert isinstance(model.weight.grad, DTensor) + + +@pytest.mark.distributed +def test_meta_parameters_initialize_with_reset_parameters(setup: DistributedSetup): + """Meta parameters should be replaced by sharded DTensors and initialized in place.""" + if setup.world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + + mesh = init_device_mesh(setup.device.type, (setup.world_size,)) + model = nn.Linear(4, 4, bias=False, device="meta") + + fully_shard(model, mesh=mesh, placements=_flat_placements(), init_model_with_meta_device=True) + + assert isinstance(model.weight, DTensor) + assert not model.weight.to_local().is_meta + (group,) = model.parameter_groups() + assert not group.main_weight.local_buffer.is_meta + assert group.main_weight.local_buffer.numel() > 0 + + +@pytest.mark.distributed +def test_non_leaf_parameter_view_survives_storage_resize(setup: DistributedSetup): + """A non-leaf parameter view saved for backward should survive full-storage resize.""" + if setup.world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + if setup.device.type != "cuda": + pytest.skip("Storage resize verification requires CUDA.") + + mesh = init_device_mesh(setup.device.type, (setup.world_size,)) + model = NonLeafViewModel().to(setup.device) + fully_shard(model, mesh=mesh, placements=_flat_placements()) + + group = model.parameter_groups()[0] + x = torch.randn(8, device=setup.device, requires_grad=True) + loss = model(x).sum() + + assert group._full_weight is not None + assert group._full_weight.local_buffer.untyped_storage().nbytes() == 0 + + loss.backward() + + assert group.main_grad is not None + assert group._full_weight is not None + assert group._full_weight.local_buffer.untyped_storage().nbytes() == 0 + + +@pytest.mark.distributed +def test_experimental_fully_shard_reduces_peak_training_memory(setup: DistributedSetup): + """Per-layer experimental FSDP should reduce peak CUDA memory during a train step.""" + if setup.world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + if setup.device.type != "cuda": + pytest.skip("Peak memory verification requires CUDA.") + + mesh = init_device_mesh(setup.device.type, (setup.world_size,)) + dim = 1024 + layers = 8 + batch = 8 + + torch.manual_seed(4321) + baseline = nn.Sequential(*[nn.Linear(dim, dim, bias=False) for _ in range(layers)]).to( + setup.device + ) + x = torch.randn(batch, dim, device=setup.device) + torch.cuda.reset_peak_memory_stats(setup.device) + baseline(x).sum().backward() + torch.cuda.synchronize(setup.device) + baseline_peak = torch.cuda.max_memory_allocated(setup.device) + + del baseline + del x + gc.collect() + torch.cuda.empty_cache() + + torch.manual_seed(4321) + model = nn.Sequential(*[nn.Linear(dim, dim, bias=False) for _ in range(layers)]).to( + setup.device + ) + for layer in model: + fully_shard(layer, mesh=mesh, placements=_flat_placements()) + + x = torch.randn(batch, dim, device=setup.device) + torch.cuda.reset_peak_memory_stats(setup.device) + model(x).sum().backward() + torch.cuda.synchronize(setup.device) + sharded_peak = torch.cuda.max_memory_allocated(setup.device) + + assert sharded_peak < baseline_peak From 832af4673de7fe38fda6cc3f4eed1b8a8bfbb864 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sat, 23 May 2026 06:44:32 +0000 Subject: [PATCH 02/23] Refine minimal experimental FSDP path Summarize today's FSDP work: - Add DBuffer storage release/reallocate support and an in-place fully_allgather_into path for materializing replicated buffers. - Simplify ParameterGroup and FsdpModule around sharded DTensor parameters, reused unsharded Parameters, meta materialization, and default-stream unshard/reshard/reduce behavior. - Remove unused optimizer/offload/state/helper surface area from the minimal path and keep version-counter preservation scoped to unsharded model-weight materialization. - Expand DBuffer and experimental FSDP tests for layouts, storage lifecycle, DTensor contracts, meta reset, nested ownership, train-step parity, and peak-memory reduction. Signed-off-by: Jingyue Wu --- .../src/megatron_fsdp/experimental/dbuffer.py | 14 + .../megatron_fsdp/experimental/fully_shard.py | 439 +++++++----------- .../distributed/megatron_fsdp/test_dbuffer.py | 29 ++ .../test_experimental_fully_shard.py | 107 ++++- 4 files changed, 308 insertions(+), 281 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py index 51f52451089..ea4e5ce0bda 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py @@ -116,6 +116,20 @@ def device(self) -> torch.device: """Device of the local buffer.""" return self.local_buffer.device + def reallocate_storage(self) -> None: + """Restore the local buffer's backing storage to its logical size.""" + self._resize_storage(self.local_buffer.numel()) + + def release_storage(self) -> None: + """Release local buffer storage without replacing the Storage object.""" + # Autograd may save views that share this Storage object. Resizing the + # existing Storage releases the allocation while preserving those aliases + # for a later reallocate_storage(). + self._resize_storage(0) + + def _resize_storage(self, numel: int) -> None: + self.local_buffer.untyped_storage().resize_(numel * self.local_buffer.element_size()) + def _get_owned_range(self, tensor_index: int) -> _OwnedRange | None: """Return this buffer's owned range for logical tensor ``tensor_index``.""" tensor_start = self.layout.tensor_to_offset[tensor_index] diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py index ba664e82387..dba988ea4d9 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Minimal experimental per-module Megatron-FSDP implementation.""" +"""Minimal per-module Megatron-FSDP implementation.""" import dataclasses from collections.abc import Callable, Sequence @@ -51,53 +51,55 @@ def __post_init__(self) -> None: class ParameterGroup: """A dtype and requires-grad homogeneous group of FSDP-owned parameters.""" - module: nn.Module - parameters: dict[str, nn.Parameter] - _unsharded_parameters: dict[str, nn.Parameter] + owning_module: nn.Module + sharded_parameters: dict[str, nn.Parameter] + unsharded_parameters: dict[str, nn.Parameter] mesh: DeviceMesh + dp_mesh: DeviceMesh dtype: torch.dtype requires_grad: bool main_weight: DBuffer model_weight: DBuffer main_grad: DBuffer | None - _full_weight: DBuffer | None - _full_weight_allocated: bool + _unsharded_model_weight: DBuffer def __init__( self, - module: nn.Module, + owning_module: nn.Module, parameters: dict[str, nn.Parameter], mesh: DeviceMesh, - model_weight_placements: Sequence[Placement], - main_grad_placements: Sequence[Placement], - main_weight_placements: Sequence[Placement], + placements: Placements, mixed_precision_policy: MixedPrecisionPolicy, ) -> None: """Create persistent sharded buffers for a group of parameters. Args: - module: Closest FSDP root module that owns this parameter group. + owning_module: Closest FSDP root module that owns this parameter group. parameters: Root-module-relative FQNs and their parameters. - mesh: Device mesh used by the buffers. - model_weight_placements: Placements for compute-weight storage. - main_grad_placements: Placements for persistent main gradients. - main_weight_placements: Placements for optimizer-owned main weights. + mesh: Full device mesh. DBuffer storage is built on the DP submesh. + placements: Parameter, gradient, and optimizer placements. mixed_precision_policy: Precision policy for main weights and gradients. """ if not parameters: raise ValueError("ParameterGroup requires at least one parameter.") + dp_mesh = _dp_submesh(mesh, placements.dp_axes) + # Python dicts preserve insertion order, so values() defines the stable # tensor order used by each DBuffer built from this group. - original_parameters = parameters - self.module = module - self.parameters = {} - self._unsharded_parameters = {} + self.owning_module = owning_module + self.sharded_parameters = {} + self.unsharded_parameters = {} self.mesh = mesh - first_parameter = next(iter(original_parameters.values())) + self.dp_mesh = dp_mesh + first_parameter = next(iter(parameters.values())) self.dtype = first_parameter.dtype self.requires_grad = first_parameter.requires_grad - for name, parameter in original_parameters.items(): + for name, parameter in parameters.items(): + if parameter.is_meta: + raise ValueError( + f"Expected parameter {name!r} to be materialized before ParameterGroup construction." + ) if parameter.dtype != self.dtype: raise ValueError( f"Expected parameter {name!r} to have dtype {self.dtype}, got {parameter.dtype}." @@ -107,152 +109,87 @@ def __init__( f"Expected parameter {name!r} to have requires_grad={self.requires_grad}, " f"got {parameter.requires_grad}." ) - main_params_dtype = mixed_precision_policy.main_params_dtype - if main_params_dtype is None: + main_weight_dtype = mixed_precision_policy.main_params_dtype + if main_weight_dtype is None: raise ValueError( - "experimental FSDP requires main_params_dtype to be specified explicitly." + "FSDP requires a main weight dtype; set MixedPrecisionPolicy.main_params_dtype." ) - main_grads_dtype = mixed_precision_policy.main_grads_dtype - if main_grads_dtype is None: - main_grads_dtype = self.dtype - self.main_weight = DBuffer( - mesh=self.mesh, - placements=main_weight_placements, - tensor_shapes=[parameter.shape for parameter in original_parameters.values()], - dtype=main_params_dtype, - device=_mesh_device(self.mesh), + main_grad_dtype = mixed_precision_policy.main_grads_dtype + if main_grad_dtype is None: + main_grad_dtype = self.dtype + # Scratch initialization starts from model weights. Checkpoint loading will + # eventually initialize main weights first and derive model weights from them. + self.model_weight = DBuffer.distribute_tensors( + [parameter.detach().contiguous() for parameter in parameters.values()], + mesh=self.dp_mesh, + placements=placements.parameter, + ) + self.main_weight = DBuffer.distribute_tensors( + [ + parameter.detach().to(dtype=main_weight_dtype).contiguous() + for parameter in parameters.values() + ], + mesh=self.dp_mesh, + placements=placements.optimizer, ) - for index, parameter in enumerate(original_parameters.values()): - if parameter.is_meta: - continue - self._copy_full_tensor_to_buffer( - self.main_weight, - index, - parameter.detach().to( - dtype=main_params_dtype, device=self.main_weight.local_buffer.device - ), - ) - self.model_weight = self._make_model_weight(tuple(model_weight_placements)) self.main_grad = ( DBuffer( - mesh=self.mesh, - placements=main_grad_placements, + mesh=self.dp_mesh, + placements=placements.gradient, tensor_shapes=self.main_weight.layout.tensor_shapes, - dtype=main_grads_dtype, + dtype=main_grad_dtype, device=self.main_weight.local_buffer.device, ) if self.requires_grad else None ) - self._full_weight: DBuffer | None = None - self._full_weight_allocated = False + self._unsharded_model_weight = DBuffer( + mesh=self.dp_mesh, + placements=[Replicate()] * self.dp_mesh.ndim, + tensor_shapes=self.model_weight.layout.tensor_shapes, + dtype=self.model_weight.local_buffer.dtype, + device=self.model_weight.local_buffer.device, + ) - for index, (name, original_parameter) in enumerate(original_parameters.items()): + for index, (name, parameter) in enumerate(parameters.items()): + parameter.data = self._unsharded_model_weight.get_tensor(index) + parameter.grad = None + setattr(parameter, _CONTAINING_PARAMETER_GROUP_ATTR, self) + self.unsharded_parameters[name] = parameter sharded_parameter = nn.Parameter( - self.main_weight.get_dtensor(index), requires_grad=original_parameter.requires_grad + self.main_weight.get_dtensor(index), requires_grad=parameter.requires_grad ) setattr(sharded_parameter, _CONTAINING_PARAMETER_GROUP_ATTR, self) - self.parameters[name] = sharded_parameter - self._set_module_parameter(name, sharded_parameter) - - def _make_model_weight(self, placements: Sequence[Placement]) -> DBuffer: - model_weight = self.main_weight.redistribute(placements) - if model_weight.local_buffer.dtype != self.dtype: - converted = DBuffer( - mesh=model_weight.mesh, - placements=model_weight.placements, - tensor_shapes=model_weight.layout.tensor_shapes, - dtype=self.dtype, - device=model_weight.local_buffer.device, - ) - converted.local_buffer.copy_(model_weight.local_buffer.to(dtype=self.dtype)) - return converted - return model_weight + self.sharded_parameters[name] = sharded_parameter + self._switch_to_sharded_parameters() + self._unsharded_model_weight.release_storage() - def _copy_full_tensor_to_buffer( - self, buffer: DBuffer, index: int, tensor: torch.Tensor - ) -> None: - """Copy this rank's overlapping slice from a full tensor into a DBuffer.""" - tensor = tensor.contiguous().view(-1) - shape = buffer.layout.tensor_shapes[index] - tensor_start = buffer.layout.tensor_to_offset[index] - tensor_end = tensor_start + shape.numel() - buffer_start = buffer.offset - buffer_end = buffer.offset + buffer.local_buffer.numel() - overlap_start = max(tensor_start, buffer_start) - overlap_end = min(tensor_end, buffer_end) - if overlap_end <= overlap_start: - return - local_numel = overlap_end - overlap_start - buffer.local_buffer.narrow(0, overlap_start - buffer_start, local_numel).copy_( - tensor.narrow(0, overlap_start - tensor_start, local_numel) - ) + def _set_module_parameters(self, parameters: dict[str, nn.Parameter]) -> None: + for name, parameter in parameters.items(): + module, parameter_name = _get_parameter_owner(self.owning_module, name) + module._parameters[parameter_name] = parameter - def _set_module_parameter(self, name: str, parameter: nn.Parameter) -> None: - module, parameter_name = _get_parameter_owner(self.module, name) - module._parameters[parameter_name] = parameter + def _switch_to_sharded_parameters(self) -> None: + self._set_module_parameters(self.sharded_parameters) - def refresh_model_weight(self) -> None: - """Refresh compute-weight storage from the current main weights.""" - self.model_weight = self._make_model_weight(self.model_weight.placements) + def _switch_to_unsharded_parameters(self) -> None: + self._set_module_parameters(self.unsharded_parameters) def unshard_parameters(self) -> None: """Install full parameters for local compute.""" - self.refresh_model_weight() - self._allocate_full_weight() - assert self._full_weight is not None - self._ensure_unsharded_parameters() - for name, parameter in self._unsharded_parameters.items(): - parameter.grad = None - self._set_module_parameter(name, parameter) + self._unsharded_model_weight.reallocate_storage() + # This buffer backs unsharded Parameters whose views may be saved by autograd. + # Materializing FSDP-managed storage should not look like a user mutation. + with torch.autograd._unsafe_preserve_version_counter( + self._unsharded_model_weight.local_buffer + ): + self.model_weight.fully_allgather_into(self._unsharded_model_weight) + self._switch_to_unsharded_parameters() def reshard_parameters(self) -> None: """Install sharded DTensor parameters on the owning modules.""" - for name, parameter in self.parameters.items(): - self._set_module_parameter(name, parameter) - self._release_full_weight_storage() - - def _ensure_unsharded_parameters(self) -> None: - """Create stable full-size autograd leaf parameters.""" - if self._unsharded_parameters: - return - assert self._full_weight is not None - for index, (name, sharded_parameter) in enumerate(self.parameters.items()): - self._unsharded_parameters[name] = nn.Parameter( - self._full_weight.get_tensor(index), requires_grad=sharded_parameter.requires_grad - ) - - def _allocate_full_weight(self) -> None: - """Ensure the stable full-weight storage is allocated and current.""" - if self._full_weight is None: - self._full_weight = self.model_weight.redistribute([Replicate()] * self.mesh.ndim) - self._full_weight_allocated = True - return - if self._full_weight_allocated: - return - - self._resize_full_weight_storage(self._full_weight.local_buffer.numel()) - refreshed = self.model_weight.redistribute([Replicate()] * self.mesh.ndim) - with torch.autograd._unsafe_preserve_version_counter(self._full_weight.local_buffer): - self._full_weight.local_buffer.copy_(refreshed.local_buffer) - self._full_weight_allocated = True - - def _release_full_weight_storage(self) -> None: - """Free full-weight storage without replacing the Storage object.""" - if self._full_weight is None or not self._full_weight_allocated: - return - # Non-leaf parameter views saved by autograd keep their Storage object. - # Retaining that object and resizing it back before backward avoids - # stale-storage failures while still releasing the allocation. - self._resize_full_weight_storage(0) - self._full_weight_allocated = False - - def _resize_full_weight_storage(self, numel: int) -> None: - assert self._full_weight is not None - with torch.autograd._unsafe_preserve_version_counter(self._full_weight.local_buffer): - self._full_weight.local_buffer.untyped_storage().resize_( - numel * self._full_weight.local_buffer.element_size() - ) + self._switch_to_sharded_parameters() + self._unsharded_model_weight.release_storage() def reduce_gradients(self, average: bool = True) -> None: """Reduce full local gradients into the persistent sharded gradient buffer.""" @@ -262,7 +199,7 @@ def reduce_gradients(self, average: bool = True) -> None: accumulate = self.has_sharded_grad() full_grads: list[torch.Tensor] = [] - for name, parameter in self._unsharded_parameters.items(): + for name, parameter in self.unsharded_parameters.items(): if parameter.grad is None: raise RuntimeError(f"Missing gradient for FSDP parameter {name!r}.") full_grads.append( @@ -270,64 +207,60 @@ def reduce_gradients(self, average: bool = True) -> None: ) partial_grad = DBuffer.distribute_tensors( - full_grads, mesh=self.mesh, placements=[Partial()] * self.mesh.ndim + full_grads, mesh=self.dp_mesh, placements=[Partial()] * self.dp_mesh.ndim ) reduced_grad = partial_grad.redistribute(self.main_grad.placements) if average: - reduced_grad.local_buffer.div_(self.mesh.size()) + reduced_grad.local_buffer.div_(self.dp_mesh.size()) if accumulate: self.main_grad.local_buffer.add_(reduced_grad.local_buffer) else: self.main_grad.local_buffer.copy_(reduced_grad.local_buffer) - for parameter in self._unsharded_parameters.values(): + for parameter in self.unsharded_parameters.values(): parameter.grad = None self.install_sharded_gradients() def has_sharded_grad(self) -> bool: """Return whether persistent sharded gradients are currently materialized.""" - return any(parameter.grad is not None for parameter in self.parameters.values()) + return any(parameter.grad is not None for parameter in self.sharded_parameters.values()) def install_sharded_gradients(self) -> None: """Install sharded DTensor gradients backed by main_grad.""" if self.main_grad is None: return - for index, parameter in enumerate(self.parameters.values()): + for index, parameter in enumerate(self.sharded_parameters.values()): parameter.grad = self.main_grad.get_dtensor(index) - def zero_grad(self, set_to_none: bool = True) -> None: - """Clear persistent and parameter gradients.""" - if self.main_grad is not None: - self.main_grad.local_buffer.zero_() - if set_to_none: - for parameter in self.parameters.values(): - parameter.grad = None - else: - self.install_sharded_gradients() - for parameter in self._unsharded_parameters.values(): - parameter.grad = None - - def unsharded_parameters(self) -> tuple[nn.Parameter, ...]: - """Return temporary full-size parameters used for autograd.""" - return tuple(self._unsharded_parameters.values()) - - class FsdpModule: - """Mixin attached to modules managed by the minimal experimental FSDP path.""" + """Mixin attached to modules managed by the minimal FSDP path.""" _parameter_groups: tuple[ParameterGroup, ...] _ready_grad_params: set[nn.Parameter] _registered_grad_param_ids: set[int] _trainable_param_count: int - def __init__(self, parameter_groups: Sequence[ParameterGroup]) -> None: - """Initialize mixin runtime state on an already-constructed module.""" + def __init__( + self, mesh: DeviceMesh, placements: Placements, mixed_precision_policy: MixedPrecisionPolicy + ) -> None: + """Initialize FSDP runtime state on an already-constructed module.""" + owned_parameters = _materialize_and_collect_owned_parameters(self, _mesh_device(mesh)) + parameter_groups = [ + ParameterGroup( + owning_module=self, + parameters=group_parameters, + mesh=mesh, + placements=placements, + mixed_precision_policy=mixed_precision_policy, + ) + for group_parameters in _group_parameters(owned_parameters) + ] self._parameter_groups = tuple(parameter_groups) self._ready_grad_params: set[nn.Parameter] = set() self._registered_grad_param_ids: set[int] = set() self._trainable_param_count = sum( - len(group.parameters) for group in self._parameter_groups if group.requires_grad + len(group.sharded_parameters) for group in self._parameter_groups if group.requires_grad ) self._register_hooks() @@ -341,7 +274,7 @@ def _register_grad_hooks(self) -> None: for group in self._parameter_groups: if not group.requires_grad: continue - for parameter in group.unsharded_parameters(): + for parameter in group.unsharded_parameters.values(): parameter_id = id(parameter) if parameter_id in self._registered_grad_param_ids: continue @@ -391,9 +324,8 @@ def fully_shard( mesh: DeviceMesh, placements: Placements, mixed_precision_policy: MixedPrecisionPolicy | None = None, - init_model_with_meta_device: bool = False, ) -> None: - """Shard one module as an experimental per-module FSDP unit. + """Shard one module as a per-module FSDP unit. Args: module: Module whose currently unowned parameters become this FSDP unit. @@ -401,50 +333,21 @@ def fully_shard( placements: Parameter, gradient, and optimizer placements. mixed_precision_policy: Optional precision policy. Defaults to FP32 main weights and parameter-dtype main gradients. - init_model_with_meta_device: If true, initialize owned meta parameters by - calling their direct module's reset_parameters() or _reset_parameters(). """ if isinstance(module, FsdpModule): - raise ValueError("This module is already managed by experimental FSDP.") + raise ValueError("This module is already managed by FSDP.") mixed_precision_policy = mixed_precision_policy or MixedPrecisionPolicy() - model_weight_placements = _placements_in_mesh_order( - mesh, placements.dp_axes, placements.parameter - ) - main_grad_placements = _placements_in_mesh_order(mesh, placements.dp_axes, placements.gradient) - main_weight_placements = _placements_in_mesh_order( - mesh, placements.dp_axes, placements.optimizer - ) - owned_parameters = _collect_owned_parameters(module) - if ( - any(parameter.is_meta for parameter in owned_parameters.values()) - and not init_model_with_meta_device - ): - raise ValueError( - "experimental FSDP found meta parameters. Pass init_model_with_meta_device=True " - "to initialize them with reset_parameters()/_reset_parameters()." - ) - grouped_parameters = _group_parameters(owned_parameters) - parameter_groups = [ - ParameterGroup( - module, - group_parameters, - mesh=mesh, - model_weight_placements=model_weight_placements, - main_grad_placements=main_grad_placements, - main_weight_placements=main_weight_placements, - mixed_precision_policy=mixed_precision_policy, - ) - for group_parameters in grouped_parameters - ] - if init_model_with_meta_device: - _reset_owned_meta_modules(module, owned_parameters) - for group in parameter_groups: - group.refresh_model_weight() - + original_cls = module.__class__ _attach_mixin(module) - assert isinstance(module, FsdpModule) - FsdpModule.__init__(module, parameter_groups=parameter_groups) + try: + assert isinstance(module, FsdpModule) + FsdpModule.__init__( + module, mesh=mesh, placements=placements, mixed_precision_policy=mixed_precision_policy + ) + except Exception: + module.__class__ = original_cls + raise def _axis_index(mesh: DeviceMesh, axis: MeshAxis) -> int: @@ -461,23 +364,26 @@ def _axis_index(mesh: DeviceMesh, axis: MeshAxis) -> int: return dim_names.index(axis) -def _placements_in_mesh_order( - mesh: DeviceMesh, dp_axes: Sequence[MeshAxis], placements: Sequence[Placement] -) -> tuple[Placement, ...]: - if len(dp_axes) != mesh.ndim: +def _dp_submesh(mesh: DeviceMesh, dp_axes: Sequence[MeshAxis]) -> DeviceMesh: + if not dp_axes: + raise ValueError("FSDP requires at least one DP mesh axis.") + + axis_indices = tuple(_axis_index(mesh, axis) for axis in dp_axes) + if len(set(axis_indices)) != len(axis_indices): + raise ValueError(f"Duplicate DP mesh axes are not allowed: {tuple(dp_axes)!r}.") + + if axis_indices == tuple(range(mesh.ndim)): + return mesh + + dim_names = mesh.mesh_dim_names + if dim_names is None: raise ValueError( - "experimental fully_shard currently requires placements for every mesh axis: " - f"mesh ndim is {mesh.ndim}, got {len(dp_axes)} axes." + "Slicing a DP submesh from a full mesh requires named mesh dimensions unless " + "dp_axes covers every mesh axis in mesh order." ) - result: list[Placement | None] = [None] * mesh.ndim - for axis, placement in zip(dp_axes, placements, strict=True): - axis_index = _axis_index(mesh, axis) - if result[axis_index] is not None: - raise ValueError(f"Duplicate placement for mesh axis {axis!r}.") - result[axis_index] = placement - if any(placement is None for placement in result): - raise ValueError("Missing placement for at least one mesh axis.") - return tuple(placement for placement in result if placement is not None) + + dp_axis_names = tuple(dim_names[index] for index in axis_indices) + return mesh[dp_axis_names[0] if len(dp_axis_names) == 1 else dp_axis_names] def _mesh_device(mesh: DeviceMesh) -> torch.device: @@ -486,26 +392,49 @@ def _mesh_device(mesh: DeviceMesh) -> torch.device: return torch.device(mesh.device_type) -def _collect_owned_parameters(module: nn.Module) -> dict[str, nn.Parameter]: - child_prefixes = [ - f"{name}." - for name, child in module.named_modules() - if name and isinstance(child, FsdpModule) - ] +def _materialize_and_collect_owned_parameters( + root_module: nn.Module, device: torch.device +) -> dict[str, nn.Parameter]: parameters: dict[str, nn.Parameter] = {} - for module_name, child in module.named_modules(): - prefix = f"{module_name}." if module_name else "" - if any(prefix.startswith(child_prefix) for child_prefix in child_prefixes): - continue - for parameter_name, parameter in child.named_parameters(recurse=False): - name = f"{prefix}{parameter_name}" - if hasattr(parameter, _CONTAINING_PARAMETER_GROUP_ATTR): + + def visit(submodule: nn.Module, submodule_fqn: str) -> None: + direct_parameters = list(submodule.named_parameters(recurse=False)) + + if any(parameter.is_meta for _, parameter in direct_parameters): + if any(not parameter.is_meta for _, parameter in direct_parameters): raise ValueError( - f"Parameter {name!r} is already owned by an experimental FSDP unit." + f"Module {submodule_fqn!r} mixes meta and non-meta direct parameters. " + "Initialize all direct parameters on meta or none of them." ) - parameters[name] = parameter + submodule.to_empty(device=device, recurse=False) + with torch.no_grad(): + if hasattr(submodule, "reset_parameters"): + submodule.reset_parameters() + elif hasattr(submodule, "_reset_parameters"): + submodule._reset_parameters() + else: + raise ValueError( + f"Module {submodule_fqn!r} does not have " + "reset_parameters or _reset_parameters." + ) + # Module.to_empty doesn't necessarily reuse Parameters so collects direct parameters again. + direct_parameters = list(submodule.named_parameters(recurse=False)) + + for local_param_name, parameter in direct_parameters: + param_fqn = f"{submodule_fqn}.{local_param_name}" if submodule_fqn else local_param_name + if hasattr(parameter, _CONTAINING_PARAMETER_GROUP_ATTR): + raise ValueError(f"Parameter {param_fqn!r} is already owned by an FSDP unit.") + parameters[param_fqn] = parameter + + for child_name, child_module in submodule.named_children(): + if isinstance(child_module, FsdpModule): + continue + child_fqn = f"{submodule_fqn}.{child_name}" if submodule_fqn else child_name + visit(child_module, child_fqn) + + visit(root_module, "") if not parameters: - raise ValueError("experimental fully_shard requires at least one unowned parameter.") + raise ValueError("fully_shard requires at least one unowned parameter.") return parameters @@ -524,30 +453,6 @@ def _get_parameter_owner(module: nn.Module, name: str) -> tuple[nn.Module, str]: return owner, parameter_name -def _reset_owned_meta_modules( - module: nn.Module, original_parameters: dict[str, nn.Parameter] -) -> None: - """Initialize modules that originally owned meta parameters.""" - modules_to_reset: dict[int, tuple[str, nn.Module]] = {} - for name, parameter in original_parameters.items(): - if not parameter.is_meta: - continue - owner_module, _ = _get_parameter_owner(module, name) - module_name = name.rsplit(".", 1)[0] if "." in name else "" - modules_to_reset.setdefault(id(owner_module), (module_name, owner_module)) - - for module_name, module in modules_to_reset.values(): - if hasattr(module, "reset_parameters"): - module.reset_parameters() - elif hasattr(module, "_reset_parameters"): - module._reset_parameters() - else: - raise ValueError( - f"[init_model_with_meta_device=True] Module {module_name!r} does not have " - "reset_parameters or _reset_parameters." - ) - - def _attach_mixin(module: nn.Module) -> None: if isinstance(module, FsdpModule): return diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_dbuffer.py b/tests/unit_tests/distributed/megatron_fsdp/test_dbuffer.py index a6ae56dc852..e33495229cd 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/test_dbuffer.py +++ b/tests/unit_tests/distributed/megatron_fsdp/test_dbuffer.py @@ -155,6 +155,35 @@ def test_constructor_allocates_local_buffer(distributed_setup): assert sharded_buffer.local_buffer.device == distributed_setup.device +@pytest.mark.distributed +def test_release_and_reallocate_storage_preserves_buffer_views(distributed_setup): + """DBuffer storage can be released and reallocated without replacing existing views.""" + mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) + buffer = DBuffer( + mesh=mesh, + placements=[Replicate()], + tensor_shapes=[torch.Size((4, 4))], + dtype=torch.float32, + device=distributed_setup.device, + ) + tensor_view = buffer.get_local_tensor(0) + buffer_data_ptr = buffer.local_buffer.data_ptr() + tensor_view_data_ptr = tensor_view.data_ptr() + + buffer.release_storage() + assert buffer.local_buffer.untyped_storage().nbytes() == 0 + + buffer.reallocate_storage() + assert ( + buffer.local_buffer.untyped_storage().nbytes() + == buffer.local_buffer.numel() * buffer.local_buffer.element_size() + ) + assert buffer.local_buffer.data_ptr() == buffer_data_ptr + assert tensor_view.data_ptr() == tensor_view_data_ptr + buffer.local_buffer.fill_(7.0) + torch.testing.assert_close(tensor_view, torch.full_like(tensor_view, 7.0)) + + @pytest.mark.distributed def test_from_local_reuses_required_local_buffer(distributed_setup): """DBuffer.from_local reuses caller-provided local storage without allocation.""" diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py index 9f7f6926773..dddbe5766cb 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py +++ b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py @@ -1,6 +1,6 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -"""Unit tests for the experimental minimal Megatron-FSDP path.""" +"""Unit tests for the minimal Megatron-FSDP path.""" import gc import os @@ -17,6 +17,7 @@ from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental import ( Flat, FsdpModule, + ParameterGroup, Placements, fully_shard, ) @@ -108,6 +109,32 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return SaveNonLeafWeightView.apply(x, self.weight.view_as(self.weight)) +class ConstantMetaModel(nn.Module): + """Model whose meta parameter is initialized by reset_parameters().""" + + def __init__(self) -> None: + super().__init__() + self.weight = nn.Parameter(torch.empty(4, 4, device="meta")) + + def reset_parameters(self) -> None: + """Initialize the weight to a deterministic value.""" + self.weight.fill_(3.0) + + +class MixedMetaRealModel(nn.Module): + """Model with unsupported mixed meta and real direct parameters.""" + + def __init__(self, device: torch.device) -> None: + super().__init__() + self.weight = nn.Parameter(torch.empty(4, 4, device="meta")) + self.bias = nn.Parameter(torch.ones(4, device=device)) + + def reset_parameters(self) -> None: + """Initialize parameters.""" + self.weight.fill_(1.0) + self.bias.zero_() + + def _flat_placements() -> Placements: return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()]) @@ -120,11 +147,15 @@ def _full_named_parameters(module: nn.Module) -> dict[str, torch.Tensor]: prefix = f"{module_name}." if module_name else "" for group in child.parameter_groups(): full_weight = group.main_weight.redistribute([Flat()]).allgather(0) - for index, name in enumerate(group.parameters): + for index, name in enumerate(group.sharded_parameters): result[f"{prefix}{name}"] = full_weight.get_tensor(index).detach().clone() return result +def _full_model_weight(group: ParameterGroup) -> torch.Tensor: + return group.model_weight.redistribute([Flat()]).allgather(0).get_tensor(0) + + @pytest.mark.distributed def test_experimental_fully_shard_train_step_matches_baseline(setup: DistributedSetup): """A minimal per-module FSDP train step should match single-rank SGD.""" @@ -171,8 +202,10 @@ def test_nested_fully_shard_excludes_child_owned_parameters(setup: DistributedSe fully_shard(model.inner, mesh=mesh, placements=_flat_placements()) fully_shard(model, mesh=mesh, placements=_flat_placements()) - inner_names = [name for group in model.inner.parameter_groups() for name in group.parameters] - outer_names = [name for group in model.parameter_groups() for name in group.parameters] + inner_names = [ + name for group in model.inner.parameter_groups() for name in group.sharded_parameters + ] + outer_names = [name for group in model.parameter_groups() for name in group.sharded_parameters] assert inner_names == ["weight"] assert outer_names == ["bias"] @@ -213,6 +246,32 @@ def test_default_main_buffer_dtypes_follow_policy_contract(setup: DistributedSet assert group.main_grad.local_buffer.dtype is torch.float64 +@pytest.mark.distributed +def test_full_mesh_is_distinct_from_dbuffer_dp_submesh(setup: DistributedSetup): + """ParameterGroup keeps the full mesh while DBuffers use only DP axes.""" + if setup.world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + + mesh = init_device_mesh(setup.device.type, (1, setup.world_size), mesh_dim_names=("tp", "dp")) + placements = Placements( + dp_axes=["dp"], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()] + ) + model = nn.Linear(4, 4, bias=False).to(setup.device) + + fully_shard(model, mesh=mesh, placements=placements) + + (group,) = model.parameter_groups() + assert group.mesh is mesh + assert group.dp_mesh.ndim == 1 + assert group.dp_mesh.size() == setup.world_size + assert group.main_weight.mesh is group.dp_mesh + assert group.model_weight.mesh is group.dp_mesh + assert group.main_grad is not None + assert group.main_grad.mesh is group.dp_mesh + assert isinstance(model.weight, DTensor) + assert model.weight.device_mesh is group.dp_mesh + + @pytest.mark.distributed def test_sharded_parameter_contract_uses_dtensors(setup: DistributedSetup): """Resting sharded parameters and gradients should be DTensors.""" @@ -221,9 +280,12 @@ def test_sharded_parameter_contract_uses_dtensors(setup: DistributedSetup): mesh = init_device_mesh(setup.device.type, (setup.world_size,)) model = nn.Linear(4, 4, bias=False).to(setup.device) + original_weight = model.weight fully_shard(model, mesh=mesh, placements=_flat_placements()) + (group,) = model.parameter_groups() + assert group.unsharded_parameters["weight"] is original_weight assert isinstance(model.weight, DTensor) assert isinstance(model.weight.data, DTensor) model(torch.randn(2, 4, device=setup.device)).sum().backward() @@ -238,15 +300,30 @@ def test_meta_parameters_initialize_with_reset_parameters(setup: DistributedSetu pytest.skip("This test requires at least 2 ranks.") mesh = init_device_mesh(setup.device.type, (setup.world_size,)) - model = nn.Linear(4, 4, bias=False, device="meta") + model = ConstantMetaModel() - fully_shard(model, mesh=mesh, placements=_flat_placements(), init_model_with_meta_device=True) + fully_shard(model, mesh=mesh, placements=_flat_placements()) assert isinstance(model.weight, DTensor) assert not model.weight.to_local().is_meta (group,) = model.parameter_groups() - assert not group.main_weight.local_buffer.is_meta - assert group.main_weight.local_buffer.numel() > 0 + assert not group.model_weight.local_buffer.is_meta + assert group.model_weight.local_buffer.numel() > 0 + full_weight = _full_model_weight(group) + torch.testing.assert_close(full_weight, torch.full_like(full_weight, 3.0)) + + +@pytest.mark.distributed +def test_meta_parameters_reject_mixed_direct_parameters(setup: DistributedSetup): + """A module cannot mix meta and real direct parameters when reset is required.""" + if setup.world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + + mesh = init_device_mesh(setup.device.type, (setup.world_size,)) + model = MixedMetaRealModel(setup.device) + + with pytest.raises(ValueError, match="mixes meta and non-meta direct parameters"): + fully_shard(model, mesh=mesh, placements=_flat_placements()) @pytest.mark.distributed @@ -265,19 +342,19 @@ def test_non_leaf_parameter_view_survives_storage_resize(setup: DistributedSetup x = torch.randn(8, device=setup.device, requires_grad=True) loss = model(x).sum() - assert group._full_weight is not None - assert group._full_weight.local_buffer.untyped_storage().nbytes() == 0 + assert group._unsharded_model_weight is not None + assert group._unsharded_model_weight.local_buffer.untyped_storage().nbytes() == 0 loss.backward() assert group.main_grad is not None - assert group._full_weight is not None - assert group._full_weight.local_buffer.untyped_storage().nbytes() == 0 + assert group._unsharded_model_weight is not None + assert group._unsharded_model_weight.local_buffer.untyped_storage().nbytes() == 0 @pytest.mark.distributed def test_experimental_fully_shard_reduces_peak_training_memory(setup: DistributedSetup): - """Per-layer experimental FSDP should reduce peak CUDA memory during a train step.""" + """Per-layer FSDP should reduce peak CUDA memory during a train step.""" if setup.world_size < 2: pytest.skip("This test requires at least 2 ranks.") if setup.device.type != "cuda": @@ -285,7 +362,7 @@ def test_experimental_fully_shard_reduces_peak_training_memory(setup: Distribute mesh = init_device_mesh(setup.device.type, (setup.world_size,)) dim = 1024 - layers = 8 + layers = 16 batch = 8 torch.manual_seed(4321) @@ -309,6 +386,8 @@ def test_experimental_fully_shard_reduces_peak_training_memory(setup: Distribute ) for layer in model: fully_shard(layer, mesh=mesh, placements=_flat_placements()) + gc.collect() + torch.cuda.empty_cache() x = torch.randn(batch, dim, device=setup.device) torch.cuda.reset_peak_memory_stats(setup.device) From afbe9fa220450bacda869c7a83aba512c053bec6 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sun, 24 May 2026 07:21:34 +0000 Subject: [PATCH 03/23] Refine experimental FSDP buffer APIs Add out= support to DBuffer redistribution and primitive communication ops, keeping axis inference in redistribute only. Use preallocated model and gradient buffers in the minimal FSDP path where possible, including direct first-gradient reduce-scatter into main_grad. Update DBuffer and experimental FSDP tests for AVG reductions, explicit primitive axes, storage reuse, and gradient accumulation behavior. Signed-off-by: Jingyue Wu --- .../megatron_fsdp/experimental/fully_shard.py | 167 +++++++++--------- .../distributed/megatron_fsdp/test_dbuffer.py | 18 ++ .../test_experimental_fully_shard.py | 60 ++++++- 3 files changed, 149 insertions(+), 96 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py index dba988ea4d9..c67a6a19335 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py @@ -15,9 +15,10 @@ """Minimal per-module Megatron-FSDP implementation.""" import dataclasses -from collections.abc import Callable, Sequence +from collections.abc import Callable, Iterable, Sequence import torch +import torch.distributed as dist from torch import nn from torch.distributed import DeviceMesh @@ -48,6 +49,20 @@ def __post_init__(self) -> None: raise ValueError(f"Expected {axis_count} {name} placements, got {len(placements)}.") +def has_grad(parameters: Iterable[nn.Parameter]) -> bool: + """Return whether parameters have gradients, requiring all-or-none state.""" + has_any_grad = False + has_any_missing_grad = False + for parameter in parameters: + if parameter.grad is None: + has_any_missing_grad = True + else: + has_any_grad = True + if has_any_grad and has_any_missing_grad: + raise RuntimeError("FSDP sharded gradients must be either all set or all None.") + return has_any_grad + + class ParameterGroup: """A dtype and requires-grad homogeneous group of FSDP-owned parameters.""" @@ -55,7 +70,6 @@ class ParameterGroup: sharded_parameters: dict[str, nn.Parameter] unsharded_parameters: dict[str, nn.Parameter] mesh: DeviceMesh - dp_mesh: DeviceMesh dtype: torch.dtype requires_grad: bool main_weight: DBuffer @@ -76,22 +90,26 @@ def __init__( Args: owning_module: Closest FSDP root module that owns this parameter group. parameters: Root-module-relative FQNs and their parameters. - mesh: Full device mesh. DBuffer storage is built on the DP submesh. + mesh: Device mesh used for all DBuffer storage in this version. placements: Parameter, gradient, and optimizer placements. mixed_precision_policy: Precision policy for main weights and gradients. """ + # --------------------------------------------------------------------- + # Validate group inputs and mesh contract + # --------------------------------------------------------------------- if not parameters: raise ValueError("ParameterGroup requires at least one parameter.") - dp_mesh = _dp_submesh(mesh, placements.dp_axes) + axis_indices = tuple(_axis_index(mesh, axis) for axis in placements.dp_axes) + assert axis_indices == tuple(range(mesh.ndim)), ( + "FSDP requires dp_axes to match every mesh axis in mesh order for now." + ) - # Python dicts preserve insertion order, so values() defines the stable - # tensor order used by each DBuffer built from this group. + # --------------------------------------------------------------------- + # Record shared metadata + # --------------------------------------------------------------------- self.owning_module = owning_module - self.sharded_parameters = {} - self.unsharded_parameters = {} self.mesh = mesh - self.dp_mesh = dp_mesh first_parameter = next(iter(parameters.values())) self.dtype = first_parameter.dtype self.requires_grad = first_parameter.requires_grad @@ -109,32 +127,38 @@ def __init__( f"Expected parameter {name!r} to have requires_grad={self.requires_grad}, " f"got {parameter.requires_grad}." ) - main_weight_dtype = mixed_precision_policy.main_params_dtype - if main_weight_dtype is None: - raise ValueError( - "FSDP requires a main weight dtype; set MixedPrecisionPolicy.main_params_dtype." - ) - main_grad_dtype = mixed_precision_policy.main_grads_dtype - if main_grad_dtype is None: - main_grad_dtype = self.dtype + + # --------------------------------------------------------------------- + # Initialize DBuffers from original parameter values + # --------------------------------------------------------------------- + # Python dicts preserve insertion order, so values() defines the stable + # tensor order used by each DBuffer built from this group. + self._unsharded_model_weight = DBuffer.distribute_tensors( + parameters.values(), + mesh=self.mesh, + placements=[Replicate()] * self.mesh.ndim, + ) # Scratch initialization starts from model weights. Checkpoint loading will # eventually initialize main weights first and derive model weights from them. - self.model_weight = DBuffer.distribute_tensors( - [parameter.detach().contiguous() for parameter in parameters.values()], - mesh=self.dp_mesh, - placements=placements.parameter, + self.model_weight = self._unsharded_model_weight.redistribute(placements.parameter) + model_weight_storage = self.model_weight.local_buffer.untyped_storage() + unsharded_model_weight_storage = self._unsharded_model_weight.local_buffer.untyped_storage() + assert model_weight_storage.data_ptr() != unsharded_model_weight_storage.data_ptr(), ( + "model_weight must not share storage with _unsharded_model_weight because " + "_unsharded_model_weight storage is resized and released after resharding." ) + + main_weight_dtype = mixed_precision_policy.main_params_dtype or torch.float32 self.main_weight = DBuffer.distribute_tensors( - [ - parameter.detach().to(dtype=main_weight_dtype).contiguous() - for parameter in parameters.values() - ], - mesh=self.dp_mesh, + [parameter.to(dtype=main_weight_dtype) for parameter in parameters.values()], + mesh=self.mesh, placements=placements.optimizer, ) + + main_grad_dtype = mixed_precision_policy.main_grads_dtype or self.dtype self.main_grad = ( DBuffer( - mesh=self.dp_mesh, + mesh=self.mesh, placements=placements.gradient, tensor_shapes=self.main_weight.layout.tensor_shapes, dtype=main_grad_dtype, @@ -143,14 +167,12 @@ def __init__( if self.requires_grad else None ) - self._unsharded_model_weight = DBuffer( - mesh=self.dp_mesh, - placements=[Replicate()] * self.dp_mesh.ndim, - tensor_shapes=self.model_weight.layout.tensor_shapes, - dtype=self.model_weight.local_buffer.dtype, - device=self.model_weight.local_buffer.device, - ) + # --------------------------------------------------------------------- + # Build parameter maps for module swapping + # --------------------------------------------------------------------- + self.sharded_parameters = {} + self.unsharded_parameters = {} for index, (name, parameter) in enumerate(parameters.items()): parameter.data = self._unsharded_model_weight.get_tensor(index) parameter.grad = None @@ -161,6 +183,10 @@ def __init__( ) setattr(sharded_parameter, _CONTAINING_PARAMETER_GROUP_ATTR, self) self.sharded_parameters[name] = sharded_parameter + + # --------------------------------------------------------------------- + # Install resting sharded parameters + # --------------------------------------------------------------------- self._switch_to_sharded_parameters() self._unsharded_model_weight.release_storage() @@ -183,7 +209,9 @@ def unshard_parameters(self) -> None: with torch.autograd._unsafe_preserve_version_counter( self._unsharded_model_weight.local_buffer ): - self.model_weight.fully_allgather_into(self._unsharded_model_weight) + self.model_weight.redistribute( + self._unsharded_model_weight.placements, out=self._unsharded_model_weight + ) self._switch_to_unsharded_parameters() def reshard_parameters(self) -> None: @@ -191,47 +219,34 @@ def reshard_parameters(self) -> None: self._switch_to_sharded_parameters() self._unsharded_model_weight.release_storage() - def reduce_gradients(self, average: bool = True) -> None: - """Reduce full local gradients into the persistent sharded gradient buffer.""" - if not self.requires_grad: - return + def reduce_gradients(self) -> None: + """Reduce full local gradients into sharded parameter gradients.""" assert self.main_grad is not None - accumulate = self.has_sharded_grad() full_grads: list[torch.Tensor] = [] for name, parameter in self.unsharded_parameters.items(): if parameter.grad is None: raise RuntimeError(f"Missing gradient for FSDP parameter {name!r}.") - full_grads.append( - parameter.grad.detach().to(dtype=self.main_grad.local_buffer.dtype).contiguous() - ) + full_grads.append(parameter.grad.to(dtype=self.main_grad.local_buffer.dtype)) partial_grad = DBuffer.distribute_tensors( - full_grads, mesh=self.dp_mesh, placements=[Partial()] * self.dp_mesh.ndim + full_grads, + mesh=self.mesh, + placements=[Partial(dist.ReduceOp.AVG)] * self.mesh.ndim, ) - reduced_grad = partial_grad.redistribute(self.main_grad.placements) - if average: - reduced_grad.local_buffer.div_(self.dp_mesh.size()) - if accumulate: + sharded_parameters = self.sharded_parameters.values() + if has_grad(sharded_parameters): + reduced_grad = partial_grad.redistribute(self.main_grad.placements) self.main_grad.local_buffer.add_(reduced_grad.local_buffer) else: - self.main_grad.local_buffer.copy_(reduced_grad.local_buffer) + partial_grad.redistribute(self.main_grad.placements, out=self.main_grad) + for index, parameter in enumerate(sharded_parameters): + parameter.grad = self.main_grad.get_dtensor(index) for parameter in self.unsharded_parameters.values(): parameter.grad = None - self.install_sharded_gradients() - def has_sharded_grad(self) -> bool: - """Return whether persistent sharded gradients are currently materialized.""" - return any(parameter.grad is not None for parameter in self.sharded_parameters.values()) - - def install_sharded_gradients(self) -> None: - """Install sharded DTensor gradients backed by main_grad.""" - if self.main_grad is None: - return - for index, parameter in enumerate(self.sharded_parameters.values()): - parameter.grad = self.main_grad.get_dtensor(index) class FsdpModule: """Mixin attached to modules managed by the minimal FSDP path.""" @@ -239,7 +254,7 @@ class FsdpModule: _parameter_groups: tuple[ParameterGroup, ...] _ready_grad_params: set[nn.Parameter] _registered_grad_param_ids: set[int] - _trainable_param_count: int + num_trainable_params: int def __init__( self, mesh: DeviceMesh, placements: Placements, mixed_precision_policy: MixedPrecisionPolicy @@ -259,7 +274,7 @@ def __init__( self._parameter_groups = tuple(parameter_groups) self._ready_grad_params: set[nn.Parameter] = set() self._registered_grad_param_ids: set[int] = set() - self._trainable_param_count = sum( + self.num_trainable_params = sum( len(group.sharded_parameters) for group in self._parameter_groups if group.requires_grad ) self._register_hooks() @@ -284,7 +299,7 @@ def _register_grad_hooks(self) -> None: def _make_grad_hook(self, parameter: nn.Parameter) -> Callable[[nn.Parameter], None]: def grad_hook(_parameter: nn.Parameter) -> None: self._ready_grad_params.add(parameter) - if len(self._ready_grad_params) == self._trainable_param_count: + if len(self._ready_grad_params) == self.num_trainable_params: self.post_backward() return grad_hook @@ -309,9 +324,9 @@ def pre_backward(self) -> None: def post_backward(self) -> None: """Reduce gradients and return parameters to their sharded resting state.""" for group in self._parameter_groups: - group.reduce_gradients() + if group.requires_grad: + group.reduce_gradients() group.reshard_parameters() - group.install_sharded_gradients() self._ready_grad_params.clear() def parameter_groups(self) -> tuple[ParameterGroup, ...]: @@ -364,28 +379,6 @@ def _axis_index(mesh: DeviceMesh, axis: MeshAxis) -> int: return dim_names.index(axis) -def _dp_submesh(mesh: DeviceMesh, dp_axes: Sequence[MeshAxis]) -> DeviceMesh: - if not dp_axes: - raise ValueError("FSDP requires at least one DP mesh axis.") - - axis_indices = tuple(_axis_index(mesh, axis) for axis in dp_axes) - if len(set(axis_indices)) != len(axis_indices): - raise ValueError(f"Duplicate DP mesh axes are not allowed: {tuple(dp_axes)!r}.") - - if axis_indices == tuple(range(mesh.ndim)): - return mesh - - dim_names = mesh.mesh_dim_names - if dim_names is None: - raise ValueError( - "Slicing a DP submesh from a full mesh requires named mesh dimensions unless " - "dp_axes covers every mesh axis in mesh order." - ) - - dp_axis_names = tuple(dim_names[index] for index in axis_indices) - return mesh[dp_axis_names[0] if len(dp_axis_names) == 1 else dp_axis_names] - - def _mesh_device(mesh: DeviceMesh) -> torch.device: if mesh.device_type == "cuda": return torch.device("cuda", torch.cuda.current_device()) diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_dbuffer.py b/tests/unit_tests/distributed/megatron_fsdp/test_dbuffer.py index e33495229cd..6bd9bb8ab82 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/test_dbuffer.py +++ b/tests/unit_tests/distributed/megatron_fsdp/test_dbuffer.py @@ -232,6 +232,24 @@ def test_distribute_tensors_moves_inputs_to_mesh_device(distributed_setup): ) +@pytest.mark.distributed +def test_distribute_tensors_detaches_and_contiguizes_inputs(distributed_setup): + """distribute_tensors treats input tensors as detached contiguous values.""" + mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) + parameter = torch.nn.Parameter( + torch.arange(12, dtype=torch.float32, device=distributed_setup.device).view(3, 4).t() + ) + + buffer = DBuffer.distribute_tensors([parameter], mesh, [Replicate()]) + + assert not parameter.is_contiguous() + assert buffer.get_local_tensor(0).is_contiguous() + assert not buffer.local_buffer.requires_grad + torch.testing.assert_close( + buffer.get_local_tensor(0), parameter.detach().contiguous(), rtol=0, atol=0 + ) + + @pytest.mark.distributed def test_sharded_allgather_round_trip(distributed_setup): """Sharded buffers round-trip through all-gather as contiguous tensor fragments.""" diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py index dddbe5766cb..7c97d200779 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py +++ b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py @@ -19,6 +19,7 @@ FsdpModule, ParameterGroup, Placements, + Replicate, fully_shard, ) @@ -247,14 +248,17 @@ def test_default_main_buffer_dtypes_follow_policy_contract(setup: DistributedSet @pytest.mark.distributed -def test_full_mesh_is_distinct_from_dbuffer_dp_submesh(setup: DistributedSetup): - """ParameterGroup keeps the full mesh while DBuffers use only DP axes.""" +def test_full_mesh_is_used_for_dbuffers(setup: DistributedSetup): + """This version uses the full mesh for DBuffer storage.""" if setup.world_size < 2: pytest.skip("This test requires at least 2 ranks.") mesh = init_device_mesh(setup.device.type, (1, setup.world_size), mesh_dim_names=("tp", "dp")) placements = Placements( - dp_axes=["dp"], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()] + dp_axes=["tp", "dp"], + parameter=[Replicate(), Flat()], + gradient=[Replicate(), Flat()], + optimizer=[Replicate(), Flat()], ) model = nn.Linear(4, 4, bias=False).to(setup.device) @@ -262,14 +266,12 @@ def test_full_mesh_is_distinct_from_dbuffer_dp_submesh(setup: DistributedSetup): (group,) = model.parameter_groups() assert group.mesh is mesh - assert group.dp_mesh.ndim == 1 - assert group.dp_mesh.size() == setup.world_size - assert group.main_weight.mesh is group.dp_mesh - assert group.model_weight.mesh is group.dp_mesh + assert group.main_weight.mesh is mesh + assert group.model_weight.mesh is mesh assert group.main_grad is not None - assert group.main_grad.mesh is group.dp_mesh + assert group.main_grad.mesh is mesh assert isinstance(model.weight, DTensor) - assert model.weight.device_mesh is group.dp_mesh + assert model.weight.device_mesh is mesh @pytest.mark.distributed @@ -291,6 +293,30 @@ def test_sharded_parameter_contract_uses_dtensors(setup: DistributedSetup): model(torch.randn(2, 4, device=setup.device)).sum().backward() assert isinstance(model.weight, DTensor) assert isinstance(model.weight.grad, DTensor) + assert original_weight.grad is None + + +@pytest.mark.distributed +def test_backward_averages_across_dp_and_accumulates_across_calls(setup: DistributedSetup): + """Each backward averages over DP ranks; repeated backwards accumulate by summing.""" + if setup.world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + + mesh = init_device_mesh(setup.device.type, (setup.world_size,)) + model = nn.Linear(1, setup.world_size, bias=False).to(setup.device) + with torch.no_grad(): + model.weight.fill_(1.0) + + fully_shard(model, mesh=mesh, placements=_flat_placements()) + + x = torch.full((1, 1), float(setup.rank + 1), device=setup.device) + model(x).sum().backward() + model(x).sum().backward() + + assert isinstance(model.weight.grad, DTensor) + local_grad = model.weight.grad.to_local() + expected = torch.full_like(local_grad, float(setup.world_size + 1)) + torch.testing.assert_close(local_grad, expected, rtol=0, atol=0) @pytest.mark.distributed @@ -313,6 +339,22 @@ def test_meta_parameters_initialize_with_reset_parameters(setup: DistributedSetu torch.testing.assert_close(full_weight, torch.full_like(full_weight, 3.0)) +@pytest.mark.distributed +def test_model_weight_alias_is_rejected(setup: DistributedSetup): + """Model weights cannot alias the unsharded parameter buffer.""" + if setup.world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + + mesh = init_device_mesh(setup.device.type, (setup.world_size,)) + placements = Placements( + dp_axes=[0], parameter=[Replicate()], gradient=[Flat()], optimizer=[Flat()] + ) + model = nn.Linear(4, 4, bias=False).to(setup.device) + + with pytest.raises(AssertionError, match="storage is resized and released"): + fully_shard(model, mesh=mesh, placements=placements) + + @pytest.mark.distributed def test_meta_parameters_reject_mixed_direct_parameters(setup: DistributedSetup): """A module cannot mix meta and real direct parameters when reset is required.""" From 63987db411e8e0b1d79d6dff69d7a294cd89cff2 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Mon, 25 May 2026 03:30:38 +0000 Subject: [PATCH 04/23] Refine experimental FSDP gradient contract Use ordered parameter tuples for FSDP parameter swapping and keep the sharded data path aligned with main_weight storage. Set grad_dtype for FSDP-managed parameters so BF16 main gradients can be reduced without pre-reduce casts, and update tests to verify sharded parameter data and grad backing buffers. Clean up hook naming, local gradient accumulation handling, and memory/test assertions for the minimal experimental path. Signed-off-by: Jingyue Wu --- .../megatron_fsdp/experimental/fully_shard.py | 127 ++++++----- .../test_experimental_fully_shard.py | 202 ++++++------------ 2 files changed, 132 insertions(+), 197 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py index c67a6a19335..a2ffb9f3e83 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py @@ -49,26 +49,13 @@ def __post_init__(self) -> None: raise ValueError(f"Expected {axis_count} {name} placements, got {len(placements)}.") -def has_grad(parameters: Iterable[nn.Parameter]) -> bool: - """Return whether parameters have gradients, requiring all-or-none state.""" - has_any_grad = False - has_any_missing_grad = False - for parameter in parameters: - if parameter.grad is None: - has_any_missing_grad = True - else: - has_any_grad = True - if has_any_grad and has_any_missing_grad: - raise RuntimeError("FSDP sharded gradients must be either all set or all None.") - return has_any_grad - - class ParameterGroup: """A dtype and requires-grad homogeneous group of FSDP-owned parameters.""" owning_module: nn.Module - sharded_parameters: dict[str, nn.Parameter] - unsharded_parameters: dict[str, nn.Parameter] + parameter_names: tuple[str, ...] + sharded_parameters: tuple[nn.Parameter, ...] + unsharded_parameters: tuple[nn.Parameter, ...] mesh: DeviceMesh dtype: torch.dtype requires_grad: bool @@ -110,6 +97,7 @@ def __init__( # --------------------------------------------------------------------- self.owning_module = owning_module self.mesh = mesh + self.parameter_names = tuple(parameters) first_parameter = next(iter(parameters.values())) self.dtype = first_parameter.dtype self.requires_grad = first_parameter.requires_grad @@ -131,8 +119,8 @@ def __init__( # --------------------------------------------------------------------- # Initialize DBuffers from original parameter values # --------------------------------------------------------------------- - # Python dicts preserve insertion order, so values() defines the stable - # tensor order used by each DBuffer built from this group. + # Python dicts preserve insertion order, so parameter_names and + # parameters.values() define the same stable DBuffer tensor order. self._unsharded_model_weight = DBuffer.distribute_tensors( parameters.values(), mesh=self.mesh, @@ -150,7 +138,7 @@ def __init__( main_weight_dtype = mixed_precision_policy.main_params_dtype or torch.float32 self.main_weight = DBuffer.distribute_tensors( - [parameter.to(dtype=main_weight_dtype) for parameter in parameters.values()], + (parameter.to(dtype=main_weight_dtype) for parameter in parameters.values()), mesh=self.mesh, placements=placements.optimizer, ) @@ -169,20 +157,28 @@ def __init__( ) # --------------------------------------------------------------------- - # Build parameter maps for module swapping + # Build parameter tuples for module swapping # --------------------------------------------------------------------- - self.sharded_parameters = {} - self.unsharded_parameters = {} - for index, (name, parameter) in enumerate(parameters.items()): + sharded_parameters: list[nn.Parameter] = [] + unsharded_parameters: list[nn.Parameter] = [] + grad_dtype = self.main_grad.local_buffer.dtype if self.requires_grad else None + for index, parameter in enumerate(parameters.values()): parameter.data = self._unsharded_model_weight.get_tensor(index) parameter.grad = None + if grad_dtype: + parameter.grad_dtype = grad_dtype setattr(parameter, _CONTAINING_PARAMETER_GROUP_ATTR, self) - self.unsharded_parameters[name] = parameter + unsharded_parameters.append(parameter) + sharded_parameter = nn.Parameter( self.main_weight.get_dtensor(index), requires_grad=parameter.requires_grad ) + if grad_dtype: + sharded_parameter.grad_dtype = grad_dtype setattr(sharded_parameter, _CONTAINING_PARAMETER_GROUP_ATTR, self) - self.sharded_parameters[name] = sharded_parameter + sharded_parameters.append(sharded_parameter) + self.sharded_parameters = tuple(sharded_parameters) + self.unsharded_parameters = tuple(unsharded_parameters) # --------------------------------------------------------------------- # Install resting sharded parameters @@ -190,8 +186,8 @@ def __init__( self._switch_to_sharded_parameters() self._unsharded_model_weight.release_storage() - def _set_module_parameters(self, parameters: dict[str, nn.Parameter]) -> None: - for name, parameter in parameters.items(): + def _set_module_parameters(self, parameters: tuple[nn.Parameter, ...]) -> None: + for name, parameter in zip(self.parameter_names, parameters, strict=True): module, parameter_name = _get_parameter_owner(self.owning_module, name) module._parameters[parameter_name] = parameter @@ -223,28 +219,47 @@ def reduce_gradients(self) -> None: """Reduce full local gradients into sharded parameter gradients.""" assert self.main_grad is not None - full_grads: list[torch.Tensor] = [] - for name, parameter in self.unsharded_parameters.items(): + def has_grad(parameters: Iterable[nn.Parameter]) -> bool: + has_any_grad = False + has_any_missing_grad = False + for parameter in parameters: + if parameter.grad is None: + has_any_missing_grad = True + else: + has_any_grad = True + if has_any_grad and has_any_missing_grad: + raise RuntimeError("FSDP sharded gradients must be either all set or all None.") + return has_any_grad + + grads: list[torch.Tensor] = [] + for name, parameter in zip(self.parameter_names, self.unsharded_parameters, strict=True): if parameter.grad is None: raise RuntimeError(f"Missing gradient for FSDP parameter {name!r}.") - full_grads.append(parameter.grad.to(dtype=self.main_grad.local_buffer.dtype)) + assert parameter.grad.dtype == self.main_grad.local_buffer.dtype, ( + "FSDP unsharded parameter grad dtype should be guaranteed by grad_dtype: " + f"parameter {name!r} has grad dtype {parameter.grad.dtype}, " + f"expected main_grad dtype {self.main_grad.local_buffer.dtype}." + ) + grads.append(parameter.grad) partial_grad = DBuffer.distribute_tensors( - full_grads, + grads, mesh=self.mesh, placements=[Partial(dist.ReduceOp.AVG)] * self.mesh.ndim, ) - sharded_parameters = self.sharded_parameters.values() - if has_grad(sharded_parameters): + # zero_grad(set_to_none=True) clears sharded parameter grads, so the next + # backward can reduce directly into main_grad. zero_grad(set_to_none=False) + # leaves sharded grads installed, so this backward accumulates into main_grad. + if has_grad(self.sharded_parameters): reduced_grad = partial_grad.redistribute(self.main_grad.placements) self.main_grad.local_buffer.add_(reduced_grad.local_buffer) else: partial_grad.redistribute(self.main_grad.placements, out=self.main_grad) - for index, parameter in enumerate(sharded_parameters): + for index, parameter in enumerate(self.sharded_parameters): parameter.grad = self.main_grad.get_dtensor(index) - for parameter in self.unsharded_parameters.values(): + for parameter in self.unsharded_parameters: parameter.grad = None @@ -252,9 +267,8 @@ class FsdpModule: """Mixin attached to modules managed by the minimal FSDP path.""" _parameter_groups: tuple[ParameterGroup, ...] - _ready_grad_params: set[nn.Parameter] - _registered_grad_param_ids: set[int] - num_trainable_params: int + _ready_grad_parameters: set[nn.Parameter] + num_training_parameters: int def __init__( self, mesh: DeviceMesh, placements: Placements, mixed_precision_policy: MixedPrecisionPolicy @@ -272,9 +286,8 @@ def __init__( for group_parameters in _group_parameters(owned_parameters) ] self._parameter_groups = tuple(parameter_groups) - self._ready_grad_params: set[nn.Parameter] = set() - self._registered_grad_param_ids: set[int] = set() - self.num_trainable_params = sum( + self._ready_grad_parameters = set() + self.num_training_parameters = sum( len(group.sharded_parameters) for group in self._parameter_groups if group.requires_grad ) self._register_hooks() @@ -283,33 +296,29 @@ def _register_hooks(self) -> None: self.register_forward_pre_hook(lambda _module, _args: self.pre_forward()) self.register_forward_hook(lambda _module, _args, _output: self.post_forward()) self.register_full_backward_pre_hook(lambda _module, _grad_output: self.pre_backward()) - - def _register_grad_hooks(self) -> None: - """Register post-accumulate hooks on full-size autograd leaf parameters.""" + # Gradient reduction is parameter-completion based: once every owned + # Parameter has accumulated its grad, this FSDP unit can reduce and + # reshard. Module full-backward hooks can fire before that when module + # inputs do not require grad. for group in self._parameter_groups: if not group.requires_grad: continue - for parameter in group.unsharded_parameters.values(): - parameter_id = id(parameter) - if parameter_id in self._registered_grad_param_ids: - continue + for parameter in group.unsharded_parameters: parameter.register_post_accumulate_grad_hook(self._make_grad_hook(parameter)) - self._registered_grad_param_ids.add(parameter_id) def _make_grad_hook(self, parameter: nn.Parameter) -> Callable[[nn.Parameter], None]: def grad_hook(_parameter: nn.Parameter) -> None: - self._ready_grad_params.add(parameter) - if len(self._ready_grad_params) == self.num_trainable_params: + self._ready_grad_parameters.add(parameter) + if len(self._ready_grad_parameters) == self.num_training_parameters: self.post_backward() return grad_hook def pre_forward(self) -> None: """Prepare full parameters for forward compute.""" - self._ready_grad_params.clear() + self._ready_grad_parameters.clear() for group in self._parameter_groups: group.unshard_parameters() - self._register_grad_hooks() def post_forward(self) -> None: """Return parameters to their sharded resting state after forward compute.""" @@ -327,7 +336,7 @@ def post_backward(self) -> None: if group.requires_grad: group.reduce_gradients() group.reshard_parameters() - self._ready_grad_params.clear() + self._ready_grad_parameters.clear() def parameter_groups(self) -> tuple[ParameterGroup, ...]: """Return parameter groups owned by this FSDP unit.""" @@ -413,11 +422,13 @@ def visit(submodule: nn.Module, submodule_fqn: str) -> None: # Module.to_empty doesn't necessarily reuse Parameters so collects direct parameters again. direct_parameters = list(submodule.named_parameters(recurse=False)) - for local_param_name, parameter in direct_parameters: - param_fqn = f"{submodule_fqn}.{local_param_name}" if submodule_fqn else local_param_name + for local_parameter_name, parameter in direct_parameters: + parameter_fqn = ( + f"{submodule_fqn}.{local_parameter_name}" if submodule_fqn else local_parameter_name + ) if hasattr(parameter, _CONTAINING_PARAMETER_GROUP_ATTR): - raise ValueError(f"Parameter {param_fqn!r} is already owned by an FSDP unit.") - parameters[param_fqn] = parameter + raise ValueError(f"Parameter {parameter_fqn!r} is already owned by an FSDP unit.") + parameters[parameter_fqn] = parameter for child_name, child_module in submodule.named_children(): if isinstance(child_module, FsdpModule): diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py index 7c97d200779..eaedb85a9d6 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py +++ b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py @@ -2,7 +2,7 @@ """Unit tests for the minimal Megatron-FSDP path.""" -import gc +import logging import os from collections.abc import Iterator from dataclasses import dataclass @@ -17,12 +17,13 @@ from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental import ( Flat, FsdpModule, - ParameterGroup, Placements, Replicate, fully_shard, ) +logger = logging.getLogger(__name__) + @dataclass(frozen=True) class DistributedSetup: @@ -107,7 +108,10 @@ def __init__(self) -> None: def forward(self, x: torch.Tensor) -> torch.Tensor: """Run using a non-leaf view of the parameter.""" - return SaveNonLeafWeightView.apply(x, self.weight.view_as(self.weight)) + weight_view = self.weight.view_as(self.weight) + assert self.weight.is_leaf + assert not weight_view.is_leaf + return SaveNonLeafWeightView.apply(x, weight_view) class ConstantMetaModel(nn.Module): @@ -122,47 +126,32 @@ def reset_parameters(self) -> None: self.weight.fill_(3.0) -class MixedMetaRealModel(nn.Module): - """Model with unsupported mixed meta and real direct parameters.""" - - def __init__(self, device: torch.device) -> None: - super().__init__() - self.weight = nn.Parameter(torch.empty(4, 4, device="meta")) - self.bias = nn.Parameter(torch.ones(4, device=device)) - - def reset_parameters(self) -> None: - """Initialize parameters.""" - self.weight.fill_(1.0) - self.bias.zero_() - - def _flat_placements() -> Placements: return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()]) -def _full_named_parameters(module: nn.Module) -> dict[str, torch.Tensor]: - result = {} - for module_name, child in module.named_modules(): - if not isinstance(child, FsdpModule): - continue - prefix = f"{module_name}." if module_name else "" - for group in child.parameter_groups(): - full_weight = group.main_weight.redistribute([Flat()]).allgather(0) - for index, name in enumerate(group.sharded_parameters): - result[f"{prefix}{name}"] = full_weight.get_tensor(index).detach().clone() - return result - - -def _full_model_weight(group: ParameterGroup) -> torch.Tensor: - return group.model_weight.redistribute([Flat()]).allgather(0).get_tensor(0) +def _mb(num_bytes: int) -> str: + return f"{num_bytes / 1024**2:.2f} MB" @pytest.mark.distributed -def test_experimental_fully_shard_train_step_matches_baseline(setup: DistributedSetup): +def test_fully_shard_train_step_matches_baseline(setup: DistributedSetup): """A minimal per-module FSDP train step should match single-rank SGD.""" if setup.world_size < 2: pytest.skip("This test requires at least 2 ranks.") + def full_named_parameters(module: nn.Module) -> dict[str, torch.Tensor]: + result = {} + for module_name, child in module.named_modules(): + if not isinstance(child, FsdpModule): + continue + prefix = f"{module_name}." if module_name else "" + for group in child.parameter_groups(): + full_weight = group.main_weight.allgather(0) + for index, name in enumerate(group.parameter_names): + result[f"{prefix}{name}"] = full_weight.get_tensor(index).detach().clone() + return result + mesh = init_device_mesh(setup.device.type, (setup.world_size,)) torch.manual_seed(1234) baseline = TinyModel().to(setup.device) @@ -186,9 +175,9 @@ def test_experimental_fully_shard_train_step_matches_baseline(setup: Distributed loss.backward() optimizer.step() - sharded_params = _full_named_parameters(model) + full_params = full_named_parameters(model) for name, expected in baseline.named_parameters(): - torch.testing.assert_close(sharded_params[name], expected, rtol=1e-5, atol=1e-6) + torch.testing.assert_close(full_params[name], expected, rtol=1e-5, atol=1e-6) @pytest.mark.distributed @@ -203,10 +192,8 @@ def test_nested_fully_shard_excludes_child_owned_parameters(setup: DistributedSe fully_shard(model.inner, mesh=mesh, placements=_flat_placements()) fully_shard(model, mesh=mesh, placements=_flat_placements()) - inner_names = [ - name for group in model.inner.parameter_groups() for name in group.sharded_parameters - ] - outer_names = [name for group in model.parameter_groups() for name in group.sharded_parameters] + inner_names = [name for group in model.inner.parameter_groups() for name in group.parameter_names] + outer_names = [name for group in model.parameter_groups() for name in group.parameter_names] assert inner_names == ["weight"] assert outer_names == ["bias"] @@ -230,73 +217,34 @@ def test_frozen_parameter_group_does_not_allocate_main_grad(setup: DistributedSe @pytest.mark.distributed -def test_default_main_buffer_dtypes_follow_policy_contract(setup: DistributedSetup): - """Default main weights are FP32 while default main gradients match parameter dtype.""" +def test_bfloat16_main_grads_follow_grad_dtype(setup: DistributedSetup): + """BF16 full grads should reduce into main_grad without casting before reduce-scatter.""" if setup.world_size < 2: pytest.skip("This test requires at least 2 ranks.") mesh = init_device_mesh(setup.device.type, (setup.world_size,)) - model = nn.Linear(4, 4, bias=False, dtype=torch.float64).to(setup.device) - + model = nn.Linear(4, 4, bias=False, dtype=torch.bfloat16).to(setup.device) fully_shard(model, mesh=mesh, placements=_flat_placements()) (group,) = model.parameter_groups() - assert group.dtype is torch.float64 assert group.main_weight.local_buffer.dtype is torch.float32 assert group.main_grad is not None - assert group.main_grad.local_buffer.dtype is torch.float64 - - -@pytest.mark.distributed -def test_full_mesh_is_used_for_dbuffers(setup: DistributedSetup): - """This version uses the full mesh for DBuffer storage.""" - if setup.world_size < 2: - pytest.skip("This test requires at least 2 ranks.") - - mesh = init_device_mesh(setup.device.type, (1, setup.world_size), mesh_dim_names=("tp", "dp")) - placements = Placements( - dp_axes=["tp", "dp"], - parameter=[Replicate(), Flat()], - gradient=[Replicate(), Flat()], - optimizer=[Replicate(), Flat()], - ) - model = nn.Linear(4, 4, bias=False).to(setup.device) - - fully_shard(model, mesh=mesh, placements=placements) - - (group,) = model.parameter_groups() - assert group.mesh is mesh - assert group.main_weight.mesh is mesh - assert group.model_weight.mesh is mesh - assert group.main_grad is not None - assert group.main_grad.mesh is mesh - assert isinstance(model.weight, DTensor) - assert model.weight.device_mesh is mesh - - -@pytest.mark.distributed -def test_sharded_parameter_contract_uses_dtensors(setup: DistributedSetup): - """Resting sharded parameters and gradients should be DTensors.""" - if setup.world_size < 2: - pytest.skip("This test requires at least 2 ranks.") + assert group.main_grad.local_buffer.dtype is torch.bfloat16 + assert group.unsharded_parameters[0].grad_dtype is torch.bfloat16 + assert group.sharded_parameters[0].dtype is group.main_weight.local_buffer.dtype + assert group.sharded_parameters[0].grad_dtype is group.main_grad.local_buffer.dtype - mesh = init_device_mesh(setup.device.type, (setup.world_size,)) - model = nn.Linear(4, 4, bias=False).to(setup.device) - original_weight = model.weight + x = torch.randn(2, 4, device=setup.device, dtype=torch.bfloat16) - fully_shard(model, mesh=mesh, placements=_flat_placements()) - - (group,) = model.parameter_groups() - assert group.unsharded_parameters["weight"] is original_weight - assert isinstance(model.weight, DTensor) - assert isinstance(model.weight.data, DTensor) - model(torch.randn(2, 4, device=setup.device)).sum().backward() + model.zero_grad(set_to_none=True) + model(x).sum().backward() assert isinstance(model.weight, DTensor) + assert model.weight.dtype is group.main_weight.local_buffer.dtype assert isinstance(model.weight.grad, DTensor) - assert original_weight.grad is None + assert model.weight.grad.dtype is group.main_grad.local_buffer.dtype -@pytest.mark.distributed +pytest.mark.distributed def test_backward_averages_across_dp_and_accumulates_across_calls(setup: DistributedSetup): """Each backward averages over DP ranks; repeated backwards accumulate by summing.""" if setup.world_size < 2: @@ -330,51 +278,17 @@ def test_meta_parameters_initialize_with_reset_parameters(setup: DistributedSetu fully_shard(model, mesh=mesh, placements=_flat_placements()) - assert isinstance(model.weight, DTensor) - assert not model.weight.to_local().is_meta (group,) = model.parameter_groups() - assert not group.model_weight.local_buffer.is_meta - assert group.model_weight.local_buffer.numel() > 0 - full_weight = _full_model_weight(group) + full_weight = group.model_weight.allgather(0).get_tensor(0) + assert not full_weight.is_meta torch.testing.assert_close(full_weight, torch.full_like(full_weight, 3.0)) -@pytest.mark.distributed -def test_model_weight_alias_is_rejected(setup: DistributedSetup): - """Model weights cannot alias the unsharded parameter buffer.""" - if setup.world_size < 2: - pytest.skip("This test requires at least 2 ranks.") - - mesh = init_device_mesh(setup.device.type, (setup.world_size,)) - placements = Placements( - dp_axes=[0], parameter=[Replicate()], gradient=[Flat()], optimizer=[Flat()] - ) - model = nn.Linear(4, 4, bias=False).to(setup.device) - - with pytest.raises(AssertionError, match="storage is resized and released"): - fully_shard(model, mesh=mesh, placements=placements) - - -@pytest.mark.distributed -def test_meta_parameters_reject_mixed_direct_parameters(setup: DistributedSetup): - """A module cannot mix meta and real direct parameters when reset is required.""" - if setup.world_size < 2: - pytest.skip("This test requires at least 2 ranks.") - - mesh = init_device_mesh(setup.device.type, (setup.world_size,)) - model = MixedMetaRealModel(setup.device) - - with pytest.raises(ValueError, match="mixes meta and non-meta direct parameters"): - fully_shard(model, mesh=mesh, placements=_flat_placements()) - - @pytest.mark.distributed def test_non_leaf_parameter_view_survives_storage_resize(setup: DistributedSetup): """A non-leaf parameter view saved for backward should survive full-storage resize.""" if setup.world_size < 2: pytest.skip("This test requires at least 2 ranks.") - if setup.device.type != "cuda": - pytest.skip("Storage resize verification requires CUDA.") mesh = init_device_mesh(setup.device.type, (setup.world_size,)) model = NonLeafViewModel().to(setup.device) @@ -395,8 +309,8 @@ def test_non_leaf_parameter_view_survives_storage_resize(setup: DistributedSetup @pytest.mark.distributed -def test_experimental_fully_shard_reduces_peak_training_memory(setup: DistributedSetup): - """Per-layer FSDP should reduce peak CUDA memory during a train step.""" +def test_fully_shard_reduces_peak_training_memory(setup: DistributedSetup): + """Per-layer FSDP should reduce peak CUDA memory during training.""" if setup.world_size < 2: pytest.skip("This test requires at least 2 ranks.") if setup.device.type != "cuda": @@ -406,35 +320,45 @@ def test_experimental_fully_shard_reduces_peak_training_memory(setup: Distribute dim = 1024 layers = 16 batch = 8 + steps = 2 + + def train_steps(model: nn.Module, optimizer: torch.optim.Optimizer, x: torch.Tensor) -> None: + for _ in range(steps): + optimizer.zero_grad(set_to_none=True) + model(x).sum().backward() + optimizer.step() torch.manual_seed(4321) - baseline = nn.Sequential(*[nn.Linear(dim, dim, bias=False) for _ in range(layers)]).to( - setup.device - ) + baseline = nn.Sequential(*[nn.Linear(dim, dim) for _ in range(layers)]).to(setup.device) + baseline_optimizer = torch.optim.AdamW(baseline.parameters(), lr=0.01) x = torch.randn(batch, dim, device=setup.device) torch.cuda.reset_peak_memory_stats(setup.device) - baseline(x).sum().backward() + train_steps(baseline, baseline_optimizer, x) torch.cuda.synchronize(setup.device) baseline_peak = torch.cuda.max_memory_allocated(setup.device) + del baseline_optimizer del baseline del x - gc.collect() torch.cuda.empty_cache() torch.manual_seed(4321) - model = nn.Sequential(*[nn.Linear(dim, dim, bias=False) for _ in range(layers)]).to( - setup.device - ) + model = nn.Sequential(*[nn.Linear(dim, dim) for _ in range(layers)]).to(setup.device) for layer in model: fully_shard(layer, mesh=mesh, placements=_flat_placements()) - gc.collect() + optimizer = torch.optim.AdamW(model.parameters(), lr=0.01) torch.cuda.empty_cache() x = torch.randn(batch, dim, device=setup.device) torch.cuda.reset_peak_memory_stats(setup.device) - model(x).sum().backward() + train_steps(model, optimizer, x) torch.cuda.synchronize(setup.device) sharded_peak = torch.cuda.max_memory_allocated(setup.device) + logger.info( + "FSDP peak memory: rank=%s, baseline=%s, sharded=%s", + setup.rank, + _mb(baseline_peak), + _mb(sharded_peak), + ) assert sharded_peak < baseline_peak From 14e807b54e292152cba367e406a41df2ad4281d8 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Mon, 25 May 2026 07:11:32 +0000 Subject: [PATCH 05/23] Require matching FSDP main grad dtype Signed-off-by: Jingyue Wu --- .../megatron_fsdp/experimental/fully_shard.py | 44 ++++++++++----- .../test_experimental_fully_shard.py | 56 ++++++++----------- 2 files changed, 51 insertions(+), 49 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py index a2ffb9f3e83..bb137c34bdf 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py @@ -88,9 +88,9 @@ def __init__( raise ValueError("ParameterGroup requires at least one parameter.") axis_indices = tuple(_axis_index(mesh, axis) for axis in placements.dp_axes) - assert axis_indices == tuple(range(mesh.ndim)), ( - "FSDP requires dp_axes to match every mesh axis in mesh order for now." - ) + assert axis_indices == tuple( + range(mesh.ndim) + ), "FSDP requires dp_axes to match every mesh axis in mesh order for now." # --------------------------------------------------------------------- # Record shared metadata @@ -122,9 +122,7 @@ def __init__( # Python dicts preserve insertion order, so parameter_names and # parameters.values() define the same stable DBuffer tensor order. self._unsharded_model_weight = DBuffer.distribute_tensors( - parameters.values(), - mesh=self.mesh, - placements=[Replicate()] * self.mesh.ndim, + parameters.values(), mesh=self.mesh, placements=[Replicate()] * self.mesh.ndim ) # Scratch initialization starts from model weights. Checkpoint loading will # eventually initialize main weights first and derive model weights from them. @@ -143,18 +141,36 @@ def __init__( placements=placements.optimizer, ) - main_grad_dtype = mixed_precision_policy.main_grads_dtype or self.dtype - self.main_grad = ( - DBuffer( + self.main_grad = None + if self.requires_grad: + main_grad_dtype = mixed_precision_policy.main_grads_dtype or self.dtype + self.main_grad = DBuffer( mesh=self.mesh, placements=placements.gradient, tensor_shapes=self.main_weight.layout.tensor_shapes, dtype=main_grad_dtype, device=self.main_weight.local_buffer.device, ) - if self.requires_grad - else None - ) + if self.main_grad.layout != self.main_weight.layout: + raise ValueError( + "FSDP temporarily requires main_grad and main_weight to have the same " + "layout until HSDP/HFSDP support is implemented." + ) + if self.main_grad.placements != self.main_weight.placements: + raise ValueError( + "FSDP temporarily requires main_grad and main_weight to have the same " + "placements until HSDP/HFSDP support is implemented. " + f"Got main_grad placements {self.main_grad.placements} and " + f"main_weight placements {self.main_weight.placements}." + ) + if self.main_grad.local_buffer.dtype != self.main_weight.local_buffer.dtype: + raise ValueError( + "FSDP temporarily requires main_grad and main_weight to have the same " + "dtype until optimizer wrapping supports optimizer-visible gradient " + "conversion. " + f"Got main_grad dtype {self.main_grad.local_buffer.dtype} and " + f"main_weight dtype {self.main_weight.local_buffer.dtype}." + ) # --------------------------------------------------------------------- # Build parameter tuples for module swapping @@ -243,9 +259,7 @@ def has_grad(parameters: Iterable[nn.Parameter]) -> bool: grads.append(parameter.grad) partial_grad = DBuffer.distribute_tensors( - grads, - mesh=self.mesh, - placements=[Partial(dist.ReduceOp.AVG)] * self.mesh.ndim, + grads, mesh=self.mesh, placements=[Partial(dist.ReduceOp.AVG)] * self.mesh.ndim ) # zero_grad(set_to_none=True) clears sharded parameter grads, so the next diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py index eaedb85a9d6..8f2167e0714 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py +++ b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py @@ -18,9 +18,9 @@ Flat, FsdpModule, Placements, - Replicate, fully_shard, ) +from megatron.core.distributed.fsdp.src.megatron_fsdp.mixed_precision import MixedPrecisionPolicy logger = logging.getLogger(__name__) @@ -192,7 +192,9 @@ def test_nested_fully_shard_excludes_child_owned_parameters(setup: DistributedSe fully_shard(model.inner, mesh=mesh, placements=_flat_placements()) fully_shard(model, mesh=mesh, placements=_flat_placements()) - inner_names = [name for group in model.inner.parameter_groups() for name in group.parameter_names] + inner_names = [ + name for group in model.inner.parameter_groups() for name in group.parameter_names + ] outer_names = [name for group in model.parameter_groups() for name in group.parameter_names] assert inner_names == ["weight"] @@ -216,35 +218,9 @@ def test_frozen_parameter_group_does_not_allocate_main_grad(setup: DistributedSe assert group.main_grad is None -@pytest.mark.distributed -def test_bfloat16_main_grads_follow_grad_dtype(setup: DistributedSetup): - """BF16 full grads should reduce into main_grad without casting before reduce-scatter.""" - if setup.world_size < 2: - pytest.skip("This test requires at least 2 ranks.") - - mesh = init_device_mesh(setup.device.type, (setup.world_size,)) - model = nn.Linear(4, 4, bias=False, dtype=torch.bfloat16).to(setup.device) - fully_shard(model, mesh=mesh, placements=_flat_placements()) - - (group,) = model.parameter_groups() - assert group.main_weight.local_buffer.dtype is torch.float32 - assert group.main_grad is not None - assert group.main_grad.local_buffer.dtype is torch.bfloat16 - assert group.unsharded_parameters[0].grad_dtype is torch.bfloat16 - assert group.sharded_parameters[0].dtype is group.main_weight.local_buffer.dtype - assert group.sharded_parameters[0].grad_dtype is group.main_grad.local_buffer.dtype - - x = torch.randn(2, 4, device=setup.device, dtype=torch.bfloat16) - - model.zero_grad(set_to_none=True) - model(x).sum().backward() - assert isinstance(model.weight, DTensor) - assert model.weight.dtype is group.main_weight.local_buffer.dtype - assert isinstance(model.weight.grad, DTensor) - assert model.weight.grad.dtype is group.main_grad.local_buffer.dtype +pytest.mark.distributed -pytest.mark.distributed def test_backward_averages_across_dp_and_accumulates_across_calls(setup: DistributedSetup): """Each backward averages over DP ranks; repeated backwards accumulate by summing.""" if setup.world_size < 2: @@ -321,6 +297,7 @@ def test_fully_shard_reduces_peak_training_memory(setup: DistributedSetup): layers = 16 batch = 8 steps = 2 + dtype = torch.bfloat16 def train_steps(model: nn.Module, optimizer: torch.optim.Optimizer, x: torch.Tensor) -> None: for _ in range(steps): @@ -329,9 +306,11 @@ def train_steps(model: nn.Module, optimizer: torch.optim.Optimizer, x: torch.Ten optimizer.step() torch.manual_seed(4321) - baseline = nn.Sequential(*[nn.Linear(dim, dim) for _ in range(layers)]).to(setup.device) + baseline = nn.Sequential(*[nn.Linear(dim, dim, dtype=dtype) for _ in range(layers)]).to( + setup.device + ) baseline_optimizer = torch.optim.AdamW(baseline.parameters(), lr=0.01) - x = torch.randn(batch, dim, device=setup.device) + x = torch.randn(batch, dim, device=setup.device, dtype=dtype) torch.cuda.reset_peak_memory_stats(setup.device) train_steps(baseline, baseline_optimizer, x) torch.cuda.synchronize(setup.device) @@ -343,13 +322,22 @@ def train_steps(model: nn.Module, optimizer: torch.optim.Optimizer, x: torch.Ten torch.cuda.empty_cache() torch.manual_seed(4321) - model = nn.Sequential(*[nn.Linear(dim, dim) for _ in range(layers)]).to(setup.device) + model = nn.Sequential(*[nn.Linear(dim, dim, dtype=dtype) for _ in range(layers)]).to( + setup.device + ) for layer in model: - fully_shard(layer, mesh=mesh, placements=_flat_placements()) + fully_shard( + layer, + mesh=mesh, + placements=_flat_placements(), + mixed_precision_policy=MixedPrecisionPolicy( + main_params_dtype=dtype, main_grads_dtype=dtype + ), + ) optimizer = torch.optim.AdamW(model.parameters(), lr=0.01) torch.cuda.empty_cache() - x = torch.randn(batch, dim, device=setup.device) + x = torch.randn(batch, dim, device=setup.device, dtype=dtype) torch.cuda.reset_peak_memory_stats(setup.device) train_steps(model, optimizer, x) torch.cuda.synchronize(setup.device) From 680f17bf028378ed94e5c2f5c5d9c52618b03cff Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Mon, 25 May 2026 21:28:11 +0000 Subject: [PATCH 06/23] Reuse FSDP model weights for matching main weights Signed-off-by: Jingyue Wu --- .../megatron_fsdp/experimental/fully_shard.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py index bb137c34bdf..4f8592578a6 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py @@ -135,20 +135,23 @@ def __init__( ) main_weight_dtype = mixed_precision_policy.main_params_dtype or torch.float32 - self.main_weight = DBuffer.distribute_tensors( - (parameter.to(dtype=main_weight_dtype) for parameter in parameters.values()), - mesh=self.mesh, - placements=placements.optimizer, - ) + if main_weight_dtype == self.dtype and placements.optimizer == placements.parameter: + self.main_weight = self.model_weight + else: + self.main_weight = DBuffer.distribute_tensors( + (parameter.to(dtype=main_weight_dtype) for parameter in parameters.values()), + mesh=self.mesh, + placements=placements.optimizer, + ) self.main_grad = None if self.requires_grad: - main_grad_dtype = mixed_precision_policy.main_grads_dtype or self.dtype + grad_dtype = mixed_precision_policy.main_grads_dtype or self.dtype self.main_grad = DBuffer( mesh=self.mesh, placements=placements.gradient, tensor_shapes=self.main_weight.layout.tensor_shapes, - dtype=main_grad_dtype, + dtype=grad_dtype, device=self.main_weight.local_buffer.device, ) if self.main_grad.layout != self.main_weight.layout: From 5f06688e18a86ad36c58a768a45317d702ff2b86 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Tue, 26 May 2026 00:27:51 +0000 Subject: [PATCH 07/23] Sync FSDP model weights before unshard Add DBuffer casting for dtype conversion before model-weight sync, refresh model weights from main weights before unshard, and cover the next-forward optimizer update path with FP32 main weights and default BF16 main grads on SGD's non-foreach path. Signed-off-by: Jingyue Wu --- .../src/megatron_fsdp/experimental/dbuffer.py | 15 ++++ .../megatron_fsdp/experimental/fully_shard.py | 74 ++++++++++--------- .../distributed/megatron_fsdp/test_dbuffer.py | 30 ++++++++ .../test_experimental_fully_shard.py | 36 +++++++++ 4 files changed, 121 insertions(+), 34 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py index ea4e5ce0bda..438fe5f5b26 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py @@ -268,6 +268,21 @@ def _create_or_validate_out( raise ValueError(f"Expected out device {self.device}, got {out.device}.") return out + def cast(self, dtype: torch.dtype) -> "DBuffer": + """Return this buffer with the same layout and placements in ``dtype``.""" + if self.dtype == dtype: + return self + + destination = DBuffer( + mesh=self.mesh, + placements=self.placements, + tensor_shapes=self.layout.tensor_shapes, + dtype=dtype, + device=self.device, + ) + destination.local_buffer.copy_(self.local_buffer) + return destination + def redistribute( self, new_placements: Iterable[Placement], *, out: "DBuffer | None" = None ) -> "DBuffer": diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py index 4f8592578a6..d9b17280ea1 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py @@ -15,7 +15,7 @@ """Minimal per-module Megatron-FSDP implementation.""" import dataclasses -from collections.abc import Callable, Iterable, Sequence +from collections.abc import Callable, Iterable import torch import torch.distributed as dist @@ -104,11 +104,13 @@ def __init__( for name, parameter in parameters.items(): if parameter.is_meta: raise ValueError( - f"Expected parameter {name!r} to be materialized before ParameterGroup construction." + f"Expected parameter {name!r} to be materialized before " + "ParameterGroup construction." ) if parameter.dtype != self.dtype: raise ValueError( - f"Expected parameter {name!r} to have dtype {self.dtype}, got {parameter.dtype}." + f"Expected parameter {name!r} to have dtype {self.dtype}, " + f"got {parameter.dtype}." ) if parameter.requires_grad != self.requires_grad: raise ValueError( @@ -117,31 +119,34 @@ def __init__( ) # --------------------------------------------------------------------- - # Initialize DBuffers from original parameter values + # Initialize persistent DBuffers # --------------------------------------------------------------------- # Python dicts preserve insertion order, so parameter_names and # parameters.values() define the same stable DBuffer tensor order. - self._unsharded_model_weight = DBuffer.distribute_tensors( - parameters.values(), mesh=self.mesh, placements=[Replicate()] * self.mesh.ndim - ) - # Scratch initialization starts from model weights. Checkpoint loading will - # eventually initialize main weights first and derive model weights from them. - self.model_weight = self._unsharded_model_weight.redistribute(placements.parameter) - model_weight_storage = self.model_weight.local_buffer.untyped_storage() - unsharded_model_weight_storage = self._unsharded_model_weight.local_buffer.untyped_storage() - assert model_weight_storage.data_ptr() != unsharded_model_weight_storage.data_ptr(), ( - "model_weight must not share storage with _unsharded_model_weight because " - "_unsharded_model_weight storage is resized and released after resharding." + tensor_shapes = tuple(parameter.shape for parameter in parameters.values()) + main_weight_dtype = mixed_precision_policy.main_params_dtype or torch.float32 + self.main_weight = DBuffer.distribute_tensors( + (parameter.to(dtype=main_weight_dtype) for parameter in parameters.values()), + mesh=self.mesh, + placements=placements.optimizer, ) - main_weight_dtype = mixed_precision_policy.main_params_dtype or torch.float32 + self._unsharded_model_weight = DBuffer( + mesh=self.mesh, + placements=[Replicate()] * self.mesh.ndim, + tensor_shapes=tensor_shapes, + dtype=self.dtype, + device=self.main_weight.local_buffer.device, + ) if main_weight_dtype == self.dtype and placements.optimizer == placements.parameter: - self.main_weight = self.model_weight + self.model_weight = self.main_weight else: - self.main_weight = DBuffer.distribute_tensors( - (parameter.to(dtype=main_weight_dtype) for parameter in parameters.values()), + self.model_weight = DBuffer( mesh=self.mesh, - placements=placements.optimizer, + placements=placements.parameter, + tensor_shapes=tensor_shapes, + dtype=self.dtype, + device=self.main_weight.local_buffer.device, ) self.main_grad = None @@ -154,11 +159,10 @@ def __init__( dtype=grad_dtype, device=self.main_weight.local_buffer.device, ) - if self.main_grad.layout != self.main_weight.layout: - raise ValueError( - "FSDP temporarily requires main_grad and main_weight to have the same " - "layout until HSDP/HFSDP support is implemented." - ) + assert self.main_grad.layout == self.main_weight.layout, ( + "main_grad is built from main_weight tensor shapes on the same mesh, " + "and DBuffer layouts are deterministic from those shapes and mesh size." + ) if self.main_grad.placements != self.main_weight.placements: raise ValueError( "FSDP temporarily requires main_grad and main_weight to have the same " @@ -166,14 +170,6 @@ def __init__( f"Got main_grad placements {self.main_grad.placements} and " f"main_weight placements {self.main_weight.placements}." ) - if self.main_grad.local_buffer.dtype != self.main_weight.local_buffer.dtype: - raise ValueError( - "FSDP temporarily requires main_grad and main_weight to have the same " - "dtype until optimizer wrapping supports optimizer-visible gradient " - "conversion. " - f"Got main_grad dtype {self.main_grad.local_buffer.dtype} and " - f"main_weight dtype {self.main_weight.local_buffer.dtype}." - ) # --------------------------------------------------------------------- # Build parameter tuples for module swapping @@ -216,8 +212,18 @@ def _switch_to_sharded_parameters(self) -> None: def _switch_to_unsharded_parameters(self) -> None: self._set_module_parameters(self.unsharded_parameters) + def sync_model_weight_from_main_weight(self) -> None: + """Refresh compute weights from optimizer weights.""" + if self.main_weight is self.model_weight: + return + + self.main_weight.cast(self.model_weight.local_buffer.dtype).redistribute( + self.model_weight.placements, out=self.model_weight + ) + def unshard_parameters(self) -> None: """Install full parameters for local compute.""" + self.sync_model_weight_from_main_weight() self._unsharded_model_weight.reallocate_storage() # This buffer backs unsharded Parameters whose views may be saved by autograd. # Materializing FSDP-managed storage should not look like a user mutation. @@ -436,7 +442,7 @@ def visit(submodule: nn.Module, submodule_fqn: str) -> None: f"Module {submodule_fqn!r} does not have " "reset_parameters or _reset_parameters." ) - # Module.to_empty doesn't necessarily reuse Parameters so collects direct parameters again. + # Module.to_empty may replace Parameters, so collect direct parameters again. direct_parameters = list(submodule.named_parameters(recurse=False)) for local_parameter_name, parameter in direct_parameters: diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_dbuffer.py b/tests/unit_tests/distributed/megatron_fsdp/test_dbuffer.py index 6bd9bb8ab82..d267298b94d 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/test_dbuffer.py +++ b/tests/unit_tests/distributed/megatron_fsdp/test_dbuffer.py @@ -155,6 +155,36 @@ def test_constructor_allocates_local_buffer(distributed_setup): assert sharded_buffer.local_buffer.device == distributed_setup.device +@pytest.mark.distributed +def test_cast_to_same_dtype_returns_self(distributed_setup): + """DBuffer.cast returns self when the dtype already matches.""" + mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) + tensors = _same_tensors_on_all_ranks(distributed_setup.device) + buffer = DBuffer.distribute_tensors(tensors, mesh, [Replicate()]) + + assert buffer.cast(torch.float32) is buffer + + +@pytest.mark.distributed +def test_cast_preserves_layout_and_casts_values(distributed_setup): + """DBuffer.cast preserves layout metadata and casts local values.""" + mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) + tensors = _same_tensors_on_all_ranks(distributed_setup.device) + buffer = DBuffer.distribute_tensors(tensors, mesh, [Replicate()]) + + cast_buffer = buffer.cast(torch.bfloat16) + + assert cast_buffer is not buffer + assert cast_buffer.mesh == buffer.mesh + assert cast_buffer.placements == buffer.placements + assert cast_buffer.layout == buffer.layout + assert cast_buffer.device == buffer.device + assert cast_buffer.dtype is torch.bfloat16 + _assert_dbuffer_local_tensors_close( + cast_buffer, [tensor.to(dtype=torch.bfloat16) for tensor in tensors] + ) + + @pytest.mark.distributed def test_release_and_reallocate_storage_preserves_buffer_views(distributed_setup): """DBuffer storage can be released and reallocated without replacing existing views.""" diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py index 8f2167e0714..e1e09d69346 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py +++ b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py @@ -243,6 +243,42 @@ def test_backward_averages_across_dp_and_accumulates_across_calls(setup: Distrib torch.testing.assert_close(local_grad, expected, rtol=0, atol=0) +@pytest.mark.distributed +def test_next_forward_uses_optimizer_updated_weights(setup: DistributedSetup): + """The next forward should observe weights updated by the previous optimizer step.""" + if setup.world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + + mesh = init_device_mesh(setup.device.type, (setup.world_size,)) + model = nn.Linear(1, setup.world_size, bias=False, dtype=torch.bfloat16).to(setup.device) + with torch.no_grad(): + model.weight.fill_(1.0) + + fully_shard( + model, + mesh=mesh, + placements=_flat_placements(), + mixed_precision_policy=MixedPrecisionPolicy(main_params_dtype=torch.float32), + ) + # SGD's foreach/fused CUDA paths require matching parameter and gradient dtypes. + # Use the scalar path to exercise FP32 main weights with default BF16 main grads. + optimizer = torch.optim.SGD(model.parameters(), lr=0.25, foreach=False) + x = torch.ones(1, 1, device=setup.device, dtype=torch.bfloat16) + + def train_iteration() -> torch.Tensor: + optimizer.zero_grad(set_to_none=True) + loss = model(x).sum() + loss.backward() + optimizer.step() + return loss.detach().float() + + first_loss = train_iteration() + second_loss = train_iteration() + + with pytest.raises(AssertionError): + torch.testing.assert_close(second_loss, first_loss) + + @pytest.mark.distributed def test_meta_parameters_initialize_with_reset_parameters(setup: DistributedSetup): """Meta parameters should be replaced by sharded DTensors and initialized in place.""" From 09909a2a457ebf269ddf24f4ec27e8ada29b6d21 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Mon, 1 Jun 2026 02:47:46 +0000 Subject: [PATCH 08/23] Preserve autograd grad dtype in minimal FSDP Signed-off-by: Jingyue Wu --- .../megatron_fsdp/experimental/fully_shard.py | 48 +++++++++++-------- 1 file changed, 27 insertions(+), 21 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py index d9b17280ea1..44bceda5496 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py @@ -16,6 +16,7 @@ import dataclasses from collections.abc import Callable, Iterable +from typing import cast import torch import torch.distributed as dist @@ -176,20 +177,18 @@ def __init__( # --------------------------------------------------------------------- sharded_parameters: list[nn.Parameter] = [] unsharded_parameters: list[nn.Parameter] = [] - grad_dtype = self.main_grad.local_buffer.dtype if self.requires_grad else None + main_grad_dtype = self.main_grad.local_buffer.dtype if self.main_grad is not None else None for index, parameter in enumerate(parameters.values()): parameter.data = self._unsharded_model_weight.get_tensor(index) parameter.grad = None - if grad_dtype: - parameter.grad_dtype = grad_dtype setattr(parameter, _CONTAINING_PARAMETER_GROUP_ATTR, self) unsharded_parameters.append(parameter) sharded_parameter = nn.Parameter( self.main_weight.get_dtensor(index), requires_grad=parameter.requires_grad ) - if grad_dtype: - sharded_parameter.grad_dtype = grad_dtype + if main_grad_dtype: + sharded_parameter.grad_dtype = main_grad_dtype setattr(sharded_parameter, _CONTAINING_PARAMETER_GROUP_ATTR, self) sharded_parameters.append(sharded_parameter) self.sharded_parameters = tuple(sharded_parameters) @@ -260,11 +259,6 @@ def has_grad(parameters: Iterable[nn.Parameter]) -> bool: for name, parameter in zip(self.parameter_names, self.unsharded_parameters, strict=True): if parameter.grad is None: raise RuntimeError(f"Missing gradient for FSDP parameter {name!r}.") - assert parameter.grad.dtype == self.main_grad.local_buffer.dtype, ( - "FSDP unsharded parameter grad dtype should be guaranteed by grad_dtype: " - f"parameter {name!r} has grad dtype {parameter.grad.dtype}, " - f"expected main_grad dtype {self.main_grad.local_buffer.dtype}." - ) grads.append(parameter.grad) partial_grad = DBuffer.distribute_tensors( @@ -274,11 +268,21 @@ def has_grad(parameters: Iterable[nn.Parameter]) -> bool: # zero_grad(set_to_none=True) clears sharded parameter grads, so the next # backward can reduce directly into main_grad. zero_grad(set_to_none=False) # leaves sharded grads installed, so this backward accumulates into main_grad. - if has_grad(self.sharded_parameters): - reduced_grad = partial_grad.redistribute(self.main_grad.placements) - self.main_grad.local_buffer.add_(reduced_grad.local_buffer) - else: + has_sharded_grads = has_grad(self.sharded_parameters) + can_reduce_into_main_grad = ( + not has_sharded_grads + and partial_grad.local_buffer.dtype == self.main_grad.local_buffer.dtype + ) + if can_reduce_into_main_grad: partial_grad.redistribute(self.main_grad.placements, out=self.main_grad) + else: + reduced_grad = partial_grad.redistribute(self.main_grad.placements) + if has_sharded_grads: + self.main_grad.local_buffer.add_(reduced_grad.local_buffer) + else: + self.main_grad.local_buffer.copy_(reduced_grad.local_buffer) + + if not has_sharded_grads: for index, parameter in enumerate(self.sharded_parameters): parameter.grad = self.main_grad.get_dtensor(index) @@ -316,9 +320,10 @@ def __init__( self._register_hooks() def _register_hooks(self) -> None: - self.register_forward_pre_hook(lambda _module, _args: self.pre_forward()) - self.register_forward_hook(lambda _module, _args, _output: self.post_forward()) - self.register_full_backward_pre_hook(lambda _module, _grad_output: self.pre_backward()) + module = cast(nn.Module, self) + module.register_forward_pre_hook(lambda _module, _args: self.pre_forward()) + module.register_forward_hook(lambda _module, _args, _output: self.post_forward()) + module.register_full_backward_pre_hook(lambda _module, _grad_output: self.pre_backward()) # Gradient reduction is parameter-completion based: once every owned # Parameter has accumulated its grad, this FSDP unit can reduce and # reshard. Module full-backward hooks can fire before that when module @@ -399,11 +404,12 @@ def fully_shard( def _axis_index(mesh: DeviceMesh, axis: MeshAxis) -> int: if isinstance(axis, int): - if axis < 0: - axis += mesh.ndim - if axis < 0 or axis >= mesh.ndim: + axis_index = axis + if axis_index < 0: + axis_index += mesh.ndim + if axis_index < 0 or axis_index >= mesh.ndim: raise ValueError(f"Mesh axis {axis} is out of bounds for mesh ndim {mesh.ndim}.") - return axis + return axis_index dim_names = mesh.mesh_dim_names if dim_names is None or axis not in dim_names: From c4d13a8fe09a376629bf45994a5666831f6d8a1b Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Mon, 1 Jun 2026 03:09:29 +0000 Subject: [PATCH 09/23] Compare minimal FSDP loss curve with baseline Signed-off-by: Jingyue Wu --- .../test_experimental_fully_shard.py | 48 ++++++++----------- 1 file changed, 21 insertions(+), 27 deletions(-) diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py index e1e09d69346..42511460773 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py +++ b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py @@ -16,7 +16,6 @@ from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental import ( Flat, - FsdpModule, Placements, fully_shard, ) @@ -135,23 +134,11 @@ def _mb(num_bytes: int) -> str: @pytest.mark.distributed -def test_fully_shard_train_step_matches_baseline(setup: DistributedSetup): - """A minimal per-module FSDP train step should match single-rank SGD.""" +def test_fully_shard_losses_match_baseline(setup: DistributedSetup): + """Minimal per-module FSDP training should match single-rank SGD.""" if setup.world_size < 2: pytest.skip("This test requires at least 2 ranks.") - def full_named_parameters(module: nn.Module) -> dict[str, torch.Tensor]: - result = {} - for module_name, child in module.named_modules(): - if not isinstance(child, FsdpModule): - continue - prefix = f"{module_name}." if module_name else "" - for group in child.parameter_groups(): - full_weight = group.main_weight.allgather(0) - for index, name in enumerate(group.parameter_names): - result[f"{prefix}{name}"] = full_weight.get_tensor(index).detach().clone() - return result - mesh = init_device_mesh(setup.device.type, (setup.world_size,)) torch.manual_seed(1234) baseline = TinyModel().to(setup.device) @@ -160,24 +147,31 @@ def full_named_parameters(module: nn.Module) -> dict[str, torch.Tensor]: fully_shard(model.fc1, mesh=mesh, placements=_flat_placements()) fully_shard(model.fc2, mesh=mesh, placements=_flat_placements()) + baseline_optimizer = torch.optim.SGD(baseline.parameters(), lr=0.05) optimizer = torch.optim.SGD(model.parameters(), lr=0.05) x = torch.randn(3, 8, device=setup.device) target = torch.randn(3, 4, device=setup.device) - baseline_loss = torch.nn.functional.mse_loss(baseline(x), target) - baseline_loss.backward() - with torch.no_grad(): - for parameter in baseline.parameters(): - parameter.add_(parameter.grad, alpha=-0.05) - - loss = torch.nn.functional.mse_loss(model(x), target) - loss.backward() - optimizer.step() + for step in range(5): + baseline_optimizer.zero_grad() + optimizer.zero_grad() + + baseline_loss = torch.nn.functional.mse_loss(baseline(x), target) + loss = torch.nn.functional.mse_loss(model(x), target) + logger.info( + "FSDP train parity: rank=%s, step=%s, baseline_loss=%s, sharded_loss=%s", + setup.rank, + step, + baseline_loss.item(), + loss.item(), + ) + torch.testing.assert_close(loss, baseline_loss, msg=f"Loss mismatch at step {step}.") - full_params = full_named_parameters(model) - for name, expected in baseline.named_parameters(): - torch.testing.assert_close(full_params[name], expected, rtol=1e-5, atol=1e-6) + baseline_loss.backward() + loss.backward() + baseline_optimizer.step() + optimizer.step() @pytest.mark.distributed From f1c2e571cc7194b048ead286c5a1992e27a83387 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Wed, 10 Jun 2026 22:05:09 +0000 Subject: [PATCH 10/23] Adapt minimal FSDP to split DBuffer API Signed-off-by: Jingyue Wu --- .../fsdp/src/megatron_fsdp/experimental/fully_shard.py | 6 ++++-- .../megatron_fsdp/test_experimental_fully_shard.py | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py index 44bceda5496..707a062c53c 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py @@ -24,9 +24,11 @@ from torch.distributed import DeviceMesh from ..mixed_precision import MixedPrecisionPolicy -from .dbuffer import DBuffer, MeshAxis, Partial, Placement, Replicate +from .dbuffer import DBuffer +from .placement import Partial, Placement, Replicate _CONTAINING_PARAMETER_GROUP_ATTR = "_mfsdp_parameter_group" +MeshAxis = int | str @dataclasses.dataclass(frozen=True) @@ -179,7 +181,7 @@ def __init__( unsharded_parameters: list[nn.Parameter] = [] main_grad_dtype = self.main_grad.local_buffer.dtype if self.main_grad is not None else None for index, parameter in enumerate(parameters.values()): - parameter.data = self._unsharded_model_weight.get_tensor(index) + parameter.data = self._unsharded_model_weight.get_local_tensor(index) parameter.grad = None setattr(parameter, _CONTAINING_PARAMETER_GROUP_ATTR, self) unsharded_parameters.append(parameter) diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py index 42511460773..d29e1ce844b 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py +++ b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py @@ -285,7 +285,7 @@ def test_meta_parameters_initialize_with_reset_parameters(setup: DistributedSetu fully_shard(model, mesh=mesh, placements=_flat_placements()) (group,) = model.parameter_groups() - full_weight = group.model_weight.allgather(0).get_tensor(0) + full_weight = group.model_weight.allgather(0).get_local_tensor(0) assert not full_weight.is_meta torch.testing.assert_close(full_weight, torch.full_like(full_weight, 3.0)) From 195b4ca47dccb34e427688a1e2c755f9686ffeaf Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Thu, 11 Jun 2026 03:39:59 +0000 Subject: [PATCH 11/23] Split minimal FSDP runtime modules Signed-off-by: Jingyue Wu --- .../megatron_fsdp/experimental/__init__.py | 5 +- .../src/megatron_fsdp/experimental/dbuffer.py | 1 + .../megatron_fsdp/experimental/fully_shard.py | 285 +----------------- .../experimental/parameter_group.py | 264 ++++++++++++++++ .../megatron_fsdp/experimental/placement.py | 24 ++ 5 files changed, 300 insertions(+), 279 deletions(-) create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py index 1cb67e172da..13f6ec25285 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py @@ -15,8 +15,9 @@ """Experimental Megatron-FSDP implementation.""" from .dbuffer import DBuffer -from .fully_shard import FsdpModule, ParameterGroup, Placements, fully_shard -from .placement import Flat, Partial, Placement, Replicate +from .fully_shard import FsdpModule, fully_shard +from .parameter_group import ParameterGroup +from .placement import Flat, Partial, Placement, Placements, Replicate __all__ = [ "DBuffer", diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py index 438fe5f5b26..5305da55994 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py @@ -42,6 +42,7 @@ def _validate_mesh_axis(mesh: DeviceMesh, axis: int) -> None: def _validate_placements(placements: Iterable[Placement]) -> None: + """Validate DBuffer placements form a supported contiguous local layout.""" seen_flat = False for placement in placements: if not isinstance(placement, (Replicate, Partial, Flat)): diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py index 707a062c53c..8fe35a1f000 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py @@ -14,282 +14,16 @@ """Minimal per-module Megatron-FSDP implementation.""" -import dataclasses -from collections.abc import Callable, Iterable +from collections.abc import Callable from typing import cast import torch -import torch.distributed as dist from torch import nn from torch.distributed import DeviceMesh from ..mixed_precision import MixedPrecisionPolicy -from .dbuffer import DBuffer -from .placement import Partial, Placement, Replicate - -_CONTAINING_PARAMETER_GROUP_ATTR = "_mfsdp_parameter_group" -MeshAxis = int | str - - -@dataclasses.dataclass(frozen=True) -class Placements: - """Per-mesh-axis placements for parameter, gradient, and optimizer buffers.""" - - dp_axes: list[MeshAxis] - parameter: list[Placement] - gradient: list[Placement] - optimizer: list[Placement] - - def __post_init__(self) -> None: - """Validate placement list lengths.""" - axis_count = len(self.dp_axes) - for name, placements in ( - ("parameter", self.parameter), - ("gradient", self.gradient), - ("optimizer", self.optimizer), - ): - if len(placements) != axis_count: - raise ValueError(f"Expected {axis_count} {name} placements, got {len(placements)}.") - - -class ParameterGroup: - """A dtype and requires-grad homogeneous group of FSDP-owned parameters.""" - - owning_module: nn.Module - parameter_names: tuple[str, ...] - sharded_parameters: tuple[nn.Parameter, ...] - unsharded_parameters: tuple[nn.Parameter, ...] - mesh: DeviceMesh - dtype: torch.dtype - requires_grad: bool - main_weight: DBuffer - model_weight: DBuffer - main_grad: DBuffer | None - _unsharded_model_weight: DBuffer - - def __init__( - self, - owning_module: nn.Module, - parameters: dict[str, nn.Parameter], - mesh: DeviceMesh, - placements: Placements, - mixed_precision_policy: MixedPrecisionPolicy, - ) -> None: - """Create persistent sharded buffers for a group of parameters. - - Args: - owning_module: Closest FSDP root module that owns this parameter group. - parameters: Root-module-relative FQNs and their parameters. - mesh: Device mesh used for all DBuffer storage in this version. - placements: Parameter, gradient, and optimizer placements. - mixed_precision_policy: Precision policy for main weights and gradients. - """ - # --------------------------------------------------------------------- - # Validate group inputs and mesh contract - # --------------------------------------------------------------------- - if not parameters: - raise ValueError("ParameterGroup requires at least one parameter.") - - axis_indices = tuple(_axis_index(mesh, axis) for axis in placements.dp_axes) - assert axis_indices == tuple( - range(mesh.ndim) - ), "FSDP requires dp_axes to match every mesh axis in mesh order for now." - - # --------------------------------------------------------------------- - # Record shared metadata - # --------------------------------------------------------------------- - self.owning_module = owning_module - self.mesh = mesh - self.parameter_names = tuple(parameters) - first_parameter = next(iter(parameters.values())) - self.dtype = first_parameter.dtype - self.requires_grad = first_parameter.requires_grad - for name, parameter in parameters.items(): - if parameter.is_meta: - raise ValueError( - f"Expected parameter {name!r} to be materialized before " - "ParameterGroup construction." - ) - if parameter.dtype != self.dtype: - raise ValueError( - f"Expected parameter {name!r} to have dtype {self.dtype}, " - f"got {parameter.dtype}." - ) - if parameter.requires_grad != self.requires_grad: - raise ValueError( - f"Expected parameter {name!r} to have requires_grad={self.requires_grad}, " - f"got {parameter.requires_grad}." - ) - - # --------------------------------------------------------------------- - # Initialize persistent DBuffers - # --------------------------------------------------------------------- - # Python dicts preserve insertion order, so parameter_names and - # parameters.values() define the same stable DBuffer tensor order. - tensor_shapes = tuple(parameter.shape for parameter in parameters.values()) - main_weight_dtype = mixed_precision_policy.main_params_dtype or torch.float32 - self.main_weight = DBuffer.distribute_tensors( - (parameter.to(dtype=main_weight_dtype) for parameter in parameters.values()), - mesh=self.mesh, - placements=placements.optimizer, - ) - - self._unsharded_model_weight = DBuffer( - mesh=self.mesh, - placements=[Replicate()] * self.mesh.ndim, - tensor_shapes=tensor_shapes, - dtype=self.dtype, - device=self.main_weight.local_buffer.device, - ) - if main_weight_dtype == self.dtype and placements.optimizer == placements.parameter: - self.model_weight = self.main_weight - else: - self.model_weight = DBuffer( - mesh=self.mesh, - placements=placements.parameter, - tensor_shapes=tensor_shapes, - dtype=self.dtype, - device=self.main_weight.local_buffer.device, - ) - - self.main_grad = None - if self.requires_grad: - grad_dtype = mixed_precision_policy.main_grads_dtype or self.dtype - self.main_grad = DBuffer( - mesh=self.mesh, - placements=placements.gradient, - tensor_shapes=self.main_weight.layout.tensor_shapes, - dtype=grad_dtype, - device=self.main_weight.local_buffer.device, - ) - assert self.main_grad.layout == self.main_weight.layout, ( - "main_grad is built from main_weight tensor shapes on the same mesh, " - "and DBuffer layouts are deterministic from those shapes and mesh size." - ) - if self.main_grad.placements != self.main_weight.placements: - raise ValueError( - "FSDP temporarily requires main_grad and main_weight to have the same " - "placements until HSDP/HFSDP support is implemented. " - f"Got main_grad placements {self.main_grad.placements} and " - f"main_weight placements {self.main_weight.placements}." - ) - - # --------------------------------------------------------------------- - # Build parameter tuples for module swapping - # --------------------------------------------------------------------- - sharded_parameters: list[nn.Parameter] = [] - unsharded_parameters: list[nn.Parameter] = [] - main_grad_dtype = self.main_grad.local_buffer.dtype if self.main_grad is not None else None - for index, parameter in enumerate(parameters.values()): - parameter.data = self._unsharded_model_weight.get_local_tensor(index) - parameter.grad = None - setattr(parameter, _CONTAINING_PARAMETER_GROUP_ATTR, self) - unsharded_parameters.append(parameter) - - sharded_parameter = nn.Parameter( - self.main_weight.get_dtensor(index), requires_grad=parameter.requires_grad - ) - if main_grad_dtype: - sharded_parameter.grad_dtype = main_grad_dtype - setattr(sharded_parameter, _CONTAINING_PARAMETER_GROUP_ATTR, self) - sharded_parameters.append(sharded_parameter) - self.sharded_parameters = tuple(sharded_parameters) - self.unsharded_parameters = tuple(unsharded_parameters) - - # --------------------------------------------------------------------- - # Install resting sharded parameters - # --------------------------------------------------------------------- - self._switch_to_sharded_parameters() - self._unsharded_model_weight.release_storage() - - def _set_module_parameters(self, parameters: tuple[nn.Parameter, ...]) -> None: - for name, parameter in zip(self.parameter_names, parameters, strict=True): - module, parameter_name = _get_parameter_owner(self.owning_module, name) - module._parameters[parameter_name] = parameter - - def _switch_to_sharded_parameters(self) -> None: - self._set_module_parameters(self.sharded_parameters) - - def _switch_to_unsharded_parameters(self) -> None: - self._set_module_parameters(self.unsharded_parameters) - - def sync_model_weight_from_main_weight(self) -> None: - """Refresh compute weights from optimizer weights.""" - if self.main_weight is self.model_weight: - return - - self.main_weight.cast(self.model_weight.local_buffer.dtype).redistribute( - self.model_weight.placements, out=self.model_weight - ) - - def unshard_parameters(self) -> None: - """Install full parameters for local compute.""" - self.sync_model_weight_from_main_weight() - self._unsharded_model_weight.reallocate_storage() - # This buffer backs unsharded Parameters whose views may be saved by autograd. - # Materializing FSDP-managed storage should not look like a user mutation. - with torch.autograd._unsafe_preserve_version_counter( - self._unsharded_model_weight.local_buffer - ): - self.model_weight.redistribute( - self._unsharded_model_weight.placements, out=self._unsharded_model_weight - ) - self._switch_to_unsharded_parameters() - - def reshard_parameters(self) -> None: - """Install sharded DTensor parameters on the owning modules.""" - self._switch_to_sharded_parameters() - self._unsharded_model_weight.release_storage() - - def reduce_gradients(self) -> None: - """Reduce full local gradients into sharded parameter gradients.""" - assert self.main_grad is not None - - def has_grad(parameters: Iterable[nn.Parameter]) -> bool: - has_any_grad = False - has_any_missing_grad = False - for parameter in parameters: - if parameter.grad is None: - has_any_missing_grad = True - else: - has_any_grad = True - if has_any_grad and has_any_missing_grad: - raise RuntimeError("FSDP sharded gradients must be either all set or all None.") - return has_any_grad - - grads: list[torch.Tensor] = [] - for name, parameter in zip(self.parameter_names, self.unsharded_parameters, strict=True): - if parameter.grad is None: - raise RuntimeError(f"Missing gradient for FSDP parameter {name!r}.") - grads.append(parameter.grad) - - partial_grad = DBuffer.distribute_tensors( - grads, mesh=self.mesh, placements=[Partial(dist.ReduceOp.AVG)] * self.mesh.ndim - ) - - # zero_grad(set_to_none=True) clears sharded parameter grads, so the next - # backward can reduce directly into main_grad. zero_grad(set_to_none=False) - # leaves sharded grads installed, so this backward accumulates into main_grad. - has_sharded_grads = has_grad(self.sharded_parameters) - can_reduce_into_main_grad = ( - not has_sharded_grads - and partial_grad.local_buffer.dtype == self.main_grad.local_buffer.dtype - ) - if can_reduce_into_main_grad: - partial_grad.redistribute(self.main_grad.placements, out=self.main_grad) - else: - reduced_grad = partial_grad.redistribute(self.main_grad.placements) - if has_sharded_grads: - self.main_grad.local_buffer.add_(reduced_grad.local_buffer) - else: - self.main_grad.local_buffer.copy_(reduced_grad.local_buffer) - - if not has_sharded_grads: - for index, parameter in enumerate(self.sharded_parameters): - parameter.grad = self.main_grad.get_dtensor(index) - - for parameter in self.unsharded_parameters: - parameter.grad = None +from .parameter_group import ParameterGroup, contained_in_parameter_group +from .placement import MeshAxis, Placements class FsdpModule: @@ -304,6 +38,10 @@ def __init__( ) -> None: """Initialize FSDP runtime state on an already-constructed module.""" owned_parameters = _materialize_and_collect_owned_parameters(self, _mesh_device(mesh)) + axis_indices = tuple(_axis_index(mesh, axis) for axis in placements.dp_axes) + assert axis_indices == tuple( + range(mesh.ndim) + ), "FSDP requires dp_axes to match every mesh axis in mesh order for now." parameter_groups = [ ParameterGroup( owning_module=self, @@ -457,7 +195,7 @@ def visit(submodule: nn.Module, submodule_fqn: str) -> None: parameter_fqn = ( f"{submodule_fqn}.{local_parameter_name}" if submodule_fqn else local_parameter_name ) - if hasattr(parameter, _CONTAINING_PARAMETER_GROUP_ATTR): + if contained_in_parameter_group(parameter): raise ValueError(f"Parameter {parameter_fqn!r} is already owned by an FSDP unit.") parameters[parameter_fqn] = parameter @@ -481,13 +219,6 @@ def _group_parameters(parameters: dict[str, nn.Parameter]) -> list[dict[str, nn. return [grouped[key] for key in grouped] -def _get_parameter_owner(module: nn.Module, name: str) -> tuple[nn.Module, str]: - """Resolve a root-module-relative parameter FQN to its direct owner.""" - module_name, separator, parameter_name = name.rpartition(".") - owner = module.get_submodule(module_name) if separator else module - return owner, parameter_name - - def _attach_mixin(module: nn.Module) -> None: if isinstance(module, FsdpModule): return diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py new file mode 100644 index 00000000000..2dad4ca7732 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py @@ -0,0 +1,264 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Parameter-group runtime state for the minimal Megatron-FSDP path.""" + +from collections.abc import Iterable + +import torch +import torch.distributed as dist +from torch import nn +from torch.distributed import DeviceMesh + +from ..mixed_precision import MixedPrecisionPolicy +from .dbuffer import DBuffer +from .placement import Partial, Placements, Replicate + +_CONTAINING_PARAMETER_GROUP_ATTR = "_mfsdp_parameter_group" + + +def contained_in_parameter_group(parameter: nn.Parameter) -> bool: + """Return whether a parameter is already owned by a ParameterGroup.""" + return hasattr(parameter, _CONTAINING_PARAMETER_GROUP_ATTR) + + +class ParameterGroup: + """A dtype and requires-grad homogeneous group of FSDP-owned parameters.""" + + owning_module: nn.Module + parameter_names: tuple[str, ...] + sharded_parameters: tuple[nn.Parameter, ...] + unsharded_parameters: tuple[nn.Parameter, ...] + mesh: DeviceMesh + dtype: torch.dtype + requires_grad: bool + main_weight: DBuffer + model_weight: DBuffer + main_grad: DBuffer | None + _unsharded_model_weight: DBuffer + + def __init__( + self, + owning_module: nn.Module, + parameters: dict[str, nn.Parameter], + mesh: DeviceMesh, + placements: Placements, + mixed_precision_policy: MixedPrecisionPolicy, + ) -> None: + """Create persistent sharded buffers for a group of parameters. + + Args: + owning_module: Closest FSDP root module that owns this parameter group. + parameters: Root-module-relative FQNs and their parameters. + mesh: Device mesh used for all DBuffer storage in this version. + placements: Parameter, gradient, and optimizer placements. + mixed_precision_policy: Precision policy for main weights and gradients. + """ + if not parameters: + raise ValueError("ParameterGroup requires at least one parameter.") + + model_weight_placements = tuple(placements.parameter) + main_grad_placements = tuple(placements.gradient) + main_weight_placements = tuple(placements.optimizer) + + # Python dicts preserve insertion order, so parameter_names and + # parameters.values() define the same stable DBuffer tensor order. + self.owning_module = owning_module + self.mesh = mesh + self.parameter_names = tuple(parameters) + first_parameter = next(iter(parameters.values())) + self.dtype = first_parameter.dtype + self.requires_grad = first_parameter.requires_grad + for name, parameter in parameters.items(): + if parameter.is_meta: + raise ValueError( + f"Expected parameter {name!r} to be materialized before " + "ParameterGroup construction." + ) + if parameter.dtype != self.dtype: + raise ValueError( + f"Expected parameter {name!r} to have dtype {self.dtype}, " + f"got {parameter.dtype}." + ) + if parameter.requires_grad != self.requires_grad: + raise ValueError( + f"Expected parameter {name!r} to have requires_grad={self.requires_grad}, " + f"got {parameter.requires_grad}." + ) + + tensor_shapes = tuple(parameter.shape for parameter in parameters.values()) + main_weight_dtype = mixed_precision_policy.main_params_dtype or torch.float32 + self.main_weight = DBuffer.distribute_tensors( + (parameter.to(dtype=main_weight_dtype) for parameter in parameters.values()), + mesh=self.mesh, + placements=main_weight_placements, + ) + + self._unsharded_model_weight = DBuffer( + mesh=self.mesh, + placements=[Replicate()] * self.mesh.ndim, + tensor_shapes=tensor_shapes, + dtype=self.dtype, + device=self.main_weight.local_buffer.device, + ) + if main_weight_dtype == self.dtype and main_weight_placements == model_weight_placements: + self.model_weight = self.main_weight + else: + self.model_weight = DBuffer( + mesh=self.mesh, + placements=model_weight_placements, + tensor_shapes=tensor_shapes, + dtype=self.dtype, + device=self.main_weight.local_buffer.device, + ) + + self.main_grad = None + if self.requires_grad: + grad_dtype = mixed_precision_policy.main_grads_dtype or self.dtype + self.main_grad = DBuffer( + mesh=self.mesh, + placements=main_grad_placements, + tensor_shapes=self.main_weight.layout.tensor_shapes, + dtype=grad_dtype, + device=self.main_weight.local_buffer.device, + ) + assert self.main_grad.layout == self.main_weight.layout, ( + "main_grad is built from main_weight tensor shapes on the same mesh, " + "and DBuffer layouts are deterministic from those shapes and mesh size." + ) + if self.main_grad.placements != self.main_weight.placements: + raise ValueError( + "FSDP temporarily requires main_grad and main_weight to have the same " + "placements until HSDP/HFSDP support is implemented. " + f"Got main_grad placements {self.main_grad.placements} and " + f"main_weight placements {self.main_weight.placements}." + ) + + sharded_parameters: list[nn.Parameter] = [] + unsharded_parameters: list[nn.Parameter] = [] + main_grad_dtype = self.main_grad.local_buffer.dtype if self.main_grad is not None else None + for index, parameter in enumerate(parameters.values()): + parameter.data = self._unsharded_model_weight.get_local_tensor(index) + parameter.grad = None + setattr(parameter, _CONTAINING_PARAMETER_GROUP_ATTR, self) + unsharded_parameters.append(parameter) + + sharded_parameter = nn.Parameter( + self.main_weight.get_dtensor(index), requires_grad=parameter.requires_grad + ) + if main_grad_dtype: + sharded_parameter.grad_dtype = main_grad_dtype + setattr(sharded_parameter, _CONTAINING_PARAMETER_GROUP_ATTR, self) + sharded_parameters.append(sharded_parameter) + self.sharded_parameters = tuple(sharded_parameters) + self.unsharded_parameters = tuple(unsharded_parameters) + + self._switch_to_sharded_parameters() + self._unsharded_model_weight.release_storage() + + def _set_module_parameters(self, parameters: tuple[nn.Parameter, ...]) -> None: + for name, parameter in zip(self.parameter_names, parameters, strict=True): + module, parameter_name = _get_parameter_owner(self.owning_module, name) + module._parameters[parameter_name] = parameter + + def _switch_to_sharded_parameters(self) -> None: + self._set_module_parameters(self.sharded_parameters) + + def _switch_to_unsharded_parameters(self) -> None: + self._set_module_parameters(self.unsharded_parameters) + + def sync_model_weight_from_main_weight(self) -> None: + """Refresh compute weights from optimizer weights.""" + if self.main_weight is self.model_weight: + return + + self.main_weight.cast(self.model_weight.local_buffer.dtype).redistribute( + self.model_weight.placements, out=self.model_weight + ) + + def unshard_parameters(self) -> None: + """Install full parameters for local compute.""" + self.sync_model_weight_from_main_weight() + self._unsharded_model_weight.reallocate_storage() + # This buffer backs unsharded Parameters whose views may be saved by autograd. + # Materializing FSDP-managed storage should not look like a user mutation. + with torch.autograd._unsafe_preserve_version_counter( + self._unsharded_model_weight.local_buffer + ): + self.model_weight.redistribute( + self._unsharded_model_weight.placements, out=self._unsharded_model_weight + ) + self._switch_to_unsharded_parameters() + + def reshard_parameters(self) -> None: + """Install sharded DTensor parameters on the owning modules.""" + self._switch_to_sharded_parameters() + self._unsharded_model_weight.release_storage() + + def reduce_gradients(self) -> None: + """Reduce full local gradients into sharded parameter gradients.""" + assert self.main_grad is not None + + def has_grad(parameters: Iterable[nn.Parameter]) -> bool: + has_any_grad = False + has_any_missing_grad = False + for parameter in parameters: + if parameter.grad is None: + has_any_missing_grad = True + else: + has_any_grad = True + if has_any_grad and has_any_missing_grad: + raise RuntimeError("FSDP sharded gradients must be either all set or all None.") + return has_any_grad + + grads: list[torch.Tensor] = [] + for name, parameter in zip(self.parameter_names, self.unsharded_parameters, strict=True): + if parameter.grad is None: + raise RuntimeError(f"Missing gradient for FSDP parameter {name!r}.") + grads.append(parameter.grad) + + partial_grad = DBuffer.distribute_tensors( + grads, mesh=self.mesh, placements=[Partial(dist.ReduceOp.AVG)] * self.mesh.ndim + ) + + # zero_grad(set_to_none=True) clears sharded parameter grads, so the next + # backward can reduce directly into main_grad. zero_grad(set_to_none=False) + # leaves sharded grads installed, so this backward accumulates into main_grad. + has_sharded_grads = has_grad(self.sharded_parameters) + can_reduce_into_main_grad = ( + not has_sharded_grads + and partial_grad.local_buffer.dtype == self.main_grad.local_buffer.dtype + ) + if can_reduce_into_main_grad: + partial_grad.redistribute(self.main_grad.placements, out=self.main_grad) + else: + reduced_grad = partial_grad.redistribute(self.main_grad.placements) + if has_sharded_grads: + self.main_grad.local_buffer.add_(reduced_grad.local_buffer) + else: + self.main_grad.local_buffer.copy_(reduced_grad.local_buffer) + + if not has_sharded_grads: + for index, parameter in enumerate(self.sharded_parameters): + parameter.grad = self.main_grad.get_dtensor(index) + + for parameter in self.unsharded_parameters: + parameter.grad = None + + +def _get_parameter_owner(module: nn.Module, name: str) -> tuple[nn.Module, str]: + """Resolve a root-module-relative parameter FQN to its direct owner.""" + module_name, separator, parameter_name = name.rpartition(".") + owner = module.get_submodule(module_name) if separator else module + return owner, parameter_name diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/placement.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/placement.py index 1b561c9634d..5e4dc6b985e 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/placement.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/placement.py @@ -38,6 +38,9 @@ class Placement: """Base class for DBuffer placements.""" +MeshAxis = int | str + + @dataclasses.dataclass(frozen=True) class Replicate(Placement): """Replicated local buffer placement.""" @@ -53,3 +56,24 @@ class Partial(Placement): @dataclasses.dataclass(frozen=True) class Flat(Placement): """Flat per-unit dim-0 sharded local buffer placement.""" + + +@dataclasses.dataclass(frozen=True) +class Placements: + """Per-mesh-axis placements for parameter, gradient, and optimizer buffers.""" + + dp_axes: list[MeshAxis] + parameter: list[Placement] + gradient: list[Placement] + optimizer: list[Placement] + + def __post_init__(self) -> None: + """Validate placement list lengths.""" + axis_count = len(self.dp_axes) + for name, placements in ( + ("parameter", self.parameter), + ("gradient", self.gradient), + ("optimizer", self.optimizer), + ): + if len(placements) != axis_count: + raise ValueError(f"Expected {axis_count} {name} placements, got {len(placements)}.") From 13055d830b2a9067abf4eed96064914f2ca2d545 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Thu, 11 Jun 2026 03:47:13 +0000 Subject: [PATCH 12/23] Split minimal FSDP module mixin Signed-off-by: Jingyue Wu --- .../megatron_fsdp/experimental/__init__.py | 3 +- .../megatron_fsdp/experimental/fsdp_module.py | 188 ++++++++++++++++++ .../megatron_fsdp/experimental/fully_shard.py | 172 +--------------- 3 files changed, 193 insertions(+), 170 deletions(-) create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fsdp_module.py diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py index 13f6ec25285..47db6401d8e 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py @@ -15,7 +15,8 @@ """Experimental Megatron-FSDP implementation.""" from .dbuffer import DBuffer -from .fully_shard import FsdpModule, fully_shard +from .fsdp_module import FsdpModule +from .fully_shard import fully_shard from .parameter_group import ParameterGroup from .placement import Flat, Partial, Placement, Placements, Replicate diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fsdp_module.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fsdp_module.py new file mode 100644 index 00000000000..e6ce5b0037f --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fsdp_module.py @@ -0,0 +1,188 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Module mixin for the minimal Megatron-FSDP path.""" + +from collections.abc import Callable +from typing import cast + +import torch +from torch import nn +from torch.distributed import DeviceMesh + +from ..mixed_precision import MixedPrecisionPolicy +from .parameter_group import ParameterGroup, contained_in_parameter_group +from .placement import MeshAxis, Placements + + +class FsdpModule: + """Mixin attached to modules managed by the minimal FSDP path.""" + + _parameter_groups: tuple[ParameterGroup, ...] + _ready_grad_parameters: set[nn.Parameter] + num_training_parameters: int + + def __init__( + self, mesh: DeviceMesh, placements: Placements, mixed_precision_policy: MixedPrecisionPolicy + ) -> None: + """Initialize FSDP runtime state on an already-constructed module.""" + owned_parameters = _materialize_and_collect_owned_parameters(self, _mesh_device(mesh)) + axis_indices = tuple(_axis_index(mesh, axis) for axis in placements.dp_axes) + assert axis_indices == tuple( + range(mesh.ndim) + ), "FSDP requires dp_axes to match every mesh axis in mesh order for now." + parameter_groups = [ + ParameterGroup( + owning_module=self, + parameters=group_parameters, + mesh=mesh, + placements=placements, + mixed_precision_policy=mixed_precision_policy, + ) + for group_parameters in _group_parameters(owned_parameters) + ] + self._parameter_groups = tuple(parameter_groups) + self._ready_grad_parameters = set() + self.num_training_parameters = sum( + len(group.sharded_parameters) for group in self._parameter_groups if group.requires_grad + ) + self._register_hooks() + + def _register_hooks(self) -> None: + module = cast(nn.Module, self) + module.register_forward_pre_hook(lambda _module, _args: self.pre_forward()) + module.register_forward_hook(lambda _module, _args, _output: self.post_forward()) + module.register_full_backward_pre_hook(lambda _module, _grad_output: self.pre_backward()) + # Gradient reduction is parameter-completion based: once every owned + # Parameter has accumulated its grad, this FSDP unit can reduce and + # reshard. Module full-backward hooks can fire before that when module + # inputs do not require grad. + for group in self._parameter_groups: + if not group.requires_grad: + continue + for parameter in group.unsharded_parameters: + parameter.register_post_accumulate_grad_hook(self._make_grad_hook(parameter)) + + def _make_grad_hook(self, parameter: nn.Parameter) -> Callable[[nn.Parameter], None]: + def grad_hook(_parameter: nn.Parameter) -> None: + self._ready_grad_parameters.add(parameter) + if len(self._ready_grad_parameters) == self.num_training_parameters: + self.post_backward() + + return grad_hook + + def pre_forward(self) -> None: + """Prepare full parameters for forward compute.""" + self._ready_grad_parameters.clear() + for group in self._parameter_groups: + group.unshard_parameters() + + def post_forward(self) -> None: + """Return parameters to their sharded resting state after forward compute.""" + for group in self._parameter_groups: + group.reshard_parameters() + + def pre_backward(self) -> None: + """Prepare full parameters for backward compute.""" + for group in self._parameter_groups: + group.unshard_parameters() + + def post_backward(self) -> None: + """Reduce gradients and return parameters to their sharded resting state.""" + for group in self._parameter_groups: + if group.requires_grad: + group.reduce_gradients() + group.reshard_parameters() + self._ready_grad_parameters.clear() + + def parameter_groups(self) -> tuple[ParameterGroup, ...]: + """Return parameter groups owned by this FSDP unit.""" + return self._parameter_groups + + +def _axis_index(mesh: DeviceMesh, axis: MeshAxis) -> int: + if isinstance(axis, int): + axis_index = axis + if axis_index < 0: + axis_index += mesh.ndim + if axis_index < 0 or axis_index >= mesh.ndim: + raise ValueError(f"Mesh axis {axis} is out of bounds for mesh ndim {mesh.ndim}.") + return axis_index + + dim_names = mesh.mesh_dim_names + if dim_names is None or axis not in dim_names: + raise ValueError(f"Mesh axis {axis!r} is not present in mesh dim names {dim_names}.") + return dim_names.index(axis) + + +def _mesh_device(mesh: DeviceMesh) -> torch.device: + if mesh.device_type == "cuda": + return torch.device("cuda", torch.cuda.current_device()) + return torch.device(mesh.device_type) + + +def _materialize_and_collect_owned_parameters( + root_module: nn.Module, device: torch.device +) -> dict[str, nn.Parameter]: + parameters: dict[str, nn.Parameter] = {} + + def visit(submodule: nn.Module, submodule_fqn: str) -> None: + direct_parameters = list(submodule.named_parameters(recurse=False)) + + if any(parameter.is_meta for _, parameter in direct_parameters): + if any(not parameter.is_meta for _, parameter in direct_parameters): + raise ValueError( + f"Module {submodule_fqn!r} mixes meta and non-meta direct parameters. " + "Initialize all direct parameters on meta or none of them." + ) + submodule.to_empty(device=device, recurse=False) + with torch.no_grad(): + if hasattr(submodule, "reset_parameters"): + submodule.reset_parameters() + elif hasattr(submodule, "_reset_parameters"): + submodule._reset_parameters() + else: + raise ValueError( + f"Module {submodule_fqn!r} does not have " + "reset_parameters or _reset_parameters." + ) + # Module.to_empty may replace Parameters, so collect direct parameters again. + direct_parameters = list(submodule.named_parameters(recurse=False)) + + for local_parameter_name, parameter in direct_parameters: + parameter_fqn = ( + f"{submodule_fqn}.{local_parameter_name}" if submodule_fqn else local_parameter_name + ) + if contained_in_parameter_group(parameter): + raise ValueError(f"Parameter {parameter_fqn!r} is already owned by an FSDP unit.") + parameters[parameter_fqn] = parameter + + for child_name, child_module in submodule.named_children(): + if isinstance(child_module, FsdpModule): + continue + child_fqn = f"{submodule_fqn}.{child_name}" if submodule_fqn else child_name + visit(child_module, child_fqn) + + visit(root_module, "") + if not parameters: + raise ValueError("fully_shard requires at least one unowned parameter.") + return parameters + + +def _group_parameters(parameters: dict[str, nn.Parameter]) -> list[dict[str, nn.Parameter]]: + grouped: dict[tuple[torch.dtype, bool], dict[str, nn.Parameter]] = {} + for name, parameter in parameters.items(): + key = (parameter.dtype, parameter.requires_grad) + grouped.setdefault(key, {})[name] = parameter + return [grouped[key] for key in grouped] diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py index 8fe35a1f000..3ab3c0e6d71 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py @@ -12,103 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Minimal per-module Megatron-FSDP implementation.""" +"""Minimal Megatron-FSDP fully_shard entrypoint.""" -from collections.abc import Callable -from typing import cast - -import torch from torch import nn from torch.distributed import DeviceMesh from ..mixed_precision import MixedPrecisionPolicy -from .parameter_group import ParameterGroup, contained_in_parameter_group -from .placement import MeshAxis, Placements - - -class FsdpModule: - """Mixin attached to modules managed by the minimal FSDP path.""" - - _parameter_groups: tuple[ParameterGroup, ...] - _ready_grad_parameters: set[nn.Parameter] - num_training_parameters: int - - def __init__( - self, mesh: DeviceMesh, placements: Placements, mixed_precision_policy: MixedPrecisionPolicy - ) -> None: - """Initialize FSDP runtime state on an already-constructed module.""" - owned_parameters = _materialize_and_collect_owned_parameters(self, _mesh_device(mesh)) - axis_indices = tuple(_axis_index(mesh, axis) for axis in placements.dp_axes) - assert axis_indices == tuple( - range(mesh.ndim) - ), "FSDP requires dp_axes to match every mesh axis in mesh order for now." - parameter_groups = [ - ParameterGroup( - owning_module=self, - parameters=group_parameters, - mesh=mesh, - placements=placements, - mixed_precision_policy=mixed_precision_policy, - ) - for group_parameters in _group_parameters(owned_parameters) - ] - self._parameter_groups = tuple(parameter_groups) - self._ready_grad_parameters = set() - self.num_training_parameters = sum( - len(group.sharded_parameters) for group in self._parameter_groups if group.requires_grad - ) - self._register_hooks() - - def _register_hooks(self) -> None: - module = cast(nn.Module, self) - module.register_forward_pre_hook(lambda _module, _args: self.pre_forward()) - module.register_forward_hook(lambda _module, _args, _output: self.post_forward()) - module.register_full_backward_pre_hook(lambda _module, _grad_output: self.pre_backward()) - # Gradient reduction is parameter-completion based: once every owned - # Parameter has accumulated its grad, this FSDP unit can reduce and - # reshard. Module full-backward hooks can fire before that when module - # inputs do not require grad. - for group in self._parameter_groups: - if not group.requires_grad: - continue - for parameter in group.unsharded_parameters: - parameter.register_post_accumulate_grad_hook(self._make_grad_hook(parameter)) - - def _make_grad_hook(self, parameter: nn.Parameter) -> Callable[[nn.Parameter], None]: - def grad_hook(_parameter: nn.Parameter) -> None: - self._ready_grad_parameters.add(parameter) - if len(self._ready_grad_parameters) == self.num_training_parameters: - self.post_backward() - - return grad_hook - - def pre_forward(self) -> None: - """Prepare full parameters for forward compute.""" - self._ready_grad_parameters.clear() - for group in self._parameter_groups: - group.unshard_parameters() - - def post_forward(self) -> None: - """Return parameters to their sharded resting state after forward compute.""" - for group in self._parameter_groups: - group.reshard_parameters() - - def pre_backward(self) -> None: - """Prepare full parameters for backward compute.""" - for group in self._parameter_groups: - group.unshard_parameters() - - def post_backward(self) -> None: - """Reduce gradients and return parameters to their sharded resting state.""" - for group in self._parameter_groups: - if group.requires_grad: - group.reduce_gradients() - group.reshard_parameters() - self._ready_grad_parameters.clear() - - def parameter_groups(self) -> tuple[ParameterGroup, ...]: - """Return parameter groups owned by this FSDP unit.""" - return self._parameter_groups +from .fsdp_module import FsdpModule +from .placement import Placements def fully_shard( @@ -142,83 +53,6 @@ def fully_shard( raise -def _axis_index(mesh: DeviceMesh, axis: MeshAxis) -> int: - if isinstance(axis, int): - axis_index = axis - if axis_index < 0: - axis_index += mesh.ndim - if axis_index < 0 or axis_index >= mesh.ndim: - raise ValueError(f"Mesh axis {axis} is out of bounds for mesh ndim {mesh.ndim}.") - return axis_index - - dim_names = mesh.mesh_dim_names - if dim_names is None or axis not in dim_names: - raise ValueError(f"Mesh axis {axis!r} is not present in mesh dim names {dim_names}.") - return dim_names.index(axis) - - -def _mesh_device(mesh: DeviceMesh) -> torch.device: - if mesh.device_type == "cuda": - return torch.device("cuda", torch.cuda.current_device()) - return torch.device(mesh.device_type) - - -def _materialize_and_collect_owned_parameters( - root_module: nn.Module, device: torch.device -) -> dict[str, nn.Parameter]: - parameters: dict[str, nn.Parameter] = {} - - def visit(submodule: nn.Module, submodule_fqn: str) -> None: - direct_parameters = list(submodule.named_parameters(recurse=False)) - - if any(parameter.is_meta for _, parameter in direct_parameters): - if any(not parameter.is_meta for _, parameter in direct_parameters): - raise ValueError( - f"Module {submodule_fqn!r} mixes meta and non-meta direct parameters. " - "Initialize all direct parameters on meta or none of them." - ) - submodule.to_empty(device=device, recurse=False) - with torch.no_grad(): - if hasattr(submodule, "reset_parameters"): - submodule.reset_parameters() - elif hasattr(submodule, "_reset_parameters"): - submodule._reset_parameters() - else: - raise ValueError( - f"Module {submodule_fqn!r} does not have " - "reset_parameters or _reset_parameters." - ) - # Module.to_empty may replace Parameters, so collect direct parameters again. - direct_parameters = list(submodule.named_parameters(recurse=False)) - - for local_parameter_name, parameter in direct_parameters: - parameter_fqn = ( - f"{submodule_fqn}.{local_parameter_name}" if submodule_fqn else local_parameter_name - ) - if contained_in_parameter_group(parameter): - raise ValueError(f"Parameter {parameter_fqn!r} is already owned by an FSDP unit.") - parameters[parameter_fqn] = parameter - - for child_name, child_module in submodule.named_children(): - if isinstance(child_module, FsdpModule): - continue - child_fqn = f"{submodule_fqn}.{child_name}" if submodule_fqn else child_name - visit(child_module, child_fqn) - - visit(root_module, "") - if not parameters: - raise ValueError("fully_shard requires at least one unowned parameter.") - return parameters - - -def _group_parameters(parameters: dict[str, nn.Parameter]) -> list[dict[str, nn.Parameter]]: - grouped: dict[tuple[torch.dtype, bool], dict[str, nn.Parameter]] = {} - for name, parameter in parameters.items(): - key = (parameter.dtype, parameter.requires_grad) - grouped.setdefault(key, {})[name] = parameter - return [grouped[key] for key in grouped] - - def _attach_mixin(module: nn.Module) -> None: if isinstance(module, FsdpModule): return From 875975084d61543ef89606c4b9dddae337427d11 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Fri, 12 Jun 2026 19:05:36 +0000 Subject: [PATCH 13/23] Remove experimental FSDP meta parameter support Signed-off-by: Jingyue Wu --- .../megatron_fsdp/experimental/fsdp_module.py | 39 +--- .../experimental/parameter_group.py | 18 +- .../test_experimental_fully_shard.py | 177 ++++++++---------- 3 files changed, 88 insertions(+), 146 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fsdp_module.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fsdp_module.py index e6ce5b0037f..5360b916169 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fsdp_module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fsdp_module.py @@ -31,13 +31,13 @@ class FsdpModule: _parameter_groups: tuple[ParameterGroup, ...] _ready_grad_parameters: set[nn.Parameter] - num_training_parameters: int + _num_training_parameters: int def __init__( self, mesh: DeviceMesh, placements: Placements, mixed_precision_policy: MixedPrecisionPolicy ) -> None: """Initialize FSDP runtime state on an already-constructed module.""" - owned_parameters = _materialize_and_collect_owned_parameters(self, _mesh_device(mesh)) + owned_parameters = _collect_owned_parameters(self) axis_indices = tuple(_axis_index(mesh, axis) for axis in placements.dp_axes) assert axis_indices == tuple( range(mesh.ndim) @@ -54,7 +54,7 @@ def __init__( ] self._parameter_groups = tuple(parameter_groups) self._ready_grad_parameters = set() - self.num_training_parameters = sum( + self._num_training_parameters = sum( len(group.sharded_parameters) for group in self._parameter_groups if group.requires_grad ) self._register_hooks() @@ -77,7 +77,7 @@ def _register_hooks(self) -> None: def _make_grad_hook(self, parameter: nn.Parameter) -> Callable[[nn.Parameter], None]: def grad_hook(_parameter: nn.Parameter) -> None: self._ready_grad_parameters.add(parameter) - if len(self._ready_grad_parameters) == self.num_training_parameters: + if len(self._ready_grad_parameters) == self._num_training_parameters: self.post_backward() return grad_hook @@ -86,6 +86,7 @@ def pre_forward(self) -> None: """Prepare full parameters for forward compute.""" self._ready_grad_parameters.clear() for group in self._parameter_groups: + group.sync_model_weight_from_main_weight() group.unshard_parameters() def post_forward(self) -> None: @@ -126,40 +127,12 @@ def _axis_index(mesh: DeviceMesh, axis: MeshAxis) -> int: return dim_names.index(axis) -def _mesh_device(mesh: DeviceMesh) -> torch.device: - if mesh.device_type == "cuda": - return torch.device("cuda", torch.cuda.current_device()) - return torch.device(mesh.device_type) - - -def _materialize_and_collect_owned_parameters( - root_module: nn.Module, device: torch.device -) -> dict[str, nn.Parameter]: +def _collect_owned_parameters(root_module: nn.Module) -> dict[str, nn.Parameter]: parameters: dict[str, nn.Parameter] = {} def visit(submodule: nn.Module, submodule_fqn: str) -> None: direct_parameters = list(submodule.named_parameters(recurse=False)) - if any(parameter.is_meta for _, parameter in direct_parameters): - if any(not parameter.is_meta for _, parameter in direct_parameters): - raise ValueError( - f"Module {submodule_fqn!r} mixes meta and non-meta direct parameters. " - "Initialize all direct parameters on meta or none of them." - ) - submodule.to_empty(device=device, recurse=False) - with torch.no_grad(): - if hasattr(submodule, "reset_parameters"): - submodule.reset_parameters() - elif hasattr(submodule, "_reset_parameters"): - submodule._reset_parameters() - else: - raise ValueError( - f"Module {submodule_fqn!r} does not have " - "reset_parameters or _reset_parameters." - ) - # Module.to_empty may replace Parameters, so collect direct parameters again. - direct_parameters = list(submodule.named_parameters(recurse=False)) - for local_parameter_name, parameter in direct_parameters: parameter_fqn = ( f"{submodule_fqn}.{local_parameter_name}" if submodule_fqn else local_parameter_name diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py index 2dad4ca7732..d5cdde1f23e 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py @@ -81,11 +81,6 @@ def __init__( self.dtype = first_parameter.dtype self.requires_grad = first_parameter.requires_grad for name, parameter in parameters.items(): - if parameter.is_meta: - raise ValueError( - f"Expected parameter {name!r} to be materialized before " - "ParameterGroup construction." - ) if parameter.dtype != self.dtype: raise ValueError( f"Expected parameter {name!r} to have dtype {self.dtype}, " @@ -110,7 +105,7 @@ def __init__( placements=[Replicate()] * self.mesh.ndim, tensor_shapes=tensor_shapes, dtype=self.dtype, - device=self.main_weight.local_buffer.device, + device=self.main_weight.device, ) if main_weight_dtype == self.dtype and main_weight_placements == model_weight_placements: self.model_weight = self.main_weight @@ -120,7 +115,7 @@ def __init__( placements=model_weight_placements, tensor_shapes=tensor_shapes, dtype=self.dtype, - device=self.main_weight.local_buffer.device, + device=self.main_weight.device, ) self.main_grad = None @@ -131,7 +126,7 @@ def __init__( placements=main_grad_placements, tensor_shapes=self.main_weight.layout.tensor_shapes, dtype=grad_dtype, - device=self.main_weight.local_buffer.device, + device=self.main_weight.device, ) assert self.main_grad.layout == self.main_weight.layout, ( "main_grad is built from main_weight tensor shapes on the same mesh, " @@ -147,7 +142,7 @@ def __init__( sharded_parameters: list[nn.Parameter] = [] unsharded_parameters: list[nn.Parameter] = [] - main_grad_dtype = self.main_grad.local_buffer.dtype if self.main_grad is not None else None + main_grad_dtype = self.main_grad.dtype if self.main_grad is not None else None for index, parameter in enumerate(parameters.values()): parameter.data = self._unsharded_model_weight.get_local_tensor(index) parameter.grad = None @@ -183,13 +178,12 @@ def sync_model_weight_from_main_weight(self) -> None: if self.main_weight is self.model_weight: return - self.main_weight.cast(self.model_weight.local_buffer.dtype).redistribute( + self.main_weight.cast(self.model_weight.dtype).redistribute( self.model_weight.placements, out=self.model_weight ) def unshard_parameters(self) -> None: """Install full parameters for local compute.""" - self.sync_model_weight_from_main_weight() self._unsharded_model_weight.reallocate_storage() # This buffer backs unsharded Parameters whose views may be saved by autograd. # Materializing FSDP-managed storage should not look like a user mutation. @@ -238,7 +232,7 @@ def has_grad(parameters: Iterable[nn.Parameter]) -> bool: has_sharded_grads = has_grad(self.sharded_parameters) can_reduce_into_main_grad = ( not has_sharded_grads - and partial_grad.local_buffer.dtype == self.main_grad.local_buffer.dtype + and partial_grad.dtype == self.main_grad.dtype ) if can_reduce_into_main_grad: partial_grad.redistribute(self.main_grad.placements, out=self.main_grad) diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py index d29e1ce844b..957f78a5665 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py +++ b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py @@ -3,13 +3,9 @@ """Unit tests for the minimal Megatron-FSDP path.""" import logging -import os -from collections.abc import Iterator -from dataclasses import dataclass import pytest import torch -import torch.distributed as dist from torch import nn from torch.distributed.device_mesh import init_device_mesh from torch.distributed.tensor import DTensor @@ -24,37 +20,6 @@ logger = logging.getLogger(__name__) -@dataclass(frozen=True) -class DistributedSetup: - """Per-rank distributed test setup.""" - - rank: int - world_size: int - device: torch.device - - -@pytest.fixture(scope="module") -def setup() -> Iterator[DistributedSetup]: - """Read torchrun rank state and set up this rank's local device.""" - if "RANK" not in os.environ or "WORLD_SIZE" not in os.environ: - pytest.skip("Not running under torchrun.") - - rank = int(os.environ["RANK"]) - world_size = int(os.environ["WORLD_SIZE"]) - local_rank = int(os.environ.get("LOCAL_RANK", rank)) - - if torch.cuda.is_available(): - torch.cuda.set_device(local_rank) - device = torch.device(f"cuda:{local_rank}") - else: - device = torch.device("cpu") - - yield DistributedSetup(rank=rank, world_size=world_size, device=device) - - if dist.is_initialized(): - dist.destroy_process_group() - - class TinyModel(nn.Module): """Small model with two separately shardable units.""" @@ -113,18 +78,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return SaveNonLeafWeightView.apply(x, weight_view) -class ConstantMetaModel(nn.Module): - """Model whose meta parameter is initialized by reset_parameters().""" - - def __init__(self) -> None: - super().__init__() - self.weight = nn.Parameter(torch.empty(4, 4, device="meta")) - - def reset_parameters(self) -> None: - """Initialize the weight to a deterministic value.""" - self.weight.fill_(3.0) - - def _flat_placements() -> Placements: return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()]) @@ -134,15 +87,18 @@ def _mb(num_bytes: int) -> str: @pytest.mark.distributed -def test_fully_shard_losses_match_baseline(setup: DistributedSetup): +def test_fully_shard_losses_match_baseline(distributed_setup): """Minimal per-module FSDP training should match single-rank SGD.""" - if setup.world_size < 2: + rank = distributed_setup.rank + world_size = distributed_setup.world_size + device = distributed_setup.device + if world_size < 2: pytest.skip("This test requires at least 2 ranks.") - mesh = init_device_mesh(setup.device.type, (setup.world_size,)) + mesh = init_device_mesh(device.type, (world_size,)) torch.manual_seed(1234) - baseline = TinyModel().to(setup.device) - model = TinyModel().to(setup.device) + baseline = TinyModel().to(device) + model = TinyModel().to(device) model.load_state_dict(baseline.state_dict()) fully_shard(model.fc1, mesh=mesh, placements=_flat_placements()) @@ -150,8 +106,8 @@ def test_fully_shard_losses_match_baseline(setup: DistributedSetup): baseline_optimizer = torch.optim.SGD(baseline.parameters(), lr=0.05) optimizer = torch.optim.SGD(model.parameters(), lr=0.05) - x = torch.randn(3, 8, device=setup.device) - target = torch.randn(3, 4, device=setup.device) + x = torch.randn(3, 8, device=device) + target = torch.randn(3, 4, device=device) for step in range(5): baseline_optimizer.zero_grad() @@ -161,7 +117,7 @@ def test_fully_shard_losses_match_baseline(setup: DistributedSetup): loss = torch.nn.functional.mse_loss(model(x), target) logger.info( "FSDP train parity: rank=%s, step=%s, baseline_loss=%s, sharded_loss=%s", - setup.rank, + rank, step, baseline_loss.item(), loss.item(), @@ -175,13 +131,15 @@ def test_fully_shard_losses_match_baseline(setup: DistributedSetup): @pytest.mark.distributed -def test_nested_fully_shard_excludes_child_owned_parameters(setup: DistributedSetup): +def test_nested_fully_shard_excludes_child_owned_parameters(distributed_setup): """An outer FSDP unit owns direct parameters but not nested child-unit parameters.""" - if setup.world_size < 2: + world_size = distributed_setup.world_size + device = distributed_setup.device + if world_size < 2: pytest.skip("This test requires at least 2 ranks.") - mesh = init_device_mesh(setup.device.type, (setup.world_size,)) - model = NestedModel().to(setup.device) + mesh = init_device_mesh(device.type, (world_size,)) + model = NestedModel().to(device) fully_shard(model.inner, mesh=mesh, placements=_flat_placements()) fully_shard(model, mesh=mesh, placements=_flat_placements()) @@ -196,13 +154,15 @@ def test_nested_fully_shard_excludes_child_owned_parameters(setup: DistributedSe @pytest.mark.distributed -def test_frozen_parameter_group_does_not_allocate_main_grad(setup: DistributedSetup): +def test_frozen_parameter_group_does_not_allocate_main_grad(distributed_setup): """A non-trainable parameter group should not allocate persistent main gradients.""" - if setup.world_size < 2: + world_size = distributed_setup.world_size + device = distributed_setup.device + if world_size < 2: pytest.skip("This test requires at least 2 ranks.") - mesh = init_device_mesh(setup.device.type, (setup.world_size,)) - model = nn.Linear(4, 4, bias=False).to(setup.device) + mesh = init_device_mesh(device.type, (world_size,)) + model = nn.Linear(4, 4, bias=False).to(device) model.weight.requires_grad_(False) fully_shard(model, mesh=mesh, placements=_flat_placements()) @@ -215,36 +175,41 @@ def test_frozen_parameter_group_does_not_allocate_main_grad(setup: DistributedSe pytest.mark.distributed -def test_backward_averages_across_dp_and_accumulates_across_calls(setup: DistributedSetup): +def test_backward_averages_across_dp_and_accumulates_across_calls(distributed_setup): """Each backward averages over DP ranks; repeated backwards accumulate by summing.""" - if setup.world_size < 2: + rank = distributed_setup.rank + world_size = distributed_setup.world_size + device = distributed_setup.device + if world_size < 2: pytest.skip("This test requires at least 2 ranks.") - mesh = init_device_mesh(setup.device.type, (setup.world_size,)) - model = nn.Linear(1, setup.world_size, bias=False).to(setup.device) + mesh = init_device_mesh(device.type, (world_size,)) + model = nn.Linear(1, world_size, bias=False).to(device) with torch.no_grad(): model.weight.fill_(1.0) fully_shard(model, mesh=mesh, placements=_flat_placements()) - x = torch.full((1, 1), float(setup.rank + 1), device=setup.device) + x = torch.full((1, 1), float(rank + 1), device=device) model(x).sum().backward() model(x).sum().backward() assert isinstance(model.weight.grad, DTensor) local_grad = model.weight.grad.to_local() - expected = torch.full_like(local_grad, float(setup.world_size + 1)) + expected = torch.full_like(local_grad, float(world_size + 1)) torch.testing.assert_close(local_grad, expected, rtol=0, atol=0) @pytest.mark.distributed -def test_next_forward_uses_optimizer_updated_weights(setup: DistributedSetup): +def test_next_forward_uses_optimizer_updated_weights(distributed_setup): """The next forward should observe weights updated by the previous optimizer step.""" - if setup.world_size < 2: + world_size = distributed_setup.world_size + device = distributed_setup.device + if world_size < 2: pytest.skip("This test requires at least 2 ranks.") - mesh = init_device_mesh(setup.device.type, (setup.world_size,)) - model = nn.Linear(1, setup.world_size, bias=False, dtype=torch.bfloat16).to(setup.device) + mesh = init_device_mesh(device.type, (world_size,)) + model = nn.Linear(1, world_size, bias=False, dtype=torch.bfloat16).to(device) with torch.no_grad(): model.weight.fill_(1.0) @@ -257,7 +222,7 @@ def test_next_forward_uses_optimizer_updated_weights(setup: DistributedSetup): # SGD's foreach/fused CUDA paths require matching parameter and gradient dtypes. # Use the scalar path to exercise FP32 main weights with default BF16 main grads. optimizer = torch.optim.SGD(model.parameters(), lr=0.25, foreach=False) - x = torch.ones(1, 1, device=setup.device, dtype=torch.bfloat16) + x = torch.ones(1, 1, device=device, dtype=torch.bfloat16) def train_iteration() -> torch.Tensor: optimizer.zero_grad(set_to_none=True) @@ -274,34 +239,41 @@ def train_iteration() -> torch.Tensor: @pytest.mark.distributed -def test_meta_parameters_initialize_with_reset_parameters(setup: DistributedSetup): - """Meta parameters should be replaced by sharded DTensors and initialized in place.""" - if setup.world_size < 2: +def test_cpu_initialized_parameters_shard_to_mesh_device(distributed_setup): + """CPU-initialized parameters should be sharded with their real values.""" + world_size = distributed_setup.world_size + device = distributed_setup.device + if world_size < 2: pytest.skip("This test requires at least 2 ranks.") - mesh = init_device_mesh(setup.device.type, (setup.world_size,)) - model = ConstantMetaModel() + mesh = init_device_mesh(device.type, (world_size,)) + model = nn.Linear(4, 4, bias=False) + with torch.no_grad(): + model.weight.fill_(3.0) + expected_weight = model.weight.detach().to(device) fully_shard(model, mesh=mesh, placements=_flat_placements()) (group,) = model.parameter_groups() full_weight = group.model_weight.allgather(0).get_local_tensor(0) - assert not full_weight.is_meta - torch.testing.assert_close(full_weight, torch.full_like(full_weight, 3.0)) + assert full_weight.device.type == device.type + torch.testing.assert_close(full_weight, expected_weight) @pytest.mark.distributed -def test_non_leaf_parameter_view_survives_storage_resize(setup: DistributedSetup): +def test_non_leaf_parameter_view_survives_storage_resize(distributed_setup): """A non-leaf parameter view saved for backward should survive full-storage resize.""" - if setup.world_size < 2: + world_size = distributed_setup.world_size + device = distributed_setup.device + if world_size < 2: pytest.skip("This test requires at least 2 ranks.") - mesh = init_device_mesh(setup.device.type, (setup.world_size,)) - model = NonLeafViewModel().to(setup.device) + mesh = init_device_mesh(device.type, (world_size,)) + model = NonLeafViewModel().to(device) fully_shard(model, mesh=mesh, placements=_flat_placements()) group = model.parameter_groups()[0] - x = torch.randn(8, device=setup.device, requires_grad=True) + x = torch.randn(8, device=device, requires_grad=True) loss = model(x).sum() assert group._unsharded_model_weight is not None @@ -315,14 +287,17 @@ def test_non_leaf_parameter_view_survives_storage_resize(setup: DistributedSetup @pytest.mark.distributed -def test_fully_shard_reduces_peak_training_memory(setup: DistributedSetup): +def test_fully_shard_reduces_peak_training_memory(distributed_setup): """Per-layer FSDP should reduce peak CUDA memory during training.""" - if setup.world_size < 2: + rank = distributed_setup.rank + world_size = distributed_setup.world_size + device = distributed_setup.device + if world_size < 2: pytest.skip("This test requires at least 2 ranks.") - if setup.device.type != "cuda": + if device.type != "cuda": pytest.skip("Peak memory verification requires CUDA.") - mesh = init_device_mesh(setup.device.type, (setup.world_size,)) + mesh = init_device_mesh(device.type, (world_size,)) dim = 1024 layers = 16 batch = 8 @@ -337,14 +312,14 @@ def train_steps(model: nn.Module, optimizer: torch.optim.Optimizer, x: torch.Ten torch.manual_seed(4321) baseline = nn.Sequential(*[nn.Linear(dim, dim, dtype=dtype) for _ in range(layers)]).to( - setup.device + device ) baseline_optimizer = torch.optim.AdamW(baseline.parameters(), lr=0.01) - x = torch.randn(batch, dim, device=setup.device, dtype=dtype) - torch.cuda.reset_peak_memory_stats(setup.device) + x = torch.randn(batch, dim, device=device, dtype=dtype) + torch.cuda.reset_peak_memory_stats(device) train_steps(baseline, baseline_optimizer, x) - torch.cuda.synchronize(setup.device) - baseline_peak = torch.cuda.max_memory_allocated(setup.device) + torch.cuda.synchronize(device) + baseline_peak = torch.cuda.max_memory_allocated(device) del baseline_optimizer del baseline @@ -353,7 +328,7 @@ def train_steps(model: nn.Module, optimizer: torch.optim.Optimizer, x: torch.Ten torch.manual_seed(4321) model = nn.Sequential(*[nn.Linear(dim, dim, dtype=dtype) for _ in range(layers)]).to( - setup.device + device ) for layer in model: fully_shard( @@ -367,14 +342,14 @@ def train_steps(model: nn.Module, optimizer: torch.optim.Optimizer, x: torch.Ten optimizer = torch.optim.AdamW(model.parameters(), lr=0.01) torch.cuda.empty_cache() - x = torch.randn(batch, dim, device=setup.device, dtype=dtype) - torch.cuda.reset_peak_memory_stats(setup.device) + x = torch.randn(batch, dim, device=device, dtype=dtype) + torch.cuda.reset_peak_memory_stats(device) train_steps(model, optimizer, x) - torch.cuda.synchronize(setup.device) - sharded_peak = torch.cuda.max_memory_allocated(setup.device) + torch.cuda.synchronize(device) + sharded_peak = torch.cuda.max_memory_allocated(device) logger.info( "FSDP peak memory: rank=%s, baseline=%s, sharded=%s", - setup.rank, + rank, _mb(baseline_peak), _mb(sharded_peak), ) From 23c47a847fcdde714903f2bffb3d630388bbd6e1 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Wed, 17 Jun 2026 06:43:15 +0000 Subject: [PATCH 14/23] Clarify FSDP version counter preservation Signed-off-by: Jingyue Wu --- .../fsdp/src/megatron_fsdp/experimental/parameter_group.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py index d5cdde1f23e..085b8c3c2c4 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py @@ -186,7 +186,10 @@ def unshard_parameters(self) -> None: """Install full parameters for local compute.""" self._unsharded_model_weight.reallocate_storage() # This buffer backs unsharded Parameters whose views may be saved by autograd. - # Materializing FSDP-managed storage should not look like a user mutation. + # Autograd records a tensor's version counter when saving it for backward, and + # in-place writes like the out= redistribution below increment that counter even + # under no_grad. Without preserving it, backward can fail with "modified by an + # inplace operation" even though FSDP only materialized internal storage. with torch.autograd._unsafe_preserve_version_counter( self._unsharded_model_weight.local_buffer ): From 08e84c62708a2a81100bc3bb820f7cc7ae3d765a Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Wed, 17 Jun 2026 07:01:06 +0000 Subject: [PATCH 15/23] Rename experimental FSDP runtime types Signed-off-by: Jingyue Wu --- .../fsdp/src/megatron_fsdp/experimental/__init__.py | 6 +++--- .../fsdp/src/megatron_fsdp/experimental/dbuffer.py | 2 +- .../fsdp/src/megatron_fsdp/experimental/fully_shard.py | 2 +- .../experimental/{fsdp_module.py => module.py} | 8 ++++---- .../src/megatron_fsdp/experimental/parameter_group.py | 9 ++++----- .../megatron_fsdp/test_experimental_fully_shard.py | 8 ++------ 6 files changed, 15 insertions(+), 20 deletions(-) rename megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/{fsdp_module.py => module.py} (96%) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py index 47db6401d8e..87fd3dac52b 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py @@ -15,16 +15,16 @@ """Experimental Megatron-FSDP implementation.""" from .dbuffer import DBuffer -from .fsdp_module import FsdpModule from .fully_shard import fully_shard -from .parameter_group import ParameterGroup +from .module import FsdpModule +from .parameter_group import FsdpParameterGroup from .placement import Flat, Partial, Placement, Placements, Replicate __all__ = [ "DBuffer", "Flat", "FsdpModule", - "ParameterGroup", + "FsdpParameterGroup", "Partial", "Placement", "Placements", diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py index 5305da55994..d690313ba30 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py @@ -66,7 +66,7 @@ class DBuffer: """ # DBuffer owns only the data-parallel sub-mesh. Higher-level callers, such as - # ParameterGroup, should extend returned DTensors with tensor-parallel mesh axes + # FsdpParameterGroup, should extend returned DTensors with tensor-parallel mesh axes # because TP sharding metadata lives on nn.Parameter in MCore/TransformerEngine. mesh: DeviceMesh placements: tuple[Placement, ...] diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py index 3ab3c0e6d71..19516ea66ff 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py @@ -18,7 +18,7 @@ from torch.distributed import DeviceMesh from ..mixed_precision import MixedPrecisionPolicy -from .fsdp_module import FsdpModule +from .module import FsdpModule from .placement import Placements diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fsdp_module.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py similarity index 96% rename from megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fsdp_module.py rename to megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py index 5360b916169..8907f0764b4 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fsdp_module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py @@ -22,14 +22,14 @@ from torch.distributed import DeviceMesh from ..mixed_precision import MixedPrecisionPolicy -from .parameter_group import ParameterGroup, contained_in_parameter_group +from .parameter_group import FsdpParameterGroup, contained_in_parameter_group from .placement import MeshAxis, Placements class FsdpModule: """Mixin attached to modules managed by the minimal FSDP path.""" - _parameter_groups: tuple[ParameterGroup, ...] + _parameter_groups: tuple[FsdpParameterGroup, ...] _ready_grad_parameters: set[nn.Parameter] _num_training_parameters: int @@ -43,7 +43,7 @@ def __init__( range(mesh.ndim) ), "FSDP requires dp_axes to match every mesh axis in mesh order for now." parameter_groups = [ - ParameterGroup( + FsdpParameterGroup( owning_module=self, parameters=group_parameters, mesh=mesh, @@ -107,7 +107,7 @@ def post_backward(self) -> None: group.reshard_parameters() self._ready_grad_parameters.clear() - def parameter_groups(self) -> tuple[ParameterGroup, ...]: + def parameter_groups(self) -> tuple[FsdpParameterGroup, ...]: """Return parameter groups owned by this FSDP unit.""" return self._parameter_groups diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py index 085b8c3c2c4..a0209f8ad92 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py @@ -29,11 +29,11 @@ def contained_in_parameter_group(parameter: nn.Parameter) -> bool: - """Return whether a parameter is already owned by a ParameterGroup.""" + """Return whether a parameter is already owned by an FsdpParameterGroup.""" return hasattr(parameter, _CONTAINING_PARAMETER_GROUP_ATTR) -class ParameterGroup: +class FsdpParameterGroup: """A dtype and requires-grad homogeneous group of FSDP-owned parameters.""" owning_module: nn.Module @@ -66,7 +66,7 @@ def __init__( mixed_precision_policy: Precision policy for main weights and gradients. """ if not parameters: - raise ValueError("ParameterGroup requires at least one parameter.") + raise ValueError("FsdpParameterGroup requires at least one parameter.") model_weight_placements = tuple(placements.parameter) main_grad_placements = tuple(placements.gradient) @@ -234,8 +234,7 @@ def has_grad(parameters: Iterable[nn.Parameter]) -> bool: # leaves sharded grads installed, so this backward accumulates into main_grad. has_sharded_grads = has_grad(self.sharded_parameters) can_reduce_into_main_grad = ( - not has_sharded_grads - and partial_grad.dtype == self.main_grad.dtype + not has_sharded_grads and partial_grad.dtype == self.main_grad.dtype ) if can_reduce_into_main_grad: partial_grad.redistribute(self.main_grad.placements, out=self.main_grad) diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py index 957f78a5665..0f20b28894b 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py +++ b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py @@ -311,9 +311,7 @@ def train_steps(model: nn.Module, optimizer: torch.optim.Optimizer, x: torch.Ten optimizer.step() torch.manual_seed(4321) - baseline = nn.Sequential(*[nn.Linear(dim, dim, dtype=dtype) for _ in range(layers)]).to( - device - ) + baseline = nn.Sequential(*[nn.Linear(dim, dim, dtype=dtype) for _ in range(layers)]).to(device) baseline_optimizer = torch.optim.AdamW(baseline.parameters(), lr=0.01) x = torch.randn(batch, dim, device=device, dtype=dtype) torch.cuda.reset_peak_memory_stats(device) @@ -327,9 +325,7 @@ def train_steps(model: nn.Module, optimizer: torch.optim.Optimizer, x: torch.Ten torch.cuda.empty_cache() torch.manual_seed(4321) - model = nn.Sequential(*[nn.Linear(dim, dim, dtype=dtype) for _ in range(layers)]).to( - device - ) + model = nn.Sequential(*[nn.Linear(dim, dim, dtype=dtype) for _ in range(layers)]).to(device) for layer in model: fully_shard( layer, From ea39f57762995d52d9f2de10d7d289a4a7ad86e0 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Wed, 17 Jun 2026 07:08:00 +0000 Subject: [PATCH 16/23] Document main_grad allocation lifetime Signed-off-by: Jingyue Wu --- .../fsdp/src/megatron_fsdp/experimental/parameter_group.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py index a0209f8ad92..05aa122ae5a 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py @@ -121,6 +121,11 @@ def __init__( self.main_grad = None if self.requires_grad: grad_dtype = mixed_precision_policy.main_grads_dtype or self.dtype + # Keep main_grad persistent for the initial implementation. For micro-batch + # size 1, this allocation could be delayed until post_backward and then + # eagerly deallocated right after optimizer.step(), avoiding main_grad + # storage during forward. That requires a separate lifetime contract with + # the optimizer, so this version keeps the simpler persistent buffer. self.main_grad = DBuffer( mesh=self.mesh, placements=main_grad_placements, From e80d570ff6c5db4f70fb98c7b68fa8ab68ebcd39 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Wed, 17 Jun 2026 07:33:47 +0000 Subject: [PATCH 17/23] Document post-backward reshard storage choice Signed-off-by: Jingyue Wu --- .../fsdp/src/megatron_fsdp/experimental/parameter_group.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py index 05aa122ae5a..a2c7bd0bccb 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py @@ -206,6 +206,12 @@ def unshard_parameters(self) -> None: def reshard_parameters(self) -> None: """Install sharded DTensor parameters on the owning modules.""" self._switch_to_sharded_parameters() + # At post-backward time, replacing unsharded parameter .data with size-0 + # empty tensors would also be safe: autograd has consumed the saved + # forward views. That alternative is not much cleaner than releasing + # this storage, and splitting post-forward and post-backward reshard + # behavior would make the caller code less clean, so keep the shared + # storage-release path. self._unsharded_model_weight.release_storage() def reduce_gradients(self) -> None: From 99b12617d7f3bf6ade5bd507f8f22b36e5eea029 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sun, 21 Jun 2026 16:24:47 +0000 Subject: [PATCH 18/23] Remove experimental FSDP distributed pytest markers Signed-off-by: Jingyue Wu --- .../megatron_fsdp/test_experimental_fully_shard.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py index 0f20b28894b..86ccea8d942 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py +++ b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py @@ -86,7 +86,6 @@ def _mb(num_bytes: int) -> str: return f"{num_bytes / 1024**2:.2f} MB" -@pytest.mark.distributed def test_fully_shard_losses_match_baseline(distributed_setup): """Minimal per-module FSDP training should match single-rank SGD.""" rank = distributed_setup.rank @@ -130,7 +129,6 @@ def test_fully_shard_losses_match_baseline(distributed_setup): optimizer.step() -@pytest.mark.distributed def test_nested_fully_shard_excludes_child_owned_parameters(distributed_setup): """An outer FSDP unit owns direct parameters but not nested child-unit parameters.""" world_size = distributed_setup.world_size @@ -153,7 +151,6 @@ def test_nested_fully_shard_excludes_child_owned_parameters(distributed_setup): assert outer_names == ["bias"] -@pytest.mark.distributed def test_frozen_parameter_group_does_not_allocate_main_grad(distributed_setup): """A non-trainable parameter group should not allocate persistent main gradients.""" world_size = distributed_setup.world_size @@ -172,9 +169,6 @@ def test_frozen_parameter_group_does_not_allocate_main_grad(distributed_setup): assert group.main_grad is None -pytest.mark.distributed - - def test_backward_averages_across_dp_and_accumulates_across_calls(distributed_setup): """Each backward averages over DP ranks; repeated backwards accumulate by summing.""" rank = distributed_setup.rank @@ -200,7 +194,6 @@ def test_backward_averages_across_dp_and_accumulates_across_calls(distributed_se torch.testing.assert_close(local_grad, expected, rtol=0, atol=0) -@pytest.mark.distributed def test_next_forward_uses_optimizer_updated_weights(distributed_setup): """The next forward should observe weights updated by the previous optimizer step.""" world_size = distributed_setup.world_size @@ -238,7 +231,6 @@ def train_iteration() -> torch.Tensor: torch.testing.assert_close(second_loss, first_loss) -@pytest.mark.distributed def test_cpu_initialized_parameters_shard_to_mesh_device(distributed_setup): """CPU-initialized parameters should be sharded with their real values.""" world_size = distributed_setup.world_size @@ -260,7 +252,6 @@ def test_cpu_initialized_parameters_shard_to_mesh_device(distributed_setup): torch.testing.assert_close(full_weight, expected_weight) -@pytest.mark.distributed def test_non_leaf_parameter_view_survives_storage_resize(distributed_setup): """A non-leaf parameter view saved for backward should survive full-storage resize.""" world_size = distributed_setup.world_size @@ -286,7 +277,6 @@ def test_non_leaf_parameter_view_survives_storage_resize(distributed_setup): assert group._unsharded_model_weight.local_buffer.untyped_storage().nbytes() == 0 -@pytest.mark.distributed def test_fully_shard_reduces_peak_training_memory(distributed_setup): """Per-layer FSDP should reduce peak CUDA memory during training.""" rank = distributed_setup.rank From c12e418d9d2c1a349126b002960abac353f106f1 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sun, 21 Jun 2026 16:26:51 +0000 Subject: [PATCH 19/23] Remove DBuffer distributed pytest markers Signed-off-by: Jingyue Wu --- .../distributed/megatron_fsdp/test_dbuffer.py | 25 ------------------- 1 file changed, 25 deletions(-) diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_dbuffer.py b/tests/unit_tests/distributed/megatron_fsdp/test_dbuffer.py index d267298b94d..a90a39b0315 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/test_dbuffer.py +++ b/tests/unit_tests/distributed/megatron_fsdp/test_dbuffer.py @@ -31,7 +31,6 @@ def _assert_dbuffer_local_tensors_close(buffer: DBuffer, expected: Iterable[torc torch.testing.assert_close(buffer.get_local_tensor(index), tensor) -@pytest.mark.distributed def test_dbuffer_layout_pads_to_lcm_times_dp_size_and_fills_gaps(distributed_setup): """DBuffer layout returns element offsets and pads to LCM * DP size.""" if distributed_setup.world_size < 2: @@ -53,7 +52,6 @@ def test_dbuffer_layout_pads_to_lcm_times_dp_size_and_fills_gaps(distributed_set assert buffer.layout.size == 48 -@pytest.mark.distributed def test_dbuffer_layout_aligns_fragment_offsets_to_rows(distributed_setup): """DBuffer layout keeps small tensors aligned to their non-leading dimensions.""" if distributed_setup.world_size < 2: @@ -74,7 +72,6 @@ def test_dbuffer_layout_aligns_fragment_offsets_to_rows(distributed_setup): assert buffer.layout.size == 24 -@pytest.mark.distributed def test_compute_layout_fills_lcm_padding_gaps(distributed_setup): """LCM packing fills row-aligned padding gaps on a 5-rank flat-sharded mesh.""" if distributed_setup.world_size < 5: @@ -118,7 +115,6 @@ def test_compute_layout_fills_lcm_padding_gaps(distributed_setup): assert buffer.get_dtensor(index).shape == shapes[index] -@pytest.mark.distributed def test_constructor_allocates_local_buffer(distributed_setup): """DBuffer allocates local storage from shape, mesh, placement, dtype, and device.""" mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) @@ -155,7 +151,6 @@ def test_constructor_allocates_local_buffer(distributed_setup): assert sharded_buffer.local_buffer.device == distributed_setup.device -@pytest.mark.distributed def test_cast_to_same_dtype_returns_self(distributed_setup): """DBuffer.cast returns self when the dtype already matches.""" mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) @@ -165,7 +160,6 @@ def test_cast_to_same_dtype_returns_self(distributed_setup): assert buffer.cast(torch.float32) is buffer -@pytest.mark.distributed def test_cast_preserves_layout_and_casts_values(distributed_setup): """DBuffer.cast preserves layout metadata and casts local values.""" mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) @@ -185,7 +179,6 @@ def test_cast_preserves_layout_and_casts_values(distributed_setup): ) -@pytest.mark.distributed def test_release_and_reallocate_storage_preserves_buffer_views(distributed_setup): """DBuffer storage can be released and reallocated without replacing existing views.""" mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) @@ -214,7 +207,6 @@ def test_release_and_reallocate_storage_preserves_buffer_views(distributed_setup torch.testing.assert_close(tensor_view, torch.full_like(tensor_view, 7.0)) -@pytest.mark.distributed def test_from_local_reuses_required_local_buffer(distributed_setup): """DBuffer.from_local reuses caller-provided local storage without allocation.""" mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) @@ -235,7 +227,6 @@ def test_from_local_reuses_required_local_buffer(distributed_setup): _assert_dbuffer_local_tensors_close(sharded_buffer.allgather(0), tensors) -@pytest.mark.distributed def test_replicate_get_local_tensor_and_dtensor(distributed_setup): """Replicated DBuffer returns full local tensors and replicated DTensors.""" mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) @@ -248,7 +239,6 @@ def test_replicate_get_local_tensor_and_dtensor(distributed_setup): torch.testing.assert_close(dtensor.to_local(), tensors[0], rtol=0, atol=0) -@pytest.mark.distributed def test_distribute_tensors_moves_inputs_to_mesh_device(distributed_setup): """distribute_tensors moves full input tensors to the mesh device type.""" mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) @@ -262,7 +252,6 @@ def test_distribute_tensors_moves_inputs_to_mesh_device(distributed_setup): ) -@pytest.mark.distributed def test_distribute_tensors_detaches_and_contiguizes_inputs(distributed_setup): """distribute_tensors treats input tensors as detached contiguous values.""" mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) @@ -280,7 +269,6 @@ def test_distribute_tensors_detaches_and_contiguizes_inputs(distributed_setup): ) -@pytest.mark.distributed def test_sharded_allgather_round_trip(distributed_setup): """Sharded buffers round-trip through all-gather as contiguous tensor fragments.""" mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) @@ -299,7 +287,6 @@ def test_sharded_allgather_round_trip(distributed_setup): _assert_dbuffer_local_tensors_close(replicated_buffer, tensors) -@pytest.mark.distributed def test_sharded_allgather_into_existing_buffer(distributed_setup): """Sharded buffers can all-gather directly into a preallocated replicated buffer.""" mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) @@ -321,7 +308,6 @@ def test_sharded_allgather_into_existing_buffer(distributed_setup): _assert_dbuffer_local_tensors_close(destination, tensors) -@pytest.mark.distributed def test_mesh_axis_must_be_non_negative_int(distributed_setup): """DBuffer communication methods require explicit non-negative integer mesh axes.""" mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) @@ -341,7 +327,6 @@ def test_mesh_axis_must_be_non_negative_int(distributed_setup): buffer.allgather(-1) -@pytest.mark.distributed def test_replicate_scatter_round_trip(distributed_setup): """Replicated buffers locally chunk into sharded buffers and all-gather back.""" mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) @@ -375,7 +360,6 @@ def test_replicate_scatter_round_trip(distributed_setup): _assert_dbuffer_local_tensors_close(sharded_buffer.allgather(0), tensors) -@pytest.mark.distributed def test_partial_allreduce(distributed_setup): """Partial buffers all-reduce into replicated buffers.""" mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) @@ -396,7 +380,6 @@ def test_partial_allreduce(distributed_setup): _assert_dbuffer_local_tensors_close(replicated_buffer, expected) -@pytest.mark.distributed def test_partial_allreduce_average(distributed_setup): """Partial buffers can all-reduce with AVG.""" mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) @@ -427,7 +410,6 @@ def test_partial_allreduce_average(distributed_setup): _assert_dbuffer_local_tensors_close(replicated_buffer, expected) -@pytest.mark.distributed def test_partial_reduce_scatter_to_flat(distributed_setup): """Partial buffers reduce-scatter into sharded buffers.""" mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) @@ -461,7 +443,6 @@ def test_partial_reduce_scatter_to_flat(distributed_setup): _assert_dbuffer_local_tensors_close(replicated_buffer, expected_tensors) -@pytest.mark.distributed def test_partial_reduce_scatter_to_flat_average(distributed_setup): """Partial buffers can reduce-scatter with AVG.""" mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) @@ -489,7 +470,6 @@ def test_partial_reduce_scatter_to_flat_average(distributed_setup): _assert_dbuffer_local_tensors_close(replicated_buffer, expected_tensors) -@pytest.mark.distributed def test_get_dtensor_from_sharded_buffer(distributed_setup): """Sharded DBuffer exposes per-tensor local shards as DTensors.""" mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) @@ -504,7 +484,6 @@ def test_get_dtensor_from_sharded_buffer(distributed_setup): assert dtensor.shape == tensors[0].shape -@pytest.mark.distributed def test_2d_mesh_replicate_flat_round_trip(distributed_setup): """A 2D mesh can replicate on one axis and flat-shard on the other.""" if distributed_setup.world_size < 4 or distributed_setup.world_size % 2 != 0: @@ -523,7 +502,6 @@ def test_2d_mesh_replicate_flat_round_trip(distributed_setup): _assert_dbuffer_local_tensors_close(replicated_buffer, tensors) -@pytest.mark.distributed def test_2d_mesh_flat_before_replicate_is_rejected(distributed_setup): """Flat axes must be a suffix to keep every local buffer contiguous.""" if distributed_setup.world_size < 4 or distributed_setup.world_size % 2 != 0: @@ -545,7 +523,6 @@ def test_2d_mesh_flat_before_replicate_is_rejected(distributed_setup): ) -@pytest.mark.distributed def test_2d_mesh_shards_across_all_ranks(distributed_setup): """Multiple Flat axes shard local storage by the product of their mesh sizes.""" if distributed_setup.world_size < 4 or distributed_setup.world_size % 2 != 0: @@ -574,7 +551,6 @@ def test_2d_mesh_shards_across_all_ranks(distributed_setup): assert fully_sharded_buffer.get_local_tensor(index).is_contiguous() -@pytest.mark.distributed def test_2d_mesh_partial_flat_reduce_scatter_to_flat_flat(distributed_setup): """Partial+Flat reduce-scatter reduces the existing Flat local shard.""" if distributed_setup.world_size < 4 or distributed_setup.world_size % 2 != 0: @@ -618,7 +594,6 @@ def test_2d_mesh_partial_flat_reduce_scatter_to_flat_flat(distributed_setup): _assert_dbuffer_local_tensors_close(replicated_buffer, expected) -@pytest.mark.distributed def test_2d_mesh_replicate_flat_scatter_to_flat_flat(distributed_setup): """Replicate+Flat scatter chunks the existing Flat local shard.""" if distributed_setup.world_size < 4 or distributed_setup.world_size % 2 != 0: From f9594792c90dd13b670e35ac3d9d3246bc6f7e49 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sun, 21 Jun 2026 22:48:35 +0000 Subject: [PATCH 20/23] Remove experimental FSDP CUDA skip guards Signed-off-by: Jingyue Wu --- .../distributed/megatron_fsdp/test_experimental_fully_shard.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py index 86ccea8d942..65109fd15ce 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py +++ b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py @@ -284,9 +284,6 @@ def test_fully_shard_reduces_peak_training_memory(distributed_setup): device = distributed_setup.device if world_size < 2: pytest.skip("This test requires at least 2 ranks.") - if device.type != "cuda": - pytest.skip("Peak memory verification requires CUDA.") - mesh = init_device_mesh(device.type, (world_size,)) dim = 1024 layers = 16 From adb4e128921dadb21965e49742bdb19807bc2e45 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sun, 21 Jun 2026 22:51:41 +0000 Subject: [PATCH 21/23] Document experimental fully_shard mixin attachment Signed-off-by: Jingyue Wu --- .../fsdp/src/megatron_fsdp/experimental/fully_shard.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py index 19516ea66ff..bd0c1d55aa8 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py @@ -57,5 +57,8 @@ def _attach_mixin(module: nn.Module) -> None: if isinstance(module, FsdpModule): return module_cls = module.__class__ + # Attach the FSDP mixin to the original module instance instead of wrapping it + # in a new child module, so parent modules do not need to replace or reorder + # their existing submodule references. fsdp_cls = type(f"ExperimentalFsdp{module_cls.__name__}", (FsdpModule, module_cls), {}) module.__class__ = fsdp_cls From 06286bc7ce48b7316b38672665ab800c9c576bf2 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sun, 21 Jun 2026 22:54:38 +0000 Subject: [PATCH 22/23] Document fully_shard mixin behavior Signed-off-by: Jingyue Wu --- .../fsdp/src/megatron_fsdp/experimental/fully_shard.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py index bd0c1d55aa8..136b600b84c 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py @@ -30,6 +30,9 @@ def fully_shard( ) -> None: """Shard one module as a per-module FSDP unit. + This attaches the FSDP mixin to the original module instance, so parent + modules do not need to replace existing child-module references. + Args: module: Module whose currently unowned parameters become this FSDP unit. mesh: Device mesh used for sharding. @@ -57,8 +60,5 @@ def _attach_mixin(module: nn.Module) -> None: if isinstance(module, FsdpModule): return module_cls = module.__class__ - # Attach the FSDP mixin to the original module instance instead of wrapping it - # in a new child module, so parent modules do not need to replace or reorder - # their existing submodule references. fsdp_cls = type(f"ExperimentalFsdp{module_cls.__name__}", (FsdpModule, module_cls), {}) module.__class__ = fsdp_cls From 0dfb83bcb80830d650d89890033947e6aea6891e Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Mon, 22 Jun 2026 17:07:17 +0000 Subject: [PATCH 23/23] Cover FSDP loss parity with microbatches Signed-off-by: Jingyue Wu --- .../test_experimental_fully_shard.py | 58 ++++++++++++------- 1 file changed, 36 insertions(+), 22 deletions(-) diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py index 65109fd15ce..b9735ccd8c9 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py +++ b/tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py @@ -86,7 +86,8 @@ def _mb(num_bytes: int) -> str: return f"{num_bytes / 1024**2:.2f} MB" -def test_fully_shard_losses_match_baseline(distributed_setup): +@pytest.mark.parametrize("num_microbatches", [1, 3]) +def test_fully_shard_losses_match_baseline(distributed_setup, num_microbatches): """Minimal per-module FSDP training should match single-rank SGD.""" rank = distributed_setup.rank world_size = distributed_setup.world_size @@ -105,28 +106,41 @@ def test_fully_shard_losses_match_baseline(distributed_setup): baseline_optimizer = torch.optim.SGD(baseline.parameters(), lr=0.05) optimizer = torch.optim.SGD(model.parameters(), lr=0.05) - x = torch.randn(3, 8, device=device) - target = torch.randn(3, 4, device=device) - - for step in range(5): - baseline_optimizer.zero_grad() - optimizer.zero_grad() - - baseline_loss = torch.nn.functional.mse_loss(baseline(x), target) - loss = torch.nn.functional.mse_loss(model(x), target) - logger.info( - "FSDP train parity: rank=%s, step=%s, baseline_loss=%s, sharded_loss=%s", - rank, - step, - baseline_loss.item(), - loss.item(), - ) - torch.testing.assert_close(loss, baseline_loss, msg=f"Loss mismatch at step {step}.") + micro_batch_size = 2 + x = torch.randn(num_microbatches, micro_batch_size, 8, device=device) + target = torch.randn(num_microbatches, micro_batch_size, 4, device=device) + microbatches = tuple(zip(x.unbind(), target.unbind())) + + def train(model, optimizer, log_prefix) -> list[torch.Tensor]: + losses = [] + for step in range(5): + optimizer.zero_grad() + + for microbatch, (microbatch_x, microbatch_target) in enumerate(microbatches): + loss = torch.nn.functional.mse_loss(model(microbatch_x), microbatch_target) + losses.append(loss.detach()) + logger.debug( + "%s train parity: rank=%s, step=%s, microbatch=%s, loss=%s", + log_prefix, + rank, + step, + microbatch, + loss, + ) + + (loss / num_microbatches).backward() - baseline_loss.backward() - loss.backward() - baseline_optimizer.step() - optimizer.step() + optimizer.step() + return losses + + baseline_losses = train(baseline, baseline_optimizer, "Baseline") + sharded_losses = train(model, optimizer, "FSDP") + + torch.testing.assert_close( + torch.stack(sharded_losses), + torch.stack(baseline_losses), + msg="Sharded losses did not match baseline losses.", + ) def test_nested_fully_shard_excludes_child_owned_parameters(distributed_setup):