Is this page helpful?

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(
tensor_shapes: collections.abc.Iterable[core.distributed.fsdp.src.megatron_fsdp.experimental.layout.Shape],
tensor_to_offset: collections.abc.Iterable[int],
size: int,
rank_to_offset: collections.abc.Iterable[int],
)#

Global tensor layout in element coordinates.

.. attribute:: tensor_shapes

Logical tensor shapes in tensor-id order.

.. attribute:: tensor_to_offset

Global element offset of each logical tensor’s first element.

.. attribute:: size

Total number of global elements, including any padding or gaps.

.. attribute:: rank_to_offset

Global element offset of each rank segment’s first element (dp_size entries). Segment k spans rank_to_offset[k] up to the next entry, or size for the last segment; see rank_size. On a 1-D mesh k is the rank; on a multi-axis mesh get_rank_range maps the ranks of the Shard axes to k. Segment sizes are equal (size // dp_size) for RowAtomic / BlockAtomic; generally uneven for TensorAtomic.

Initialization

Normalize iterable inputs to immutable tuples and validate the layout.

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

None

tensor_to_offset: tuple[int, ...]#

None

size: int#

None

rank_to_offset: tuple[int, ...]#

None

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

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

__post_init__() → None#

Validate offsets are in bounds and tensor ranges do not overlap.

validate_for_row_atomic(*, block_size: int = 1) → None#

Check equal-size shards and row/block alignment from the actual offsets.

validate_for_tensor_atomic() → None#

Check every nonempty tensor is contained in a single rank segment.

property dp_size: int#

Number of segments (data-parallel shards) this layout was planned for.

rank_size(rank: int) → int#

Number of elements in segment rank.

property has_equal_shard_sizes: bool#

Whether every rank’s segment has the same numel.

get_rank_range(
mesh: torch.distributed.DeviceMesh,
placements: collections.abc.Iterable[torch.distributed.tensor.placement_types.Placement],
rank: int,
) → core.distributed.fsdp.src.megatron_fsdp.experimental.range.Range#

Return the global buffer element range for a global process rank.

Parameters:

  • mesh – The device mesh defining the distribution.

  • placements – Per-mesh-axis placements.

  • rank – Global process rank belonging to mesh.

get_local_range(
mesh: torch.distributed.DeviceMesh,
placements: collections.abc.Iterable[torch.distributed.tensor.placement_types.Placement],
) → core.distributed.fsdp.src.megatron_fsdp.experimental.range.Range#

Return this rank’s element range in the global buffer for placements.

get_tensor_range(
tensor_index: int,
) → core.distributed.fsdp.src.megatron_fsdp.experimental.range.Range#

Return the full tensor’s element range in the global buffer.

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#