Is this page helpful?

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

Construct FSDP modules that share runtime streams and prefetch orders.

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

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

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_context(
device: torch.device | None = None,
*,
use_symmetric_memory: bool = False,
unify_communication_stream: bool = False,
) → collections.abc.Iterator[core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpContext]#

Construct FSDP modules that share runtime streams and prefetch orders.

Independent roots are ordered by their root-level fully_shard calls. Construction must finish before any of the registered modules run forward.

Parameters:

  • device – CUDA device on which to create communication streams. Defaults to the current CUDA device.

  • use_symmetric_memory – Allocate communication staging buffers from PyTorch’s NCCL symmetric-memory pool.

  • unify_communication_stream – Whether all-gathers and reduce-scatters share one communication stream to reduce peak transient memory. See https://github.com/NVIDIA/Megatron-LM/issues/6471.

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.module.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.

Parameters:

  • context – FSDP context whose roots should use this microbatch state.

  • is_last – Whether forwards in this scope are for the last microbatch.

core.distributed.fsdp.src.megatron_fsdp.experimental.fully_shard._attach_mixin(module: torch.nn.Module) → None#