Implementation:NVIDIA TransformerEngine JAX Quantize Helper
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, JAX, Quantization |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Provides the quantization configuration management system, including recipe-based configuration, the autocast/fp8_autocast context manager, hardware capability checks, and Flax collection state management for FP8 training.
Description
autocast (aliased as fp8_autocast) is the main context manager that sets a global quantization recipe and mesh resource. get_quantize_config_with_recipe maps Recipe objects (DelayedScaling, Float8CurrentScaling, MXFP8BlockScaling, NVFP4BlockScaling) to config classes that define per-tensor scaling modes and quantizer factories. Hardware support checks (is_fp8_available, is_scaling_mode_supported) query GPU compute capability, cuBLAS, and CUDA versions. update_collections manages the Flax variable collection for FP8 metadata (scales, amax history). TensorSource enum distinguishes forward (x, kernel) from backward (grad) quantization parameters.
This is the central configuration hub for the entire quantization system -- users interact with autocast to enable FP8 training, and all downstream modules query the global recipe to determine how to quantize their tensors.
Usage
Use the autocast context manager to enable FP8 quantization globally. Use is_fp8_available to check hardware support before enabling FP8. Use get_quantize_config_with_recipe for programmatic quantization configuration.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/jax/quantize/helper.py- Lines
- 1--946
Signature
class TensorSource(Enum):
X = "x"
KERNEL = "kernel"
GRAD = "grad"
class AmaxComputeAlgo(Enum):
MAX = "max"
MOST_RECENT = "most_recent"
class BaseQuantizeConfig(ABC): ...
class NoOpQuantizeConfig(BaseQuantizeConfig): ...
class DelayedScalingQuantizeConfig(BaseQuantizeConfig): ...
class CurrentScalingQuantizeConfig(BaseQuantizeConfig): ...
class BlockScalingQuantizeConfig(BaseQuantizeConfig): ...
class NVFP4ScalingQuantizeConfig(BaseQuantizeConfig): ...
def is_fp8_available() -> Tuple[bool, str]: ...
def is_scaling_mode_supported(scaling_mode: ScalingMode) -> Tuple[bool, str]: ...
def get_supported_scaling_modes() -> List[ScalingMode]: ...
def get_quantize_config_with_recipe(fp8_recipe: Recipe) -> BaseQuantizeConfig: ...
def autocast(
recipe: Optional[Recipe] = None,
mesh_resource: Optional[MeshResource] = None,
**kwargs,
) -> contextlib.contextmanager: ...
def fp8_autocast(
recipe: Optional[Recipe] = None,
mesh_resource: Optional[MeshResource] = None,
**kwargs,
) -> contextlib.contextmanager: ...
def update_collections(new: Collection, original: Collection) -> Collection: ...
Import
from transformer_engine.jax.quantize.helper import autocast, is_fp8_available, get_quantize_config_with_recipe
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| recipe | Optional[Recipe] |
No | Quantization recipe (DelayedScaling, Float8CurrentScaling, etc.) |
| mesh_resource | Optional[MeshResource] |
No | Mesh resource for distributed training |
Outputs
| Name | Type | Description |
|---|---|---|
| context | contextmanager |
Context manager that enables FP8 quantization globally |
Usage Examples
from transformer_engine.jax.quantize.helper import autocast, is_fp8_available
from transformer_engine.common.recipe import DelayedScaling
# Check FP8 availability
available, reason = is_fp8_available()
if available:
recipe = DelayedScaling()
with autocast(recipe=recipe, mesh_resource=mesh_resource):
output = model.apply(params, input_data)