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

Transformer Engine MXFP8 distributed buffers composed from physical planes.

Module Contents#

Classes#

QuantizedDBuffer

The MFSDP storage and lifecycle for one TE MXFP8 tensor.

Functions#

`_get_quantizer`	Return a shared MXFP8 quantizer with the given usages.
`effective_dtype`	Return MFSDP’s storage dtype for a parameter.
`_rowwise_scale_layout`	Derive rowwise scales from the layout shared by rowwise_data and columnwise_data.
`_columnwise_scale_layout`	Derive columnwise scales from the layout shared by rowwise_data and columnwise_data.
`_block_atomic_to_flat`	Replace BlockAtomic placements with RowAtomic for coordinates measured in blocks.
`_pad_rowwise_scale`	Pad rowwise scales to TE’s physical allocation shape.
`_pad_columnwise_scale`	Pad columnwise scales to TE’s physical allocation shape.
`update_mxfp8_tensor`	Update `dst` in place from `src`, a compact view of the allocated planes.

Data#

`_MXFP8_DTYPE`
`_MXFP8_BLOCK_SIZE`

API#

core.distributed.fsdp.src.megatron_fsdp.experimental.quantized_dbuffer._MXFP8_DTYPE#: None

core.distributed.fsdp.src.megatron_fsdp.experimental.quantized_dbuffer._MXFP8_BLOCK_SIZE#: 32

core.distributed.fsdp.src.megatron_fsdp.experimental.quantized_dbuffer._get_quantizer( *, rowwise: bool, columnwise: bool, ) → transformer_engine.pytorch.tensor.mxfp8_tensor.MXFP8Quantizer#: Return a shared MXFP8 quantizer with the given usages.

core.distributed.fsdp.src.megatron_fsdp.experimental.quantized_dbuffer.effective_dtype(tensor: torch.Tensor) → torch.dtype#: Return MFSDP’s storage dtype for a parameter.

core.distributed.fsdp.src.megatron_fsdp.experimental.quantized_dbuffer._rowwise_scale_layout( data_layout: core.distributed.fsdp.src.megatron_fsdp.experimental.layout.GlobalLayout, ) → core.distributed.fsdp.src.megatron_fsdp.experimental.layout.GlobalLayout#: Derive rowwise scales from the layout shared by rowwise_data and columnwise_data.

core.distributed.fsdp.src.megatron_fsdp.experimental.quantized_dbuffer._columnwise_scale_layout( data_layout: core.distributed.fsdp.src.megatron_fsdp.experimental.layout.GlobalLayout, ) → core.distributed.fsdp.src.megatron_fsdp.experimental.layout.GlobalLayout#: Derive columnwise scales from the layout shared by rowwise_data and columnwise_data.

core.distributed.fsdp.src.megatron_fsdp.experimental.quantized_dbuffer._block_atomic_to_flat( placements: collections.abc.Iterable[torch.distributed.tensor.placement_types.Placement], ) → tuple[torch.distributed.tensor.placement_types.Placement, ...]#

Replace BlockAtomic placements with RowAtomic for coordinates measured in blocks.

For example, one MXFP8 columnwise scale row represents a 32-row weight block, so RowAtomic preserves the shard boundaries of BlockAtomic(32).

core.distributed.fsdp.src.megatron_fsdp.experimental.quantized_dbuffer._pad_rowwise_scale(scale: torch.Tensor) → torch.Tensor#: Pad rowwise scales to TE’s physical allocation shape.

core.distributed.fsdp.src.megatron_fsdp.experimental.quantized_dbuffer._pad_columnwise_scale(scale: torch.Tensor) → torch.Tensor#: Pad columnwise scales to TE’s physical allocation shape.

core.distributed.fsdp.src.megatron_fsdp.experimental.quantized_dbuffer.update_mxfp8_tensor( src: transformer_engine.pytorch.tensor.mxfp8_tensor.MXFP8Tensor, dst: transformer_engine.pytorch.tensor.mxfp8_tensor.MXFP8Tensor, *, rowwise: bool = True, columnwise: bool = True, ) → None#

Update dst in place from src, a compact view of the allocated planes.

rowwise and columnwise select the newly gathered planes, which are taken from src with padded scales. Planes gathered earlier for the other phase stay as installed on dst, and planes absent from src are cleared. src is not modified.

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

The MFSDP storage and lifecycle for one TE MXFP8 tensor.

The data layout owns the parameter-to-rank partition. Scale planes use its compact MXFP8 coordinates, so every plane on a rank describes the same local data rows. TE-only padding is added when materializing a wrapper for compute, not to distributed storage.

Initialization

Allocate MXFP8 data and scale planes from a shared weight layout.

rowwise_data: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer#: None

columnwise_data: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer#: None

rowwise_scale: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer#: None

columnwise_scale: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer#: 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], device: torch.device | str, *, shard_layout: core.distributed.fsdp.src.megatron_fsdp.experimental.placement.ShardLayout = ShardLayout.PADDED, ) → core.distributed.fsdp.src.megatron_fsdp.experimental.quantized_dbuffer.QuantizedDBuffer#

Build an MXFP8 layout from logical tensor shapes and allocate its planes.

Parameters:

mesh – Device mesh whose dimensions correspond to placements.
placements – Per-mesh-axis DBuffer placements.
tensor_shapes – Logical 2D shapes that the buffer manages.
device – Device for the local planes.
shard_layout – Padded (equal) or unpadded (32-row-block-aligned, unequal) rank segments; see DBuffer.empty.

property mesh: torch.distributed.DeviceMesh#: Device mesh shared by all physical planes.

property placements: tuple[torch.distributed.tensor.placement_types.Placement, ...]#: Logical weight placements; scale planes map these to scale coordinates.

classmethod _from_planes( rowwise_data: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer, columnwise_data: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer, rowwise_scale: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer, columnwise_scale: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer, ) → core.distributed.fsdp.src.megatron_fsdp.experimental.quantized_dbuffer.QuantizedDBuffer#: Create a composed view from already-allocated physical planes.

property is_rowwise_allocated: bool#: Whether rowwise data and scales are allocated.

property is_columnwise_allocated: bool#: Whether columnwise data and scales are allocated.

get_tensor_view( index: int, ) → transformer_engine.pytorch.tensor.mxfp8_tensor.MXFP8Tensor#

Return a compact, unswizzled MXFP8 view that aliases every allocated plane pair.

Released planes are omitted, so the view never reads storage that release_storage() freed.

quantize_( main_weight: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer, ) → None#: Quantize a local master shard with matching mesh, placements, layout, and device.

property planes: tuple[core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer, core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer, core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer, core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer]#: Physical planes in TE’s rowwise-data-first order.

planes_for( *, rowwise: bool, columnwise: bool, ) → tuple[core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer, ...]#: Physical data and scale planes of the selected representations.

property is_symmetric_memory: bool#: Whether every plane is backed by symmetric memory.

reallocate_storage() → None#: Restore every plane’s backing storage.

release_storage(*, rowwise_only: bool = False) → None#: Release all plane storage, or only rowwise storage, while retaining aliases.

view( placements: collections.abc.Iterable[torch.distributed.tensor.placement_types.Placement], ) → core.distributed.fsdp.src.megatron_fsdp.experimental.quantized_dbuffer.QuantizedDBuffer#: Return a storage-sharing view of every physical plane.

redistribute( new_placements: collections.abc.Iterable[torch.distributed.tensor.placement_types.Placement], *, out: QuantizedDBuffer | None = None, rowwise: bool = True, columnwise: bool = True, ) → core.distributed.fsdp.src.megatron_fsdp.experimental.quantized_dbuffer.QuantizedDBuffer#: Redistribute the selected planes, returning out or a new wrapper.

allgather( mesh_axis: int | collections.abc.Iterable[int], *, out: QuantizedDBuffer | None = None, ) → core.distributed.fsdp.src.megatron_fsdp.experimental.quantized_dbuffer.QuantizedDBuffer#: Gather one or more axes in every plane, reusing out when supplied.

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