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

Global tensor layout metadata for DBuffer.

Module Contents#

Classes#

GlobalLayout

Global tensor layout in element coordinates.

Functions#

`non_leading_numel`	Return the number of elements after dim 0 for a non-scalar shape.
`_pad_to_multiple`

Data#

Shape

API#

core.distributed.fsdp.src.megatron_fsdp.experimental.layout.Shape: TypeAlias#: None

class core.distributed.fsdp.src.megatron_fsdp.experimental.layout.GlobalLayout#

Global tensor layout in element coordinates.

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

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

size: int#: None

classmethod build( shapes: collections.abc.Iterable[core.distributed.fsdp.src.megatron_fsdp.experimental.layout.Shape], dp_size: int, ) → core.distributed.fsdp.src.megatron_fsdp.experimental.layout.GlobalLayout#

Compute global tensor element offsets and padded size.

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.

The computed layout is compatible with Flat, TensorAtomic, and BlockAtomic, even though the latter two are not implemented.

chunk_size is the least common multiple of each tensor’s row size (shape[1:].numel()). 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.

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.

__post_init__() → None#: Validate tensor offsets are row-aligned, in bounds, and non-overlapping.

get_local_range( mesh: torch.distributed.DeviceMesh, placements: collections.abc.Iterable[core.distributed.fsdp.src.megatron_fsdp.experimental.placement.Placement], ) → tuple[int, int]#: Return this rank’s local element offset and length for placements.

core.distributed.fsdp.src.megatron_fsdp.experimental.layout.non_leading_numel(shape: torch.Size) → int#: Return the number of elements after dim 0 for a non-scalar shape.

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

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