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

Module mixin for the minimal Megatron-FSDP path.

Module Contents#

Classes#

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_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.FsdpModule( context: core.distributed.fsdp.src.megatron_fsdp.experimental.context.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(…)’

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

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

_trainable_parameter_countdown: core.distributed.fsdp.src.megatron_fsdp.experimental.countdown.Countdown#: 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.context.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[torch.nn.Module], 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_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, ...], group_dtype: torch.dtype, ) → tuple[torch.distributed.tensor.placement_types.Placement, ...]#

Specialize public placements for one homogeneous parameter group.

Map Torch’s user-facing Shard(0) to BlockAtomic for MXFP8 groups and RowAtomic for ordinary floating-point groups.

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

Module Contents#

Classes#

Functions#

API#

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