core.distributed.fsdp.src.megatron_fsdp.experimental.module_utils

Utilities for parameter ownership, metadata, and autograd state in FSDP modules.

Module Contents

Functions

is_in_backward	Return whether the current thread is executing an autograd GraphTask.
get_parameter_owner	Resolve a root-module-relative parameter FQN to its direct owner.
copy_parameter_attributes	Copy model metadata between Parameters, including false or None values.

Data

_PARAMETER_ATTRIBUTES_TO_PRESERVE

API

core.distributed.fsdp.src.megatron_fsdp.experimental.module_utils.is_in_backward() → bool

Return whether the current thread is executing an autograd GraphTask.

core.distributed.fsdp.src.megatron_fsdp.experimental.module_utils.get_parameter_owner(

root_module: torch.nn.Module,

parameter_fqn: str,

) → tuple[torch.nn.Module, str]

Resolve a root-module-relative parameter FQN to its direct owner.

core.distributed.fsdp.src.megatron_fsdp.experimental.module_utils._PARAMETER_ATTRIBUTES_TO_PRESERVE

(‘is_embedding_or_output_parameter’, ‘is_embedding_parameter’, ‘use_muon’, ‘is_wide_residual_retenti…

core.distributed.fsdp.src.megatron_fsdp.experimental.module_utils.copy_parameter_attributes(

from_: torch.nn.Parameter,

to_: torch.nn.Parameter,

) → None

Copy model metadata between Parameters, including false or None values.

Extend the shared attribute list when adding a new model-parameter contract.