`core.distributed.fsdp.src.megatron_fsdp.experimental.fully_shard`#

Minimal Megatron-FSDP fully_shard entrypoint.

Module Contents#

Classes#

Placements

Per-data-parallel-axis placements for MFSDP buffers.

Functions#

`fully_shard`	Apply FSDP to a module in place.
`_validate_dp_axes`	Validate the parent mesh’s data-parallel axes.
`_axis_index`
`microbatch`	Mark an FSDP microbatch as the last accumulation microbatch.
`_attach_mixin`
`_get_fsdp_class`	Reuse the subclass so classmethods share lazy state, such as CUDA streams.

Data#

`fully_shard_context`
`MeshAxis`

API#

core.distributed.fsdp.src.megatron_fsdp.experimental.fully_shard.fully_shard_context#: None

core.distributed.fsdp.src.megatron_fsdp.experimental.fully_shard.MeshAxis#: None

class core.distributed.fsdp.src.megatron_fsdp.experimental.fully_shard.Placements#

Per-data-parallel-axis placements for MFSDP buffers.

dp_axes identifies the parent-mesh axes that form MFSDP’s data-parallel mesh. Placement sequences are ordered to match those axes. Use Torch’s Shard(0) for public parameter, gradient, and optimizer sharding.

dp_axes: collections.abc.Sequence[core.distributed.fsdp.src.megatron_fsdp.experimental.fully_shard.MeshAxis]#: None

parameter: collections.abc.Sequence[torch.distributed.tensor.placement_types.Placement]#: None

gradient: collections.abc.Sequence[torch.distributed.tensor.placement_types.Placement]#: None

optimizer: collections.abc.Sequence[torch.distributed.tensor.placement_types.Placement]#: None

__post_init__() → None#: Validate placement sequence lengths.

core.distributed.fsdp.src.megatron_fsdp.experimental.fully_shard.fully_shard( module: torch.nn.Module, mesh: torch.distributed.DeviceMesh, placements: core.distributed.fsdp.src.megatron_fsdp.experimental.fully_shard.Placements, mixed_precision_policy: core.distributed.fsdp.src.megatron_fsdp.mixed_precision.MixedPrecisionPolicy | None = None, grad_divisor: int = 1, schedule_policy: core.distributed.fsdp.src.megatron_fsdp.experimental.schedule.SchedulePolicy = SchedulePolicy(), register_hooks: bool = True, ) → None#

Apply FSDP to a module in place.

This attaches the FSDP mixin to the original module instance, so parent modules do not need to replace existing child-module references.

Parameters:

module – Module whose currently unowned parameters are managed by FSDP.
mesh – Parent device mesh containing the data-parallel axes.
placements – Parameter, gradient, and optimizer placements.
mixed_precision_policy – Optional precision policy. Defaults to FP32 main weights and parameter-dtype main gradients.
grad_divisor –
Additional divisor applied to the reduced gradient, on top of the averaging the mesh already performs. Defaults to 1, which is correct whenever each mesh rank contributes exactly one term to the gradient.

Expert parallelism is the motivating case. A rank’s experts process tokens routed to them from every rank in the expert-parallel group, and the backward pass routes those tokens’ gradients back, so a rank’s expert gradient already sums over ep_size ranks’ data before any reduction happens. Averaging over the expert-data-parallel mesh alone therefore divides by too little, and grad_divisor=ep_size makes up the difference. Dense parameters see only their own rank’s tokens and need no divisor.
schedule_policy – Communication scheduling policy for this FSDP module.
register_hooks – Whether to register the automatic forward and backward execution hooks on module. Disable this when an external scheduler invokes the corresponding FSDP lifecycle methods explicitly. The state-dict safety hook is registered independently.

core.distributed.fsdp.src.megatron_fsdp.experimental.fully_shard._validate_dp_axes( mesh: torch.distributed.DeviceMesh, dp_axes: collections.abc.Sequence[core.distributed.fsdp.src.megatron_fsdp.experimental.fully_shard.MeshAxis], ) → None#: Validate the parent mesh’s data-parallel axes.

core.distributed.fsdp.src.megatron_fsdp.experimental.fully_shard._axis_index( mesh: torch.distributed.DeviceMesh, axis: core.distributed.fsdp.src.megatron_fsdp.experimental.fully_shard.MeshAxis, ) → int#

core.distributed.fsdp.src.megatron_fsdp.experimental.fully_shard.microbatch( context: core.distributed.fsdp.src.megatron_fsdp.experimental.context.FsdpContext, is_last: bool, ) → collections.abc.Iterator[None]#

Mark an FSDP microbatch as the last accumulation microbatch.

At present, this is only needed for HSDP/HFSDP gradient accumulation, so FSDP finalizes gradients only on the last backward. Plain all-Shard(0) data parallelism finalizes gradients on every backward and does not need it.