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

Owner assignment and P2P packing for MFSDP v2’s all-RowAtomic layout.

Tensor shapes and shard ranges come from the group’s existing GlobalLayout. Owner assignments and peer dictionaries use global process ranks.

Module Contents#

Classes#

`_OwnerWork`	A tensor’s compute cost and eligible owner ranks in buffer order.
`GroupOwnerLayout`	A group’s shared buffer layout and owner assignments.
`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`	Estimate Newton-Schulz work from a tensor’s full shape.
`assign_owner_work`	Assign each participating tensor to a global rank holding part of it.

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.

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

Estimate Newton-Schulz work from a tensor’s full shape.

View the tensor as (rows, cols) = (shape[0], shape[1:].numel()). After normalization and transposing so that rows <= cols, a quintic Newton-Schulz step with scalar coefficients (a, b, c) is:

A = X @ X.T
X = a * X + (b * A + c * (A @ A)) @ X

The matrix products cost O(numel * min(rows, cols)) per step. We use numel * (min(rows, cols) * num_ns_steps + 1) as a relative work estimate, omitting constant factors and adding a linear term for normalization and elementwise work.

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

A tensor’s compute cost and eligible owner ranks in buffer order.

tensor_index: int#: None

cost: float#: None

candidates: list[int]#: None

core.distributed.fsdp.src.megatron_fsdp.experimental.owner_planning.assign_owner_work( layout: core.distributed.fsdp.src.megatron_fsdp.experimental.layout.GlobalLayout, mesh: torch.distributed.device_mesh.DeviceMesh, tensor_indices: collections.abc.Iterable[int], cost_fn: collections.abc.Callable[[torch.Size], float] | None = None, ) → dict[int, int]#

Assign each participating tensor to a global rank holding part of it.

Tensors with a single candidate are assigned first, followed by tensors with multiple candidates in descending cost order. Each tensor is assigned to its least-loaded holder. Ties follow buffer order.

Parameters:

layout – The group’s DBuffer layout, including padding and ineligible tensors.
mesh – The device mesh across which the buffer is all-RowAtomic sharded.
tensor_indices – Indices of participating tensors in layout.
cost_fn – Positive cost estimate from a tensor’s full shape. Defaults to Newton-Schulz with five iterations.

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

A group’s shared buffer layout and owner assignments.

.. attribute:: mesh

The device mesh over which the buffer is all-RowAtomic sharded.

.. attribute:: layout

The existing DBuffer.layout, including all tensors in the group.

.. attribute:: tensor_to_owner

Participating tensor index to owner global process rank. These keys identify the participating tensors; no filtered layout copy is stored.

mesh: torch.distributed.device_mesh.DeviceMesh#: None

layout: core.distributed.fsdp.src.megatron_fsdp.experimental.layout.GlobalLayout#: None

tensor_to_owner: dict[int, int]#: None

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

Select participating parameters and balance their owner assignments.

eligible_fn defaults to selecting ≥2D parameters. cost_fn receives full tensor shapes; see assign_owner_work for the default estimate.

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 buffer 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.
mesh.mesh.tolist() == [0, 1]  # Global process ranks.
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]}]
param_0.shape == (6, 4)  # Global shape.
param_1.shape == (4, 4)  # Global shape.
layout.tensor_to_offset == (0, 24)
layout.size == 40
# The buffer, rather than each parameter separately, is split evenly:
# rank 0 holds [0, 20); rank 1 holds [20, 40).
# Choose rank 1 as both owners for this example. It holds part of param_0
# and all of param_1; these are explicit assignments, not the default balancer's output.
owner_layout = GroupOwnerLayout(mesh=mesh, layout=layout, tensor_to_owner={0: 1, 1: 1})

param_0.local_shard.shape == (5, 4)  # Rank 0 holds param_0[0:5, ...].
param_1.local_shard.shape == (0, 4)  # Rank 0 holds none of param_1.
param_0.local_shard.numel() == 20
param_1.local_shard.numel() == 0

owner_gather_plan.send_buffers == {1: tensor(20)}  # 20 elements from param_0.
# `owner_gather_plan.send_buffers[1]` represents the following flat buffer:
#   +--------------------+
#   | param_0 (20 elems) |
#   +--------------------+
#      element 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 == (1, 4)  # Rank 1 holds param_0[5:6, ...].
param_1.local_shard.shape == (4, 4)  # Rank 1 holds all of param_1.

send_buffers = {}  # Rank 1 owns everything.
recv_sizes = {0: 20}  # 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.
    # No entry for param_1: it is fully local to rank 1.
}
# Reconstruct param_0 from the received 20 elements followed by its own 4.
# Reconstruct param_1 directly from its own 16 elements, without communication.

.. 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.

layout: core.distributed.fsdp.src.megatron_fsdp.experimental.layout.GlobalLayout#: None

rank_ranges: dict[int, core.distributed.fsdp.src.megatron_fsdp.experimental.range.Range]#: 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( owner_layout: 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:

owner_layout – 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 buffer 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 an owned parameter in the group layout.
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.

layout: core.distributed.fsdp.src.megatron_fsdp.experimental.layout.GlobalLayout#: None

local_range: core.distributed.fsdp.src.megatron_fsdp.experimental.range.Range#: None

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

recv_sizes: dict[int, int]#: None

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

classmethod pack( owner_layout: 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:

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