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

Joint owner assignment across MFSDP v2 distributed buffers.

Tensor shapes and shard ranges come from the buffer’s layout and placements. Owner assignments use global process ranks.

Module Contents

Classes

_OwnerWork

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

Functions

assign_owner_work

Jointly assign tensors with at least two dimensions to global ranks holding their shards.

API

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

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

buffer_index: int

None

tensor_index: int

None

cost: float

None

candidates: list[int]

None

core.distributed.fsdp.src.megatron_fsdp.experimental.owner_planning.assign_owner_work(

buffers: collections.abc.Iterable[core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer],

) → list[dict[int, int]]

Jointly assign tensors with at least two dimensions to global ranks holding their shards.

Account for single-candidate tensors across all buffers first, then assign the remaining tensors in descending cost order across buffers. Each tensor goes to its least-loaded holder. Equal-cost tensors follow input buffer order and tensor index; equal-load holders follow buffer offset order. Inputs must be identical across ranks.

Costs estimate five Newton-Schulz iterations from each 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) * 5 + 1) as a relative work estimate, omitting constant factors and adding a linear term for normalization and elementwise work.

Parameters:

buffers – Distributed buffers on the same DP mesh.

Returns:

One tensor-index-to-owner mapping per buffer, in input order. Buffers with no eligible tensors have an empty mapping.