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

`_contains_any_placement_type`	Return whether any placement in `placement_groups` matches `placement_types`.
`sync_model_weights_from_main_weights`	Sync MFSDP compute weights for parameter groups represented by `parameters`.
`get_containing_parameter_group`	Return the FSDP parameter group that owns `parameter`, if any.

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._contains_any_placement_type( *placement_groups: tuple[torch.distributed.tensor.placement_types.Placement, ...], placement_types: type | tuple[type, ...], ) → bool#: Return whether any placement in placement_groups matches placement_types.

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_scale: float = 1.0, use_symmetric_memory: bool = False, parameter_to_owner: dict[torch.nn.Parameter, int] | None = None, shard_layout: core.distributed.fsdp.src.megatron_fsdp.experimental.placement.ShardLayout = ShardLayout.PADDED, )#

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, in logical tensor order. Tied parameters count once.
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_scale – Multiplier applied after the mesh-size averaging. See fully_shard.
parameter_to_owner – Construction-time mapping from original parameters to owner ranks in mesh. Every TensorAtomic parameter needs an entry; entries for other parameters and non-TensorAtomic groups are ignored.
shard_layout – Whether RowAtomic/BlockAtomic layouts pad to equal rank segments or pack without padding into generally unequal ones. 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 | core.distributed.fsdp.src.megatron_fsdp.experimental.quantized_dbuffer.QuantizedDBuffer#: None

post_optimizer_model_weight: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer | core.distributed.fsdp.src.megatron_fsdp.experimental.quantized_dbuffer.QuantizedDBuffer#: 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

grad_comm_dtype: torch.dtype#: None

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

_unshard_events: dict[core.distributed.fsdp.src.megatron_fsdp.experimental.phase.Phase, torch.cuda.Event]#: None

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

grad_scale: float#: None

shard_layout: core.distributed.fsdp.src.megatron_fsdp.experimental.placement.ShardLayout#: 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.

static _validate_parameter_to_owner( parameter_to_fqns: dict[torch.nn.Parameter, list[str]], parameter_to_owner: dict[torch.nn.Parameter, int], ) → None#: Report by FQN the first parameter that has no owner assignment.

_build_layout( parameters: tuple[torch.nn.Parameter, ...], parameter_to_owner: dict[torch.nn.Parameter, int] | None, *placement_groups: tuple[torch.distributed.tensor.placement_types.Placement, ...], ) → core.distributed.fsdp.src.megatron_fsdp.experimental.layout.GlobalLayout#: Build the global layout shared by this group’s buffers.

_initialize_buffers( parameters: tuple[torch.nn.Parameter, ...], parameter_to_owner: dict[torch.nn.Parameter, int] | None, 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#: Build the shared layout and allocate weight and gradient buffers.

_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#: Expose each FsdpParameter’s unsharded Parameter on its owning module.

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

property is_rowwise_unsharded: bool#: Whether MXFP8 rowwise planes occupy temporary gathered storage.

unshard_parameters( phase: core.distributed.fsdp.src.megatron_fsdp.experimental.phase.Phase | None = None, ) → None#

Prepare this group’s Parameters for the requested compute phase.

On the caller-selected all-gather stream, refresh compute weights, gather missing data if sharded, update FsdpParameters, switch the module’s Parameters to their unsharded representations, and record readiness. Reuse existing events so compute never waits for a later prefetch when its weights are already available.

wait_for_unshard( phase: core.distributed.fsdp.src.megatron_fsdp.experimental.phase.Phase | None, ) → None#: Make the current stream wait for this group’s requested weights.

_install_unsharded_parameters( *, rowwise: bool = True, columnwise: bool = True, ) → None#

Update each FsdpParameter’s unsharded data and scales from the weight buffer.

rowwise and columnwise select the newly gathered MXFP8 planes to install, preserving planes already installed for the other phase. Regular compute weights ignore them. Replicated weights also need installation after optimizer updates because their padded MXFP8 scales are copies.

_allgather_model_weight( *, rowwise: bool = True, columnwise: bool = True, ) → None#

Allocate the gather buffer and fill it from sharded compute weights.

rowwise and columnwise select the MXFP8 planes to gather. Regular compute weights ignore them and gather their whole buffer.

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

release_unsharded_storage(*, rowwise_only: bool = False) → None#

Release gathered weights, or only their MXFP8 rowwise planes.

The caller orders this release after compute on the all-gather stream.

Storage is resized to zero instead of replacing parameter .data: autograd may have saved views into the unsharded parameters during forward, and those views must follow the Storage when a later unshard reallocates it.

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_inner_gradients( partial_grad: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer, ) → None#

Reduce a packed partial gradient across DP-inner and accumulate into main_grad.

main_grad rests at the DP-outer-Partial accumulation placement, accumulating each backward through the standard zero_grad contract, and .grad binds to it. The last-microbatch DP-outer reduction is applied separately by reduce_outer_gradients; the caller (_reduce_gradient_groups) decides when and on which stream to run it, which is what lets the DP-outer reduction overlap the next group’s DP-inner reduction.

reduce_outer_gradients() → None#

Finalize the deferred DP-outer reduction into the optimizer-layout view.

All-reduce (HSDP) or reduce-scatter (HFSDP) the accumulator across DP-outer into pre_optimizer_main_grad (the persistent buffer’s optimizer-layout view), then bind each sharded .grad to it. Runs after reduce_inner_gradients on the last microbatch; the caller can run it on its own stream to overlap the DP-inner reductions (see FsdpModule._reduce_gradient_groups). main_grad and its view are persistent, so running on a separate stream needs no cross-stream free.

_install_sharded_grads( grad_buffer: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer, ) → None#: Bind each sharded parameter’s .grad to the given reduced gradient buffer.

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

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