`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.
`_validate_layout`	Check layout coordinates satisfy each requested shard placement.

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.

core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer._validate_layout( layout: core.distributed.fsdp.src.megatron_fsdp.experimental.layout.GlobalLayout, placements: collections.abc.Iterable[torch.distributed.tensor.placement_types.Placement], ) → None#: Check layout coordinates satisfy each requested shard placement.

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, shard_layout: core.distributed.fsdp.src.megatron_fsdp.experimental.placement.ShardLayout = ShardLayout.PADDED, ) → core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer#

Build a RowAtomic or BlockAtomic layout and allocate its local buffer.

For TensorAtomic, use layout_builder.build_for_tensor_atomic and pass the resulting layout to the DBuffer constructor.

Parameters:

mesh – Device mesh whose dimensions correspond to placements.
placements – Per-mesh-axis DBuffer placements.
tensor_shapes – Logical shapes that the DBuffer manages.
dtype – Dtype for the local buffer.
device – Device for the local buffer.
block_size – Number of consecutive rows kept together on one rank.
shard_layout – ShardLayout.PADDED pads to equal segments; ShardLayout.UNPADDED packs without padding into generally unequal segments.

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.

property is_allocated: bool#

Whether the local buffer has backing storage.

release_storage() clears this and reallocate_storage() restores it. Empty local buffers are always allocated.

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#

_peer_chunks( buffer: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer, placements: collections.abc.Iterable[torch.distributed.tensor.placement_types.Placement], mesh_axis: int, ) → list[torch.Tensor]#

Slice buffer into each peer’s range under placements along mesh_axis.

Peers are the ranks of this rank’s mesh_axis process group, in group-rank order.

_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#

View a contained local range with the requested placements.

This only slices and relabels storage; it performs no communication. When changing Partial placements, callers must populate the view with values appropriate to its new placements 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], *, layout: core.distributed.fsdp.src.megatron_fsdp.experimental.layout.GlobalLayout | None = None, ) → 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.
layout – Prebuilt layout matching the input tensor shapes and order. If omitted, build a RowAtomic layout. Use the functions in layout_builder to supply a TensorAtomic or BlockAtomic layout.

Returns:

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

copy_from(index: int, tensor: torch.Tensor) → None#

Copy a full logical tensor’s local owned range into this buffer.

Meta tensors leave their owned range unspecified. Padding and layout gaps are not observable through get_tensor_view() and remain unspecified.

_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 with multi-axis gathers/views or a single-axis reduction.

Only Shard-to-Replicate gathers and Replicate-to-Shard views may change multiple axes. Reductions and other transitions must change one axis. Without out, sharding returns a view and unchanged placements return self. Other transitions allocate the destination as needed.

allgather( mesh_axis: int | collections.abc.Iterable[int], *, out: DBuffer | None = None, ) → core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer#

Gather one or more sharded axes into Replicate placements.

Gather outer-to-inner, writing each intermediate into a view of the final output. This reuses existing mesh groups without allocating scratch storage.

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.

RowAtomic 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`#