`core.distributed.fsdp.src.megatron_fsdp.experimental.module`#

Module mixin for the minimal Megatron-FSDP path.

Module Contents#

Classes#

`FsdpContext`	Runtime stream and prefetch state shared by one FSDP subtree.
`FsdpModule`	Mixin attached to modules managed by the minimal FSDP path.

Functions#

`_collect_backward_order`	Collect FsdpModules in static backward prefetch order.
`_axis_index`
`_collect_owned_parameters`
`_group_parameters`

API#

class core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpContext( device: torch.device, root_module: core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpModule, )#

Runtime stream and prefetch state shared by one FSDP subtree.

Initialization

Create rank-local runtime state for a root FSDP subtree.

Parameters:

device – Device on which this context schedules communication.
root_module – Outermost module that owns this context.

allgather_stream: torch.cuda.Stream#: None

reduce_scatter_stream: torch.cuda.Stream#: None

is_last_microbatch: bool#: None

root_module: FsdpModule#: None

forward_order: core.distributed.fsdp.src.megatron_fsdp.experimental.indexed_order.IndexedOrder[FsdpModule]#: None

backward_order: core.distributed.fsdp.src.megatron_fsdp.experimental.indexed_order.IndexedOrder[FsdpModule]#: None

current_stream() → torch.cuda.Stream#: Current stream on this context’s device.

register_post_backward_final_callback() → None#

Root post_backward() means only that root-owned parameters have accumulated gradients; it may run before descendant reductions, or not run at all when the root owns no trainable parameters. Waiting at autograd completion orders consumers after every descendant reduction.

class core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpModule( 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, use_symm_mem: bool = False, )#

Mixin attached to modules managed by the minimal FSDP path.

Initialization

Initialize FSDP runtime state on an already-constructed module.

_name: str | None#: None

_parameter_groups: tuple[core.distributed.fsdp.src.megatron_fsdp.experimental.parameter_group.FsdpParameterGroup, ...]#: None

_context: core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpContext | None#: None

_ready_grad_parameters: set[torch.nn.Parameter]#: None

_num_trainable_parameters: int#: None

_unshard_event: torch.cuda.Event | None#: None

_lazy_init_context() → None#

Initialize one shared runtime context for this FSDP root subtree.

MFSDP v2 requires users to apply fully_shard bottom-up, so child FSDP modules are constructed before their eventual root module is constructed. This method resolves the root lazily on the first forward through the outermost FSDP module and shares that one context with every FSDP descendant.

Alternatives considered:

Eagerly initialize contexts during fully_shard. When a parent is sharded, we could create a new root context and reassign it to all descendant FSDP modules. This creates transient child contexts that are never used if the parent is later sharded, and each parent shard must walk its descendants again, making nested sharding quadratic.
Store an is_root field on each FSDP module. fully_shard could mark newly sharded modules as roots and clear that flag on descendant FSDP modules when a parent is sharded. This avoids creating unused contexts but moves root tracking onto every FSDP module, adding per-module state that must stay consistent with the final sharded module hierarchy.

property context: core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpContext#: Return the initialized runtime context.

property name: str#: Return this FsdpModule’s name.

is_root() → bool#: Return whether this module is the outermost FsdpModule in its context.

_register_hooks() → None#

_make_grad_hook( parameter: torch.nn.Parameter, ) → collections.abc.Callable[[torch.nn.Parameter], None]#

pre_forward() → None#

Prepare full parameters for forward compute and prefetch the next FsdpModule.

While this FsdpModule computes, we issue the next FsdpModule’s all-gather on the comm stream, so AG_{i+1} is launched before F_i finishes.

_unshard_parameter_groups() → None#

Unshard this FsdpModule’s parameter groups on the all-gather stream.

If _unshard_event is already set, this FsdpModule was already unsharded or prefetched and this method is a no-op. Otherwise, this method records _unshard_event after materialization so compute can wait without depending on later release work.

post_forward() → None#: Return parameters to their sharded resting state after forward compute.

_reshard_parameter_groups() → None#

Reshard parameter groups and release unsharded storage after compute.

This method clears _unshard_event after queuing the release, so future users enqueue a fresh all-gather.

pre_backward() → None#: Prepare full parameters and prefetch the next FsdpModule in backward order.

post_backward() → None#: Reduce gradients and return parameters to their sharded resting state.

_reduce_gradient_groups() → None#: Pack gradients and immediately launch their reduce-scatters.

property parameter_groups: tuple[core.distributed.fsdp.src.megatron_fsdp.experimental.parameter_group.FsdpParameterGroup, ...]#: Parameter groups owned by this FsdpModule.

_nvtx_label(phase: Literal[forward, backward]) → str#

core.distributed.fsdp.src.megatron_fsdp.experimental.module._collect_backward_order( module: torch.nn.Module, order: core.distributed.fsdp.src.megatron_fsdp.experimental.indexed_order.IndexedOrder[core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpModule], ) → None#: Collect FsdpModules in static backward prefetch order.

core.distributed.fsdp.src.megatron_fsdp.experimental.module._axis_index( mesh: torch.distributed.DeviceMesh, axis: core.distributed.fsdp.src.megatron_fsdp.experimental.placement.MeshAxis, ) → int#

core.distributed.fsdp.src.megatron_fsdp.experimental.module._collect_owned_parameters( root_module: torch.nn.Module, ) → dict[str, torch.nn.Parameter]#

core.distributed.fsdp.src.megatron_fsdp.experimental.module._group_parameters( parameters: dict[str, torch.nn.Parameter], ) → list[dict[str, torch.nn.Parameter]]#

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

Module Contents#

Classes#

Functions#

API#

`core.distributed.fsdp.src.megatron_fsdp.experimental.module`#