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

Distributed tensor buffers for Megatron-FSDP.

Module Contents#

Classes#

`_OwnedRange`
`DBuffer`	A distributed buffer holding a group of logical tensors.

Functions#

`_validate_placements`	Validate DBuffer placements form a supported contiguous local layout.
`_get_reduce_op`	Convert a DTensor Partial reduction name to a torch.distributed op.

API#

class core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer._OwnedRange#

numel: int#: None

tensor_relative_offset: int#: None

buffer_relative_offset: int#: None

core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer._validate_placements( placements: collections.abc.Iterable[torch.distributed.tensor.placement_types.Placement], ) → None#: Validate DBuffer placements form a supported contiguous local layout.

core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer._get_reduce_op( partial_placement: torch.distributed.tensor.Partial, ) → torch.distributed.ReduceOp.RedOpType#: Convert a DTensor Partial reduction name to a torch.distributed op.

class core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer( mesh: torch.distributed.DeviceMesh, placements: collections.abc.Iterable[torch.distributed.tensor.placement_types.Placement], layout: core.distributed.fsdp.src.megatron_fsdp.experimental.layout.GlobalLayout, dtype: torch.dtype, device: torch.device | str, )#

A distributed buffer holding a group of logical tensors.

DBuffer is analogous to DTensor, but manages a group of logical tensors in one local storage tensor. It stores enough metadata to return per-tensor views, redistribute the buffer across mesh axes, and materialize per-tensor DTensors for optimizer state or distributed checkpointing.

Initialization

Create a DBuffer and allocate its local buffer.

Parameters:

mesh – Device mesh whose dimensions correspond to placements.
placements – Per-mesh-axis DBuffer placements.
layout – Global shapes, offsets, and allocation size for this buffer.
dtype – Dtype for the local buffer.
device – Device for the local buffer.

mesh: torch.distributed.DeviceMesh#: None

placements: tuple[torch.distributed.tensor.placement_types.Placement, ...]#: None

layout: core.distributed.fsdp.src.megatron_fsdp.experimental.layout.GlobalLayout#: None

offset: int#: None

local_buffer: torch.Tensor#: None

classmethod empty( mesh: torch.distributed.DeviceMesh, placements: collections.abc.Iterable[torch.distributed.tensor.placement_types.Placement], tensor_shapes: collections.abc.Iterable[core.distributed.fsdp.src.megatron_fsdp.experimental.layout.Shape], dtype: torch.dtype, device: torch.device | str, *, block_size: int = 1, ) → core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer#: Build a layout from logical tensor shapes and allocate its local buffer.

property dtype: torch.dtype#: Dtype of the local buffer.

property device: torch.device#: Device of the local buffer.

property is_symmetric_memory: bool#: Whether the local buffer is allocated from symmetric memory.

reallocate_storage() → None#: Restore the local buffer’s backing storage to its logical size.

release_storage() → None#: Release local buffer storage without replacing the Storage object.

rendezvous(mesh_axis: int) → None#: Rendezvous this local buffer for symmetric-memory collectives.

_resize_storage(numel: int) → None#

_get_owned_range( tensor_index: int, ) → core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer._OwnedRange | None#: Return this buffer’s owned range for logical tensor tensor_index.

classmethod from_local( local_buffer: torch.Tensor, mesh: torch.distributed.DeviceMesh, placements: collections.abc.Iterable[torch.distributed.tensor.placement_types.Placement], layout: core.distributed.fsdp.src.megatron_fsdp.experimental.layout.GlobalLayout, ) → core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer#

Create a DBuffer from an existing local buffer.

Parameters:

local_buffer – Contiguous local tensor storage for this rank. DBuffer uses it directly in collectives such as all-gather and reduce-scatter, which are efficient with contiguous tensors.
mesh – Device mesh whose dimensions correspond to placements.
placements – Per-mesh-axis DBuffer placements.
layout – Existing global layout for the logical tensors in this buffer.

Returns:

A DBuffer that reuses local_buffer without allocating storage.

view( placements: collections.abc.Iterable[torch.distributed.tensor.placement_types.Placement], ) → core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer#

Return a storage-sharing buffer with supported placements.

Views preserve placements, relabel a full local buffer, or locally slice one full local buffer to Flat. A view that changes a Partial placement is only a storage destination: callers must populate it with a reduction before reading it.

classmethod distribute_tensors( tensors: collections.abc.Iterable[torch.Tensor], mesh: torch.distributed.DeviceMesh, placements: collections.abc.Iterable[torch.distributed.tensor.placement_types.Placement], *, block_size: int = 1, ) → core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer#

Distribute full local tensors into a DBuffer.

Parameters:

tensors – Full tensors available on this rank. Meta tensors contribute shape and dtype metadata but no values.
mesh – Device mesh whose dimensions correspond to placements.
placements – Per-mesh-axis DBuffer placements.

Returns:

A DBuffer whose real local storage matches placements. Ranges corresponding to meta tensors are left uninitialized.

_create_or_validate_out( out: DBuffer | None, *, placements: collections.abc.Iterable[torch.distributed.tensor.placement_types.Placement] | None = None, dtype: torch.dtype | None = None, ) → core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer#

cast( dtype: torch.dtype, *, out: DBuffer | None = None, ) → core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer#: Return this buffer with the same layout and placements in dtype.

redistribute( new_placements: collections.abc.Iterable[torch.distributed.tensor.placement_types.Placement], *, out: DBuffer | None = None, ) → core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer#

Redistribute this buffer to new_placements.

This dispatcher supports the one-axis transitions: Flat -> Replicate, Partial -> Replicate, Partial -> Flat, Replicate -> Flat, and Replicate -> Partial. Other placement changes are intentionally unsupported.

allgather( mesh_axis: int, *, out: DBuffer | None = None, ) → core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer#: All-gather a sharded axis into Replicate placement.

allreduce( mesh_axis: int, *, out: DBuffer | None = None, ) → core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer#: All-reduce a Partial axis into Replicate placement.

reduce_scatter( mesh_axis: int, new_placement: torch.distributed.tensor.placement_types.Placement, *, out: DBuffer | None = None, ) → core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer#: Reduce-scatter a Partial axis into new_placement.

get_tensor_view(index: int) → torch.Tensor#

Return this rank’s local view for logical tensor index.

Flat placements shard dim 0, so the returned view preserves all non-leading dimensions and only changes the leading dimension.

get_dtensor(index: int) → torch.distributed.tensor.DTensor#: Return logical tensor index as a DTensor.

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

Module Contents#

Classes#

Functions#

API#

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