Implementation:NVIDIA TransformerEngine Tensor Utils
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, PyTorch, Quantization, Distributed |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Utility functions for managing FP8 primary weights in distributed training, including casting master weights to FP8 and swapping raw data buffers in quantized tensors.
Description
replace_raw_data swaps the underlying data buffer of a quantized tensor (Float8Tensor or Float8BlockwiseQTensor) to support custom memory layouts like ZeRO-2 contiguous buffers. cast_master_weights_to_fp8 handles the complex workflow of converting FP32 master weights back to FP8 model weights during distributed training, dispatching to different code paths based on quantizer type (delayed scaling, current scaling, blockwise, MXFP8) and using multi-tensor kernels for efficient batched scale computation and casting. post_all_gather_processing handles post-communication processing of gathered weights. is_custom checks whether a quantizer or storage is a custom (user-defined) type.
Usage
Critical for FP8 primary weights with distributed training (ZeRO/FSDP). The master-to-FP8 casting must handle multiple quantization recipes, distributed amax reduction, and FSDP shard-aware data placement.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/pytorch/tensor/utils.py- Lines
- 1--636
Signature
def replace_raw_data(tensor: QuantizedTensor, new_raw_data: torch.Tensor): ...
def cast_master_weights_to_fp8(
model_weights, master_weights, quantizers,
amax_reduce_group=None, ...
): ...
def post_all_gather_processing(
model_weights: Union[torch.Tensor, List[torch.Tensor]]
): ...
def is_custom(
x: Optional[Union[Quantizer, QuantizedTensorStorage]] = None
) -> bool: ...
Import
from transformer_engine.pytorch.tensor.utils import (
replace_raw_data,
cast_master_weights_to_fp8,
post_all_gather_processing,
is_custom,
)
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| model_weights | List[torch.Tensor] |
Yes | FP8 model weight tensors to update |
| master_weights | List[torch.Tensor] |
Yes | FP32 master weight tensors |
| quantizers | List[Quantizer] |
Yes | Quantizers for each weight |
| amax_reduce_group | ProcessGroup |
No | Distributed group for amax reduction |
Outputs
| Name | Type | Description |
|---|---|---|
| N/A | N/A | Model weights are updated in-place with FP8 data cast from master weights |
Usage Examples
from transformer_engine.pytorch.tensor.utils import cast_master_weights_to_fp8
# Cast master weights back to FP8 after optimizer step
cast_master_weights_to_fp8(
model_weights=fp8_params,
master_weights=fp32_params,
quantizers=quantizer_list,
amax_reduce_group=dp_group,
)