Transformer Engine (TE) Documentation

Transformer Engine (TE) is a library designed to accelerate Transformer models on NVIDIA GPUs. It provides improved performance and lower memory utilization for both training and inference.

Key features include:

• Support for FP8 (8-bit floating point) precision on Hopper and Ada GPUs.
• Support for NVFP4 (8-bit and 4-bit floating point) precision on Blackwell GPUs.
• Highly optimized building blocks for popular Transformer architectures.
• An automatic-mixed-precision-like API for seamless integration with existing deep learning code.

Supported frameworks: PyTorch and JAX.

Transformer Engine's JAX attention APIs support self-attention and cross-attention using Multi-Head Attention (MHA), Grouped-Query Attention (GQA), and Multi-Query Attention (MQA).

Supported Layouts:

• Standard BSHD (Batch, Sequence, Heads, Dimension).
• Packed THD (Total tokens, Heads, Dimension).
• Q/K/V can be supplied separately or in packed forms.

Key Features:

• Causal and padding masks.
• Bias and dropout.
• Sliding-Window Attention (SWA) and attention sinks.
• Experimental score_mod callbacks for FlexAttention-style customization.
• Support for different Q/K and V head dimensions (e.g., DeepSeek-style MLA).
• Context parallelism for long contexts using Ring or AllGather collectives (supported for selected BSHD and THD configurations).

The examples/jax/ directory provides introductory examples for using Transformer Engine with JAX, specifically focusing on FP8 training.

Key learning paths include:

• MNIST training: A starting point to learn how to use Transformer Engine and enable FP8 training.
• Encoder training: Examples demonstrating how to scale training across multiple GPUs using Transformer Engine.

When using TransformerEngine with JAX, the following conventions are used in the official documentation and examples:

• Framework: Uses Flax Linen. If using other stacks, refer to the Flax NNX/Linen or Haiku/Flax interop guides.
• Baseline dtype: Inputs and parameters are assumed to be bf16.
• Benchmarking: Performance is measured using quickstart_jax_utils.speedometer, which runs a JIT-compiled forward and backward loop with a warmup period.

When evaluating the performance impact of Transformer Engine's precision settings (BF16, FP8 Block, MXFP8, or NVFP4), note that several major operations are precision-agnostic. These operations run at the same speed regardless of the linear layer precision used:

• Attention (QK^T and softmax*V): Typically runs in BF16/FP16 via FlashAttention.
• LayerNorm / RMSNorm: Usually runs in FP32.
• Activation functions: Element-wise and memory-bound.
• AllReduce (DDP/FSDP): Communication-bound and independent of compute precision.

Only the linear projection GEMMs are directly affected by the precision setting in terms of compute performance.

The speedup achieved by using lower precision formats depends heavily on your model's matrix dimensions.

• Large GEMMs: High hidden/intermediate sizes and high token counts (micro_batch_size * sequence_length) help amortize the fixed quantization overhead, leading to meaningful speedups.
• Small GEMMs: Operations like attention output projection (where K=N=hidden_size with no expansion) may see little benefit or even a slowdown if the quantization overhead outweighs the kernel speedup.

It is recommended to benchmark with your actual model configuration to account for these shape-dependent effects before committing to a training run.

For sequence parallelism, Transformer Engine supports the all-gather of quantized tensors. This optimization targets the input and output gradient tensors. Using quantized all-gather provides three main benefits:

1. Reduced memory usage: Eliminates the need to store high-precision tensors for the backward pass.
2. Reduced communication: Smaller quantized tensors result in less data transferred across the network.
3. Parallelized quantization: Quantization workloads are distributed across multiple GPUs.

Note that support varies by recipe; for example, columnwise quantized all-gather is not available for all configurations. The actual behavior depends on your specific recipe and module configuration.

Nvidia-DL-Framework-Inspect works with Transformer Engine (TE) by utilizing hook calls embedded within TE GEMM (General Matrix Multiply) operations. The integration relies on three components:

1. TE Training: The active training process.
2. Feature Classes: User-defined or TE-provided classes that implement specific logic.
3. config.yaml: A configuration file that maps which hooks are used for specific layers.

Nvidia-DL-Framework-Inspect automates the insertion of these hooks into the correct locations. If a tensor is mentioned in config.yaml, the calls to modify_tensor_enabled() and modify_tensor() are substituted with definitions from the feature class. Other calls return default values and perform no action.

To train using low precisions like FP8, MXFP8, or NVFP4 in PyTorch, use the transformer_engine.pytorch.autocast context manager. This manager requires a recipe argument, which must be an object inheriting from transformer_engine.common.recipe.Recipe.

Key usage rules:

• Perform forward computations inside the autocast context.
• Call .backward() outside of the context (it inherits settings from the forward pass).
• Low precision training (FP8/etc.) requires SM89+ (Ada or newer) hardware.

You can apply different recipes to different parts of your model using Sequential contexts or override outer recipes using Nested contexts.

FP8 Blockwise Scaling assigns a dedicated FP32 scaling factor to each block of elements (128 for 1D, or 128x128 for 2D). This recipe is inspired by the DeepSeek-v3 quantization scheme.

Key Constraints:

• Hardware: Requires SM90 (Hopper) or later. On Blackwell, it is emulated with MXFP8.
• Tensor Dimensions: The tensor must have at least 2 dimensions. The last dimension must be divisible by 128, and the product of all dimensions except the last must be divisible by 128.
• Scaling Factors: By default, scaling factors are constrained to powers of 2. On Hopper, you can relax this using the NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 environment variable. On Blackwell, only powers of 2 are supported.
• JAX Support: Float8BlockScaling is not currently supported in JAX.

Depending on your model's sequence dimension distribution, choose one of the following tutorials:

1. Single-GPU Attention: Covers BSHD GQA + SWA, performance comparisons against native JAX baselines, and DeepSeek-style MLA head dimensions.
2. Context-Parallel Attention: Covers Packed THD GQA + SWA on four GPUs, including Ring and AllGather CP with striped load balancing, and performance comparisons against single-GPU fused attention.