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

Parameter-group runtime state for the minimal Megatron-FSDP path.

Module Contents#

Classes#

`FsdpParameter`	One physical parameter and its FSDP runtime representations.
`FsdpParameterGroup`	A dtype and requires-grad homogeneous group of FSDP-owned parameters.

Functions#

`get_containing_parameter_group`	Return the FSDP parameter group that owns `parameter`, if any.
`_get_parameter_owner`	Resolve a root-module-relative parameter FQN to its direct owner.

Data#

_CONTAINING_PARAMETER_GROUP_ATTR

API#

core.distributed.fsdp.src.megatron_fsdp.experimental.parameter_group._CONTAINING_PARAMETER_GROUP_ATTR#: ‘_mfsdp_parameter_group’

core.distributed.fsdp.src.megatron_fsdp.experimental.parameter_group.get_containing_parameter_group( parameter: torch.nn.Parameter, ) → FsdpParameterGroup | None#: Return the FSDP parameter group that owns parameter, if any.

class core.distributed.fsdp.src.megatron_fsdp.experimental.parameter_group.FsdpParameter#

One physical parameter and its FSDP runtime representations.

fqns: tuple[str, ...]#: None

sharded: torch.nn.Parameter#: None

unsharded: torch.nn.Parameter#: None

class core.distributed.fsdp.src.megatron_fsdp.experimental.parameter_group.FsdpParameterGroup( owning_module: torch.nn.Module, parameters: dict[str, torch.nn.Parameter], 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, )#

A dtype and requires-grad homogeneous group of FSDP-owned parameters.

Initialization

Create persistent sharded buffers for a group of parameters.

Parameters:

owning_module – Closest FSDP root module that owns this parameter group.
parameters – Root-module-relative FQNs and their parameters.
mesh – Device mesh used for all DBuffer storage in this version.
placements – Parameter, gradient, and optimizer placements.
mixed_precision_policy – Precision policy for main weights and gradients.
use_symm_mem – Allocate communication staging buffers from PyTorch’s NCCL symmetric-memory pool.

_owning_module: weakref.ReferenceType[torch.nn.Module]#: None

fsdp_parameters: tuple[core.distributed.fsdp.src.megatron_fsdp.experimental.parameter_group.FsdpParameter, ...]#: None

mesh: torch.distributed.DeviceMesh#: None

dtype: torch.dtype#: None

requires_grad: bool#: None

main_weight: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer#: None

model_weight: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer#: None

main_grad: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer | None#: None

_unsharded_model_weight: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer#: None

_symm_mem_pool: torch.cuda.MemPool | None#: None

_symmetric_memory_context()#

_set_module_parameter( fqns: tuple[str, ...], parameter: torch.nn.Parameter, ) → None#

_switch_to_sharded_parameters() → None#

_switch_to_unsharded_parameters() → None#

sync_model_weight_from_main_weight() → None#: Refresh compute weights from optimizer weights.

unshard_parameters() → None#: Install full parameters for local compute.

reshard_parameters() → None#: Install sharded DTensor parameters on the owning modules.

release_unsharded_storage() → None#: Release this group’s full-parameter storage.

allocate_partial_grad_buffer() → core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer#: Allocate the unreduced reduce-scatter input buffer.

copy_gradients_to_partial_buffer( partial_grad: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer, ) → None#: Pack full local gradients into an existing reduce-scatter input buffer.

_has_sharded_grads() → bool#

reduce_partial_gradients( partial_grad: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer, is_last_microbatch: bool = True, ) → None#

Reduce a packed partial gradient buffer into sharded parameter gradients.

For HSDP main_grad rests DP-outer-Partial (Partial where main_weight is Replicate) between microbatches, accumulating each backward through the standard zero_grad contract; the last microbatch reduces the DP-outer axes, finalizing main_grad to main_weight’s placements so .grad is the fully reduced gradient before optimizer.step(). With every axis Flat (plain DP) main_grad already rests finalized.

core.distributed.fsdp.src.megatron_fsdp.experimental.parameter_group._get_parameter_owner( module: torch.nn.Module, name: str, ) → tuple[torch.nn.Module, str]#: Resolve a root-module-relative parameter FQN to its direct owner.

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