`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 FSDP roots constructed together.
`FsdpModule`	Mixin attached to modules managed by the minimal FSDP path.

Functions#

`_collect_backward_order`	Collect one root’s static backward prefetch order.
`_collect_fsdp_children`	Collect the nearest FSDP descendants of `module`.
`_collect_owned_parameters`
`_group_parameters`

API#

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

Runtime stream and prefetch state shared by FSDP roots constructed together.

Initialization

Create rank-local runtime state for FSDP modules on device.

Parameters:: device – Device on which this context schedules communication.

allgather_stream: torch.cuda.Stream#: None

reduce_scatter_stream: torch.cuda.Stream#: None

is_last_microbatch: bool#: 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

register_module( module: core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpModule, ) → None#: Register a module constructed in this context.

finalize() → None#: Finalize roots, names, and cross-root prefetch orders.

ensure_finalized() → None#: Raise if construction has not completed for this context.

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( context: core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpContext, 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

_num_ready_grad_parameters: int#: None

_is_root: bool#: None

_num_trainable_parameters: int#: None

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

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

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

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

_register_hooks() → None#

_make_grad_hook() → 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 one root’s static backward prefetch order.

core.distributed.fsdp.src.megatron_fsdp.experimental.module._collect_fsdp_children( module: torch.nn.Module, children: set[core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpModule], ) → None#: Collect the nearest FSDP descendants of module.

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`#