Is this page helpful?

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

Minimal Megatron-FSDP fully_shard entrypoint.

Module Contents#

Functions#

fully_shard_context

Construct FSDP modules that share runtime streams and prefetch orders.

fully_shard

Apply FSDP to a module in place.

_normalize_placements

Return a copy with data-parallel mesh axes normalized to integer indices.

_axis_index

microbatch

Mark an FSDP microbatch as the last accumulation microbatch.

_attach_mixin

Data#

_FSDP_CONTEXT

API#

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

‘(…)’

core.distributed.fsdp.src.megatron_fsdp.experimental.fully_shard.fully_shard_context(
device: torch.device | None = None,
*,
use_symmetric_memory: 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.

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.placement.Placements,
mixed_precision_policy: core.distributed.fsdp.src.megatron_fsdp.mixed_precision.MixedPrecisionPolicy | None = None,
grad_divisor: int = 1,
) → 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 – 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.

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

core.distributed.fsdp.src.megatron_fsdp.experimental.fully_shard._normalize_placements(
mesh: torch.distributed.DeviceMesh,
placements: core.distributed.fsdp.src.megatron_fsdp.experimental.placement.Placements,
) → core.distributed.fsdp.src.megatron_fsdp.experimental.placement.Placements#

Return a copy with data-parallel mesh axes normalized to integer indices.

core.distributed.fsdp.src.megatron_fsdp.experimental.fully_shard._axis_index(
mesh: torch.distributed.DeviceMesh,
axis: core.distributed.fsdp.src.megatron_fsdp.experimental.placement.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-Flat 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#