`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.
`sync_model_weights_from_main_weights`	Sync MFSDP compute weights for parameter groups represented by `parameters`.

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.

core.distributed.fsdp.src.megatron_fsdp.experimental.parameter_group.sync_model_weights_from_main_weights( parameters: collections.abc.Iterable[torch.nn.Parameter], ) → None#

Sync MFSDP compute weights for parameter groups represented by parameters.

Parameters outside the experimental MFSDP path are ignored. A parameter group may own multiple parameters, but its compute-weight buffer is synced once.

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, fqn_to_parameter: dict[str, torch.nn.Parameter], 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, use_symmetric_memory: 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.
fqn_to_parameter – Root-module-relative FQNs and their parameters.
mesh – Parent device mesh containing the data-parallel axes.
model_weight_placements – Compute-weight buffer placements.
main_grad_placements – Main-gradient buffer placements.
main_weight_placements – Main-weight buffer placements.
mixed_precision_policy – Precision policy for main weights and gradients.
use_symmetric_memory – Allocate communication staging buffers from PyTorch’s NCCL symmetric-memory pool.
grad_divisor – Additional divisor applied on top of the mesh-size averaging. See fully_shard.

_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

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

_model_weight_is_stale: bool#: None

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

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

_main_grad_is_stale: bool#: None

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

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

grad_divisor: int#: None

static _collect_parameter_metadata( fqn_to_parameter: dict[str, torch.nn.Parameter], ) → tuple[dict[torch.nn.Parameter, list[str]], torch.dtype, bool]#: Group tied parameters and validate their shared metadata.

_initialize_buffers( parameters: tuple[torch.nn.Parameter, ...], 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, use_symmetric_memory: bool, ) → None#: Allocate weight and gradient buffers in their required dependency order.

_build_fsdp_parameters( parameter_to_fqns: dict[torch.nn.Parameter, list[str]], ) → tuple[core.distributed.fsdp.src.megatron_fsdp.experimental.parameter_group.FsdpParameter, ...]#: Materialize parameter storage and build its FSDP representations.

_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, ) → None#

Reduce a packed partial gradient buffer into sharded parameter gradients.

For HSDP/HFSDP main_grad rests DP-outer-Partial 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 (all-reduce to Replicate for HSDP, reduce-scatter to Flat for HFSDP) 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#