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

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

Data#

`_MXFP8_DTYPE`
`_MXFP8_QUANTIZER`
`_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_QUANTIZER#: ‘MXFP8Quantizer(…)’

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

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 Flat for coordinates measured in blocks.

For example, one MXFP8 columnwise scale row represents a 32-row weight block, so Flat 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.

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], tensor_shapes: collections.abc.Iterable[torch.Size], 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

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

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.

get_tensor_view( index: int, ) → transformer_engine.pytorch.tensor.mxfp8_tensor.MXFP8Tensor#: Return a compact, unswizzled MXFP8 view that aliases all local planes.

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

Return an unswizzled compute tensor with scales padded for TE’s GEMM path.

Data planes remain views. Scale planes alias storage only when no padding is needed; otherwise they are copied into padded allocations.

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.

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

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

release_storage() → None#: Release every plane’s backing 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, ) → core.distributed.fsdp.src.megatron_fsdp.experimental.quantized_dbuffer.QuantizedDBuffer#: Redistribute every plane, returning out when supplied or a new wrapper.

allgather( mesh_axis: int, *, out: QuantizedDBuffer | None = None, ) → core.distributed.fsdp.src.megatron_fsdp.experimental.quantized_dbuffer.QuantizedDBuffer#: All-gather every plane, returning out when supplied or a new wrapper.

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