Implementation:NVIDIA TransformerEngine Debug Quantization

Field	Value
Sources	TransformerEngine
Domains	Deep_Learning, PyTorch, Debug, Quantization
Last Updated	2026-02-07 14:00 GMT

Overview

Provides DebugQuantizer and DebugQuantizedTensor wrappers that intercept the quantization pipeline to enable debug API calls for tensor inspection and modification during training.

Description

This module contains two main classes:

• DebugQuantizer wraps a parent Quantizer to inject debug hooks into the quantization process. During initialization it queries the debug API to determine plans for each tensor direction (rowwise/columnwise): API_CALL_MODIFY (route through modify_tensor), STANDARD_QUANTIZE (use normal quantization), or HIGH_PRECISION (skip quantization). It tracks the next iteration when debug features will be active, enabling CPU overhead optimization when features are inactive.

• DebugQuantizedTensor stores separate rowwise and columnwise GEMM tensors, supporting mixed quantization plans where each direction can have a different strategy. It implements prepare_for_saving/restore_from_saved for autograd compatibility and provides dequantize(), get_tensor(), and update_usage() methods.

Usage

Used internally by Transformer Engine when debug mode is enabled. The DebugQuantizer replaces the standard quantizer for layers selected by the debug config, intercepting quantization to allow features to inspect or modify tensors.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/debug/pytorch/debug_quantization.py

Lines

1--675

Signature

class DebugQuantizer(Quantizer):
    def __init__(self, layer_name, tensor_name, parent_quantizer, tp_group): ...
    def quantize(self, tensor, *, out=None, dtype=None) -> Union[DebugQuantizedTensor, torch.Tensor]: ...
    def process_gemm_output(self, tensor) -> torch.Tensor: ...
    def any_feature_enabled(self) -> bool: ...
    def get_next_debug_iter(self) -> Optional[int]: ...

class DebugQuantizedTensor(QuantizedTensorStorage):
    def __init__(self, rowwise_gemm_tensor, columnwise_gemm_tensor, quantizer, layer_name=None, tensor_name=None): ...
    def get_tensor(self, transpose: bool) -> torch.Tensor: ...
    def dequantize(self, *, dtype=None) -> torch.Tensor: ...

Import

from transformer_engine.debug.pytorch.debug_quantization import DebugQuantizer, DebugQuantizedTensor

I/O Contract

Inputs

Name	Type	Required	Description
layer_name	str	Yes	Name of the TE layer
tensor_name	str	Yes	One of activation, weight, gradient, output, wgrad, dgrad
parent_quantizer	Optional[Quantizer]	Yes	The original quantizer to wrap (None if no quantization)
tp_group	ProcessGroup	Yes	Tensor parallel process group for reductions
tensor	torch.Tensor	Yes	High-precision tensor to quantize (for quantize())

Outputs

Name	Type	Description
result	DebugQuantizedTensor or torch.Tensor	Quantized tensor (or high-precision) depending on debug plans

Usage Examples

# Internal usage within Transformer Engine debug mode:
debug_quantizer = DebugQuantizer(
    layer_name="encoder.layer.0.attention.query",
    tensor_name="activation",
    parent_quantizer=fp8_quantizer,
    tp_group=tp_group,
)

# Check if any debug features are active for this tensor
if debug_quantizer.any_feature_enabled():
    result = debug_quantizer.quantize(input_tensor)

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements