`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, ) → 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.

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, use_symm_mem: bool = False, ) → 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.
use_symm_mem – Allocate all-gather and reduce-scatter staging buffers from PyTorch’s NCCL symmetric-memory pool.

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.