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

Pure parameter layout and owner-compute packing logic for MFSDP v2’s all-Flat layout.

ParameterLayout describes how a single parameter’s flat element range is split across the DP group under MFSDP v2’s all-Flat layout.
ParameterLayout.from_group builds {tensor_index: layout} for eligible parameters in an FsdpParameterGroup, keyed by each parameter’s index within the group.
assign_owner_work balances owner-compute work across owner ranks using a caller-supplied cost function.
GroupOwnerLayout.from_group builds a data structure capturing the per-group owner layout upon the above.
OwnerGatherPlan.pack/OwnerScatterPlan.pack take the GroupOwnerLayout plus this rank’s data and build the flat P2P send/recv buffers,
OwnerGatherPlan.reconstruct_full stitches gathered shards back into the full flat tensor on the owner, and
OwnerScatterPlan.unpack extracts this rank’s flat result shards from the received buffers.

Module Contents#

Classes#

`ParameterLayout`	How a single parameter’s flat element range splits across the DP group.
`GroupOwnerLayout`	Owner layout for one `FsdpParameterGroup`: its params and their owner ranks.
`OwnerGatherPlan`	Metadata and send buffers for the owner-gather P2P step of a set of parameters.
`OwnerScatterPlan`	Metadata and send buffers for the owner-scatter P2P step of a set of parameters.

Functions#

`select_ge_2d_params`	Whether the given tensor has dimensionality ≥2.
`ns_cost_fn`	Cost function matching the Newton-Schulz orthogonalization compute estimate for the given number of Newton-Schulz iterations/steps.
`assign_owner_work`	Assign one owner rank to each parameter, keyed by tensor index.

API#

core.distributed.fsdp.src.megatron_fsdp.experimental.owner_planning.select_ge_2d_params(param: torch.Tensor) → bool#: Whether the given tensor has dimensionality ≥2.

class core.distributed.fsdp.src.megatron_fsdp.experimental.owner_planning.ParameterLayout#

How a single parameter’s flat element range splits across the DP group.

MFSDP v2’s all-Flat layout gives each rank one contiguous global element range per parameter, in rank order, so rank r holds [offset, offset + count) where offset is the sum of the previous ranks’ counts. A rank with count == 0 holds no elements of this parameter.

.. attribute:: full_shape

The parameter’s global shape.

.. attribute:: flat_counts

Per-rank element count; 0 means the rank holds no elements.

full_shape: torch.Size#: None

flat_counts: tuple[int, ...]#: None

__post_init__() → None#

classmethod from_group( group: core.distributed.fsdp.src.megatron_fsdp.experimental.parameter_group.FsdpParameterGroup, *, eligible_fn: collections.abc.Callable[[torch.Tensor], bool] | None = None, ) → dict[int, Self]#

Build {tensor_index: layout} for eligible parameters in an FsdpParameterGroup.

Keys are the parameters’ indices within group.fsdp_parameters (their tensor indices in the DBuffer layout). Parameters not selected by eligible_fn are absent from the returned dict.

Parameters:

group – The FSDP parameter group whose DBuffer layout describes the parameter placements.
eligible_fn – Predicate selecting which parameters participate in owner-compute orthogonalization. Takes a parameter tensor as input and return whether the parameter is supposed to be included. When None, defaults to matching ≥2D tensors (param.ndim >= 2).

property dp_size: int#: Number of ranks in the DP group for this parameter.

full_numel() → int#: Return the total number of elements in the full (unsharded) parameter.

rank_offset(rank: int) → int#: Return the starting offset of rank’s shard within the flat parameter.

rank_numel(rank: int) → int#: Return the number of elements rank holds.

owner_candidates() → tuple[int, ...]#: Return the ranks that hold a non-empty shard of this parameter.

is_boundary() → bool#: True if more than one rank holds a non-empty shard of this parameter.

core.distributed.fsdp.src.megatron_fsdp.experimental.owner_planning.ns_cost_fn( num_ns_steps: int, ) → collections.abc.Callable[[core.distributed.fsdp.src.megatron_fsdp.experimental.owner_planning.ParameterLayout], int]#

Cost function matching the Newton-Schulz orthogonalization compute estimate for the given number of Newton-Schulz iterations/steps.

numel * (min(rows, cols) * num_steps + 1) under the DBuffer’s leading-dim view ((shape[0], shape[1:].numel())).

core.distributed.fsdp.src.megatron_fsdp.experimental.owner_planning.assign_owner_work( layouts: dict[int, core.distributed.fsdp.src.megatron_fsdp.experimental.owner_planning.ParameterLayout], cost_fn: collections.abc.Callable[[core.distributed.fsdp.src.megatron_fsdp.experimental.owner_planning.ParameterLayout], float] | None = None, ) → dict[int, int]#

Assign one owner rank to each parameter, keyed by tensor index.

Non-boundary parameters are assigned to their original rank (the only rank holding their elements, so no communication is needed) and their cost counts toward that rank’s running total cost. Boundary parameters are processed in descending cost order and each is greedily given to its eligible rank with the smallest running cost total.

Parameters:

layouts – Parameter layouts keyed by each parameter’s tensor index in its FsdpParameterGroup.
cost_fn – Callable that returns a positive cost estimate for a given parameter layout. The greedy balancer minimizes the maximum running cost total across ranks, so the cost should reflect the relative compute weight of owning each parameter (e.g., an orthogonalization cost estimate). When None, defaults to a compute estimate for orthogonalization via Newton-Schulz with 5 iterations/steps.

Returns:

Mapping from tensor index to owner rank.

class core.distributed.fsdp.src.megatron_fsdp.experimental.owner_planning.GroupOwnerLayout#

Owner layout for one FsdpParameterGroup: its params and their owner ranks.

The layouts and assignments are step-independent in the general case, so the owner layout may be cached across optimizer steps.

The tensor_index used here refers to the parameter’s/tensor’s index in the FsdpParameterGroup.

.. attribute:: group

The FSDP parameter group the layouts and owners refer to.

.. attribute:: layouts

{tensor_index: layout} for the participating parameters.

.. attribute:: owners

{tensor_index: owner_rank} with an entry for every participating parameter. Ranks are indices into the group’s mesh.

group: core.distributed.fsdp.src.megatron_fsdp.experimental.parameter_group.FsdpParameterGroup#: None

layouts: dict[int, core.distributed.fsdp.src.megatron_fsdp.experimental.owner_planning.ParameterLayout]#: None

owners: dict[int, int]#: None

classmethod from_group( group: core.distributed.fsdp.src.megatron_fsdp.experimental.parameter_group.FsdpParameterGroup, *, cost_fn: collections.abc.Callable[[core.distributed.fsdp.src.megatron_fsdp.experimental.owner_planning.ParameterLayout], float] | None = None, eligible_fn: collections.abc.Callable[[torch.Tensor], bool] | None = None, ) → Self#

Build the owner layout for one group.

Parameters:

group – The FSDP parameter group whose DBuffer layout describes the parameter placements.
cost_fn – Cost estimate per parameter layout used to balance owner assignments across ranks. When None, defaults to a compute estimate for orthogonalization via Newton-Schulz with 5 iterations/steps. See also assign_owner_work.
eligible_fn – Predicate selecting which parameters participate in owner-compute orthogonalization. When None, defaults to matching ≥2D tensors. See also ParameterLayout.from_group.

property mesh: torch.distributed.device_mesh.DeviceMesh#: Device mesh of the group.

class core.distributed.fsdp.src.megatron_fsdp.experimental.owner_planning.OwnerGatherPlan#

Metadata and send buffers for the owner-gather P2P step of a set of parameters.

The owner keeps its own shard locally (no self-send), so it only receives from the other shard-holding ranks. reconstruct_full reconstructs each owned tensor by concatenating the per-rank shards in rank order (i.e., global element order).

Example:

# We are also using some pseudocode here for brevity.

# Assume:
torch.distributed.get_world_size() == 2
torch.distributed.get_rank() == 0  # We're observing from rank 0
param_0: torch.Tensor
param_1: torch.Tensor
# Params are in this order as observed by MFSDP.
model.param_groups == [{"params": [param_0, param_1]}]
# Both params are owned by rank 1 (was previously determined using `GroupOwnerLayout`).
param_0.owner == 1
param_1.owner == 1

param_0.shape == (6, 4)  # Global shape.
param_1.shape == (4, 4)  # Global shape.
param_0.local_shard.shape == (3, 4)  # Rank 0 has shard indexed by `[0:3, ...]`.
param_1.local_shard.shape == (2, 4)  # Rank 0 has shard indexed by `[0:2, ...]`.
param_0.local_shard.numel == 12
param_1.local_shard.numel == 8

owner_gather_plan.send_buffers == {1: tensor(20)}  # 12 + 8 = 20 elements
# `owner_gather_plan.send_buffers[1]` represents the following in its packed flat buffer:
#   +--------------------+-------------------+
#   | param_0 (12 elems) | param_1 (8 elems) |
#   +--------------------+-------------------+
#                 byte order: -->

# Rank 0 owns nothing
owner_gather_plan.recv_sizes == {}
owner_gather_plan.own_shards == {}
owner_gather_plan.recv_offsets == {}

# ---

# Same settings as above, now observing from rank 1 (the owner):
torch.distributed.get_rank() == 1

param_0.local_shard.shape == (3, 4)  # Rank 1 has shard indexed by `[3:6, ...]`.
param_1.local_shard.shape == (2, 4)  # Rank 1 has shard indexed by `[2:4, ...]`.

send_buffers = {}  # Rank 1 owns everything.
recv_sizes = {0: 20}  # 12 + 8 = 20 elements from rank 0
# Rank 1's own shards, flattened (views):
own_shards = {0: param_0.local_shard.view(-1), 1: param_1.local_shard.view(-1)}
recv_offsets = {
    (0, 0): 0,  # `param_0` (tensor index 0) from rank 0: offset 0
    (1, 0): 12,  # `param_1` (tensor index 1) from rank 0: offset 12
}

.. attribute:: send_buffers

Per-destination-owner flat send buffer (this rank’s shards for that owner’s params, in tensor-index order). Only owners with non-zero send size appear.

.. attribute:: recv_sizes

Per-source-rank element count this rank (as an owner) receives. Only sources with non-zero total size appear.

.. attribute:: own_shards

This rank’s flat local shard per owned parameter, keyed by tensor index (used directly in reconstruction, not communicated).

.. attribute:: recv_offsets

Per (tensor_index, src_rank), the flat offset of this param’s shard inside the recv buffer received from src_rank. Only contains tuples for which src_rank holds elements.

layouts: dict[int, core.distributed.fsdp.src.megatron_fsdp.experimental.owner_planning.ParameterLayout]#: None

this_rank: int#: None

send_buffers: dict[int, torch.Tensor]#: None

recv_sizes: dict[int, int]#: None

own_shards: dict[int, torch.Tensor]#: None

recv_offsets: dict[tuple[int, int], int]#: None

classmethod pack( plan: core.distributed.fsdp.src.megatron_fsdp.experimental.owner_planning.GroupOwnerLayout, local_shards: dict[int, torch.Tensor], ) → Self#

Pack this rank’s local shards into per-owner P2P send buffers.

Parameters:

plan – The group’s owner layout.
local_shards – This rank’s local shard per parameter, only required for every parameter it holds elements of. Shards may be passed in any shape.

reconstruct_full( param_index: int, recv_buffers: dict[int, torch.Tensor], ) → torch.Tensor#

Reconstruct the full flat tensor for one owned parameter from its per-rank shards.

Concatenates the per-rank shards in rank order (i.e., global element order). Results can be viewed into the desired shape. For a parameter only this rank holds elements of, the own flat shard is returned directly.

Parameters:

param_index – Tensor index of the parameter (a key of the layouts dict passed to pack).
recv_buffers – Per-source-rank received buffer (only sources that sent).

class core.distributed.fsdp.src.megatron_fsdp.experimental.owner_planning.OwnerScatterPlan#

Metadata and send buffers for the owner-scatter P2P step of a set of parameters.

The owner keeps its own result shard (applied directly), so it only sends to the other shard-holding ranks.

.. attribute:: send_buffers

Per-destination-rank flat send buffer (this owner’s result shards for the params it owns, in tensor-index order). Only destinations with non-zero send size appear.

.. attribute:: recv_sizes

Per-owner-rank element count this rank (as a destination) receives. Only owners with non-zero total size appear.

.. attribute:: recv_offsets

Per (tensor_index, owner_rank), the flat offset of this param’s result shard inside the recv buffer received from owner_rank. Only contains tuples for which this rank holds elements.

layouts: dict[int, core.distributed.fsdp.src.megatron_fsdp.experimental.owner_planning.ParameterLayout]#: None

this_rank: int#: None

send_buffers: dict[int, torch.Tensor]#: None

recv_sizes: dict[int, int]#: None

recv_offsets: dict[tuple[int, int], int]#: None

classmethod pack( plan: core.distributed.fsdp.src.megatron_fsdp.experimental.owner_planning.GroupOwnerLayout, full_results: dict[int, torch.Tensor], ) → Self#

Pack this owner rank’s full results into per-destination P2P send buffers.

Parameters:

plan – The group’s owner layout.
full_results – Full result tensor per parameter this rank owns. Tensors may be passed in any shape.

unpack( recv_buffers: dict[int, torch.Tensor], ) → dict[int, torch.Tensor]#

Extract this rank’s local flat result shards from the per-owner recv buffers.

Parameters:: recv_buffers – Per-owner-rank received buffer (only owners that sent).
Returns:: Mapping from tensor index to this rank’s local flat result shard, for parameters this rank holds elements of but does NOT own.

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

Module Contents#

Classes#

Functions#

API#

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