Is this page helpful?

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

Minimal Megatron-FSDP fully_shard entrypoint.

Module Contents#

Functions#

fully_shard

Apply FSDP to a module in place.

microbatch

Mark an FSDP microbatch as the last accumulation microbatch.

_attach_mixin

_collect_fsdp_contexts

API#

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.microbatch(
module: torch.nn.Module,
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:

  • module – Module tree whose FSDP 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#
core.distributed.fsdp.src.megatron_fsdp.experimental.fully_shard._collect_fsdp_contexts(
module: torch.nn.Module,
contexts: list[core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpContext],
) → None#