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

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.

Use :mod:.layout_builder to construct layouts for specific placements.

.. 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 the padded RowAtomic / BlockAtomic layout of layout_builder.build_for_row_atomic; generally uneven for build_for_row_atomic_unpadding and build_for_tensor_atomic.

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

__post_init__() → None#: Validate offsets are in bounds and tensor ranges do not overlap.

validate_for_row_atomic(*, block_size: int = 1) → None#: Check that no rank boundary splits a row (or block_size-row block).

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#