Add experimental Megatron-FSDP fully_shard implementation

megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py

@@ -15,6 +15,19 @@
 """Experimental Megatron-FSDP implementation."""
 
 from .dbuffer import DBuffer
-from .placement import Flat, Partial, Placement, Replicate
+from .fsdp_module import FsdpModule
+from .fully_shard import fully_shard
+from .parameter_group import ParameterGroup
+from .placement import Flat, Partial, Placement, Placements, Replicate
 
-__all__ = ["DBuffer", "Flat", "Partial", "Placement", "Replicate"]
+__all__ = [
+    "DBuffer",
+    "Flat",
+    "FsdpModule",
+    "ParameterGroup",
+    "Partial",
+    "Placement",
+    "Placements",
+    "Replicate",
+    "fully_shard",
+]

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)):
@@ -116,6 +117,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]
@@ -254,6 +269,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":

megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fsdp_module.py

@@ -0,0 +1,161 @@
+# 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 = _collect_owned_parameters(self)
+        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.sync_model_weight_from_main_weight()
+            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 _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))
+
+        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]

megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py

@@ -0,0 +1,61 @@
+# 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 Megatron-FSDP fully_shard entrypoint."""
+
+from torch import nn
+from torch.distributed import DeviceMesh
+
+from ..mixed_precision import MixedPrecisionPolicy
+from .fsdp_module import FsdpModule
+from .placement import Placements
+
+
+def fully_shard(
+    module: nn.Module,
+    mesh: DeviceMesh,
+    placements: Placements,
+    mixed_precision_policy: MixedPrecisionPolicy | None = None,
+) -> None:
+    """Shard one module as a 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.
+    """
+    if isinstance(module, FsdpModule):
+        raise ValueError("This module is already managed by FSDP.")
+
+    mixed_precision_policy = mixed_precision_policy or MixedPrecisionPolicy()
+    original_cls = module.__class__
+    _attach_mixin(module)
+    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 _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