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

`_collect_owned_parameters`
`_group_parameters`
`_specialize_placements`	Specialize public placements for one homogeneous parameter group.

API#

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_scale: float = 1.0, schedule_policy: core.distributed.fsdp.src.megatron_fsdp.experimental.schedule_policy.SchedulePolicy = SchedulePolicy(), use_symmetric_memory: bool = False, shard_layout: core.distributed.fsdp.src.megatron_fsdp.experimental.placement.ShardLayout = ShardLayout.PADDED, )#

Mixin attached to modules managed by the minimal FSDP path.

Initialization

Initialize FSDP runtime state on an already-constructed module.

Parameters:: shard_layout – Padded or unpadded rank segments for every parameter group. See fully_shard.

_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_policy.SchedulePolicy#: None

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

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

property phase: core.distributed.fsdp.src.megatron_fsdp.experimental.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 expected gradient-accumulation callbacks have arrived, 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.

_trainable_fsdp_parameters() → collections.abc.Iterator[core.distributed.fsdp.src.megatron_fsdp.experimental.parameter_group.FsdpParameter]#: Iterate this module’s trainable parameters in a stable order.

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( phase: core.distributed.fsdp.src.megatron_fsdp.experimental.phase.Phase | None = None, prefetch: bool = False, ) → None#

Unshard this FsdpModule’s parameter groups immediately.

Parameters:

phase – In what phase this module runs. None unshards for every phase.
prefetch – Whether to also prefetch successors in the module order of phase, which must then be FORWARD or BACKWARD.

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.execution_order.ExecutionOrder, prefetch_size: int | None, ) → None#: Validate the current event, prefetch successors, and advance once.

_unshard_parameter_groups( phase: core.distributed.fsdp.src.megatron_fsdp.experimental.phase.Phase | None, ) → None#: Ask each group to gather missing weights on the all-gather stream.

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.

_release_unsharded_storage(*, rowwise_only: bool = False) → None#: Select and release gathered storage after the current compute finishes.

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 launch their reductions.

When overlap is enabled, the last-microbatch DP-outer reduction runs on its own stream so a group’s DP-outer reduction overlaps the next group’s DP-inner reduction; otherwise both share the reduce-scatter stream and stay serialized.

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