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

`_is_in_backward`	Return whether the current thread is executing an autograd GraphTask.
`_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`
`_specialize_placements`	Specialize public placements for one homogeneous parameter group.

API#

core.distributed.fsdp.src.megatron_fsdp.experimental.module._is_in_backward() → bool#: Return whether the current thread is executing an autograd GraphTask.

class core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpContext( device: torch.device, use_symmetric_memory: bool = False, unify_communication_stream: bool = False, )#

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.
use_symmetric_memory – Whether modules constructed in this context allocate communication staging buffers from PyTorch’s NCCL symmetric-memory pool.
unify_communication_stream – Whether all-gathers and reduce-scatters share one communication stream to reduce peak transient memory.

allgather_stream: torch.cuda.Stream#: None

reduce_scatter_stream: torch.cuda.Stream#: None

is_last_microbatch: bool#: None

use_symmetric_memory: bool#: None

unify_communication_stream: 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

_post_backward_hook_registered: bool#: 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.

post_backward() → None#: Order current-stream consumers after this context’s gradient reductions.

register_post_backward_hook() → None#

Multiple FSDP roots can share this context. Waiting for the reduce-scatter stream in each root’s post_backward() would prevent one root’s backward compute from overlapping another root’s gradient reductions. Wait once at context-level autograd completion instead.

class core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpModule( context: core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpContext, mesh: torch.distributed.DeviceMesh, model_weight_placements: tuple[torch.distributed.tensor.placement_types.Placement, ...], main_grad_placements: tuple[torch.distributed.tensor.placement_types.Placement, ...], main_weight_placements: tuple[torch.distributed.tensor.placement_types.Placement, ...], mixed_precision_policy: core.distributed.fsdp.src.megatron_fsdp.mixed_precision.MixedPrecisionPolicy, grad_divisor: int = 1, schedule_policy: core.distributed.fsdp.src.megatron_fsdp.experimental.schedule.SchedulePolicy = SchedulePolicy(), use_symmetric_memory: bool = False, register_hooks: bool = True, )#

Mixin attached to modules managed by the minimal FSDP path.

Initialization

Initialize FSDP runtime state on an already-constructed module.

class Phase(*args, **kwds)#

Bases: enum.Enum

Lifecycle phase of this FsdpModule.

Initialization

RESTING#: ‘auto(…)’

FORWARD#: ‘auto(…)’

BACKWARD#: ‘auto(…)’

_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

_trainable_parameter_countdown: core.distributed.fsdp.src.megatron_fsdp.experimental.countdown.Countdown#: None

_is_root: bool#: None

_num_trainable_parameters: int#: None

_schedule_policy: core.distributed.fsdp.src.megatron_fsdp.experimental.schedule.SchedulePolicy#: None

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

_phase: core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpModule.Phase#: None

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

property phase: Phase#: Return this module’s lifecycle phase.

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#

register_post_backward_hook( post_backward_hook: collections.abc.Callable[[core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpModule], None], ) → None#

The hook runs when this module’s backward is complete, so it can reshard this module’s parameters and reduce their gradients. It is invoked once all of this module’s trainable parameters have accumulated gradients, or via a full-backward hook when the module owns no trainable parameters.

Parameters:: post_backward_hook – Callback receiving this FSDP module after all of its trainable parameters have accumulated gradients.

static _pre_load_state_dict( _module: torch.nn.Module, _state_dict: dict[str, object], _prefix: str, local_metadata: dict[str, object], _strict: bool, _missing_keys: list[str], _unexpected_keys: list[str], _error_msgs: list[str], ) → None#: Reject state-dict loads that replace parameters managed by FSDP.

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( prefetch: Literal[forward, backward, none] = 'none', ) → None#

Unshard this FsdpModule’s parameter groups immediately.

External schedulers invoking this directly (rather than through the automatic pre_forward hook) must first synchronize the all-gather stream on the root by calling context.allgather_stream.wait_stream(context.current_stream()) before this when self.is_root(); the automatic forward path performs that root sync in pre_forward() immediately before this.

_prefetch_parameter_groups( order: core.distributed.fsdp.src.megatron_fsdp.experimental.indexed_order.IndexedOrder[core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpModule], prefetch_size: int | None, ) → None#: Prefetch successors from order according to this module’s budget.

_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() → None#: Reshard this FsdpModule’s parameter groups.

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

property num_parameter_elements: int#: Return the number of unsharded parameter elements owned by this module.

_nvtx_label(operation: str) → str#

_nvtx_range(operation: str) → collections.abc.Iterator[None]#: Scope an nvtx range to this context so an early return still pops it.

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._specialize_placements( placements: tuple[torch.distributed.tensor.placement_types.Placement, ...], dtype: torch.dtype, ) → tuple[torch.distributed.tensor.placement_types.Placement, ...]#

Specialize public placements for one homogeneous parameter group.

Today every parameter group maps Torch’s user-facing Shard(0) to the DBuffer-specific Flat format. This dtype-homogeneous group boundary is where MXFP8 groups will instead select BlockAtomic.

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

Module Contents#

Classes#

Functions#

API#

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