Implementation:NVIDIA TransformerEngine Debug Per Tensor Scaling
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, PyTorch, Debug, Quantization |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Debug feature that overrides delayed scaling with per-tensor current scaling for specific tensors and GEMMs, enabling finer-grained quantization control during training.
Description
PerTensorScaling allows switching specific tensors from delayed scaling to per-tensor current scaling within a DelayedScaling recipe autocast. It computes the scaling factor based on the tensor's current amax value using Float8CurrentScalingQuantizer, rather than the delayed (historical) amax. This is useful for diagnosing whether delayed scaling factors are causing numerical issues. The feature requires that the default quantizer is a Float8Quantizer (delayed scaling) and supports E4M3 and E5M2 FP8 formats.
Usage
Enable via YAML config, specifying which GEMMs and tensors should use per-tensor current scaling instead of delayed scaling. Only usable within a DelayedScaling recipe autocast context.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/debug/features/per_tensor_scaling.py- Lines
- 1--134
Signature
def per_tensor_cast(tensor: torch.Tensor, fp8_dtype: tex.DType, out=None) -> Float8Tensor: ...
@Registry.register_feature(namespace="transformer_engine")
class PerTensorScaling(TEConfigAPIMapper):
def fp8_gemm(self, config, layer_name, gemm, iteration) -> Tuple[bool, None]: ...
def modify_tensor_enabled(self, config, layer_name, tensor_name, gemm, iteration) -> Tuple[bool, int]: ...
def modify_tensor(self, config, layer_name, gemm, tensor_name, tensor, iteration, default_quantizer, out=None, dtype=None) -> Optional[Float8Tensor]: ...
Import
from transformer_engine.debug.features.per_tensor_scaling import PerTensorScaling, per_tensor_cast
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| config | Dict | Yes | GEMM and tensor specification |
| tensor | torch.Tensor | Yes | High-precision GPU tensor |
| default_quantizer | Float8Quantizer | Yes | Must be a Float8Quantizer (delayed scaling)
|
Outputs
| Name | Type | Description |
|---|---|---|
| result | Float8Tensor | Tensor cast to FP8 using per-tensor current scaling |
Usage Examples
# YAML configuration:
# example_per_tensor_scaling:
# enabled: True
# layers:
# layer_types: [transformer_layer.self_attn.layernorm_q]
# transformer_engine:
# PerTensorScaling:
# enabled: True
# gemms: [dgrad]
# tensors: [weight, activation]