`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_context`	Activate `context` for FSDP construction, finalizing it on successful exit.
`current_fully_shard_context`	Return the innermost active `fully_shard_context`, or `None`.
`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#

`_FSDP_CONTEXT`
`MeshAxis`

API#

core.distributed.fsdp.src.megatron_fsdp.experimental.fully_shard._FSDP_CONTEXT#: ‘(…)’

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

core.distributed.fsdp.src.megatron_fsdp.experimental.fully_shard.fully_shard_context( context: core.distributed.fsdp.src.megatron_fsdp.experimental.context.FsdpContext, ) → collections.abc.Iterator[core.distributed.fsdp.src.megatron_fsdp.experimental.context.FsdpContext]#: Activate context for FSDP construction, finalizing it on successful exit.

core.distributed.fsdp.src.megatron_fsdp.experimental.fully_shard.current_fully_shard_context() → core.distributed.fsdp.src.megatron_fsdp.experimental.context.FsdpContext | None#

Return the innermost active fully_shard_context, or None.

Read-only counterpart of :func:fully_shard_context: it never creates, joins, or finalizes a context, and returns None whenever no fully_shard_context scope is active. Callers that must share one context – for example per-chunk wrappers built by a single wrap call – use it to join the caller’s ambient context instead of opening a second one.

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_scale: float = 1.0, schedule_policy: core.distributed.fsdp.src.megatron_fsdp.experimental.schedule_policy.SchedulePolicy = SchedulePolicy(), shard_layout: core.distributed.fsdp.src.megatron_fsdp.experimental.placement.ShardLayout = ShardLayout.PADDED, ) → 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_scale –
Multiplier applied after mesh averaging. Defaults to 1, which is correct when 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_scale=1 / ep_size makes up the difference. Dense parameters see only their own rank’s tokens and need no additional scaling.

For per-token loss, grad_scale=mesh.size() cancels mesh averaging so gradients are summed. Normalize once by the global token count after accumulating all microbatches and before clipping or updating weights.
schedule_policy – Prefetching, execution-hook, and gradient-completion policy for this FSDP module.
shard_layout – How RowAtomic/BlockAtomic parameter groups are split across ranks. ShardLayout.PADDED (default) pads every group to equal shards; ShardLayout.UNPADDED packs each group without padding and splits it at row boundaries (32-row block boundaries for MXFP8) into generally unequal shards.

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.

Non-final microbatches still release unsharded weights, but defer automatic gradient synchronization. HSDP/HFSDP also defer DP-outer gradient reduction until the last microbatch. Wrap both forward and backward in this scope.