core.distributed.fsdp.src.megatron_fsdp.experimental.layout_builder

Build global tensor layouts for DBuffer.

Module Contents

Functions

build_row_atomic_layout	Build a RowAtomic/BlockAtomic layout with the builder selected by shard_layout.
build_for_row_atomic	Build equal-size rank segments for RowAtomic or BlockAtomic placements.
build_for_row_atomic_unpadding	Build padding-free, generally unequal rank segments for RowAtomic or BlockAtomic.
build_for_tensor_atomic	Compute global tensor element offsets from a per-tensor owner-rank list.
_row_atomic_alignments	Return each tensor’s row/block alignment, block_size * row_size, in elements.
_pad_to_multiple

API

core.distributed.fsdp.src.megatron_fsdp.experimental.layout_builder.build_row_atomic_layout(

shapes: collections.abc.Iterable[core.distributed.fsdp.src.megatron_fsdp.experimental.layout.Shape],

dp_size: int,

*,

block_size: int = 1,

shard_layout: core.distributed.fsdp.src.megatron_fsdp.experimental.placement.ShardLayout,

) → core.distributed.fsdp.src.megatron_fsdp.experimental.layout.GlobalLayout

Build a RowAtomic/BlockAtomic layout with the builder selected by shard_layout.

core.distributed.fsdp.src.megatron_fsdp.experimental.layout_builder.build_for_row_atomic(

shapes: collections.abc.Iterable[core.distributed.fsdp.src.megatron_fsdp.experimental.layout.Shape],

dp_size: int,

*,

block_size: int = 1,

) → core.distributed.fsdp.src.megatron_fsdp.experimental.layout.GlobalLayout

Build equal-size rank segments for RowAtomic or BlockAtomic placements.

This is a DBuffer-specific reimplementation of param_and_grad_buffer.build_data_parallel_buffer_index. It keeps only the global offset construction and final padding so each rank-local shard size is a multiple of chunk_size; DBuffer derives rank-local slices later through DTensor placements.

chunk_size is the least common multiple of each tensor’s row size (shape[1:].numel()), additionally multiplied by block_size for BlockAtomic. For example, with shapes P0=(2, 6), P1=(4, 4), P2=(4, 4), P3=(1, 2), P4=(1, 6), chunk_size = LCM(6, 4, 4, 2, 6) = 12 and a 5-rank DP layout has equal-size rank shards:

rank 0 [ 0, 12): | P0 row 0              | P0 row 1               |
rank 1 [12, 24): | P1 row 0      | P1 row 1      | P1 row 2       |
rank 2 [24, 36): | P1 row 3      | P3    | gap   | P2 row 0       |
rank 3 [36, 48): | P2 row 1      | P2 row 2      | P2 row 3       |
rank 4 [48, 60): | P4                    | pad                    |

The diagram uses four character columns per element; segment widths are proportional. Every chunk boundary is aligned to each tensor’s row size, so each DP shard owns full rows even when fragments fill regular-tensor padding gaps.

Parameters:

• shapes – Logical tensor shapes in tensor-id order.

• dp_size – Data-parallel shard count for this global layout.

• block_size – Number of consecutive rows kept on one rank. Defaults to one for RowAtomic; larger values build a BlockAtomic layout.

Returns:

Global layout with row-aligned tensor offsets and a total size padded to a multiple of chunk_size * dp_size, so every rank-local shard length is a multiple of chunk_size.

core.distributed.fsdp.src.megatron_fsdp.experimental.layout_builder.build_for_row_atomic_unpadding(

shapes: collections.abc.Iterable[core.distributed.fsdp.src.megatron_fsdp.experimental.layout.Shape],

dp_size: int,

*,

block_size: int = 1,

) → core.distributed.fsdp.src.megatron_fsdp.experimental.layout.GlobalLayout

Build padding-free, generally unequal rank segments for RowAtomic or BlockAtomic.

Here, tensors are packed back to back in tensor-id order, so size is the sum of all tensor numels. Each rank boundary is then placed as close as possible to an even size / dp_size split while landing on a row (or block_size-row) boundary of the tensor it falls in, so every rank still owns whole rows or blocks. Segment sizes therefore differ by at most one row or block.

For example, with shapes P0=(2, 6), P1=(4, 4), P2=(1, 2) on 2 ranks, the even split at 15 is rounded to P1’s row boundary at 16:

rank 0 [ 0, 16): | P0 row 0              | P0 row 1               | P1 row 0      |
rank 1 [16, 30): | P1 row 1      | P1 row 2      | P1 row 3      | P2    |

The padded layout would instead use chunk_size = LCM(6, 4, 2) = 12, pad the total to a multiple of 12 * 2, and allocate 48 elements for 30 elements of data.

Parameters:

• shapes – Logical tensor shapes in tensor-id order.

• dp_size – Data-parallel shard count for this global layout.

• block_size – Number of consecutive rows kept on one rank. Defaults to one for RowAtomic; larger values build a BlockAtomic layout.

Returns:

Global layout whose size is the total number of logical elements and whose rank boundaries never split a row or block.

core.distributed.fsdp.src.megatron_fsdp.experimental.layout_builder.build_for_tensor_atomic(

shapes: collections.abc.Iterable[core.distributed.fsdp.src.megatron_fsdp.experimental.layout.Shape],

dp_size: int,

*,

tensor_owners: collections.abc.Iterable[int],

) → core.distributed.fsdp.src.megatron_fsdp.experimental.layout.GlobalLayout

Compute global tensor element offsets from a per-tensor owner-rank list.

Pack tensors by ascending owner rank, preserving tensor-id order within each rank. Logical tensor IDs and shapes stay in their original order; only their offsets reflect the packing. Rank r’s segment contains exactly its tensors. Tensors are never split across ranks. Because segments are packed tightly, there is no padding: size equals the sum of all tensor numels. Segment lengths are in general different per rank.

For example, with shapes P0=(4, 4), P1=(2, 6), P2=(1, 2) and tensor_owners=(0, 1, 1) on 2 ranks:

rank 0 [ 0, 16): | P0                      |
rank 1 [16, 30): | P1             | P2     |

Parameters:

• shapes – Logical tensor shapes in tensor-id order.

• dp_size – Data-parallel shard count for this global layout.

• tensor_owners – Owner rank of each tensor, in tensor-id order. Every rank must be an integer in [0, dp_size).

Returns:

Global layout whose per-rank segments are gap-free and whose size is the total number of logical elements.

core.distributed.fsdp.src.megatron_fsdp.experimental.layout_builder._row_atomic_alignments(

tensor_shapes: tuple[torch.Size, ...],

block_size: int,

) → tuple[int, ...]

Return each tensor’s row/block alignment, block_size * row_size, in elements.

core.distributed.fsdp.src.megatron_fsdp.experimental.layout_builder._pad_to_multiple(value: int, multiple: int) → int