Transformer Engine

Quick start | User guide | PyTorch API | JAX API | Examples | Releases

What is Transformer Engine?

Transformer Engine (TE) is an NVIDIA library for accelerating Transformer model training on NVIDIA GPUs. It combines optimized building blocks and fused kernels with automatic mixed-precision-style APIs for PyTorch and JAX, so low-precision training can be adopted without rewriting a training stack.

Transformer Engine manages the scaling factors, amax histories, and quantization metadata required by low-precision recipes. Its modules cover attention, linear layers, normalization, Mixture-of-Experts (MoE), and communication operations used in large-scale distributed training.

Highlights

• FP8 training on NVIDIA Hopper, Ada, Blackwell, and Rubin GPUs.

• MXFP8 and NVFP4 training on NVIDIA Blackwell GPUs.

• Optimized attention, GEMM, normalization, quantization, and fused Transformer and MoE modules.

• PyTorch and JAX APIs with autocast-style contexts and configurable low-precision recipes.

• Support for tensor, sequence, context, and EP, including communication overlap.

• FP16 and BF16 optimizations on NVIDIA Ampere architecture GPUs and later.

News

• [09/2026] Transformer Engine v2.19 adds Rubin support, hybrid quantization, MXFP8 EP communication, and expanded FP8 attention support.

• [09/2026] Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine describes optimized JAX MoE training on GB200 and GB300 systems.

• [08/2026] Transformer Engine v2.18 adds FP8 block scaling in PyTorch, zero-copy EP, and CUDA Graph support for THD attention.

• [07/2026] Transformer Engine v2.17 introduces NCCL-backed EP for PyTorch and JAX, faster MoE routing, and a JAX Flax MoE block.

• [06/2026] Boosting MoE Training Throughput with Advanced Fusion Kernels shows how fused Transformer Engine operations accelerate MoE training.

See the project updates archive for earlier news.

Quick start

Install the latest stable release for your framework:

# PyTorch
pip install --no-build-isolation "transformer_engine[pytorch]"

# JAX
pip install --no-build-isolation "transformer_engine[jax]"

For a ready-to-run environment, use an NVIDIA NGC framework container. Replace <YY.MM> with a container release listed in the NVIDIA Deep Learning Frameworks Support Matrix.

docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:<YY.MM>-py3
docker run --gpus all -it --rm nvcr.io/nvidia/jax:<YY.MM>-py3

Continue with the PyTorch and JAX getting started guide. For prerequisites, source builds, environment variables, and troubleshooting, see the installation guide.

Integrations

Transformer Engine has been integrated with popular LLM frameworks such as:

• Hugging Face Accelerate

• Lightning

• NVIDIA BioNeMo Recipes

• NVIDIA JAX Toolbox

• NVIDIA Megatron-LM

• NVIDIA NeMo Megatron Bridge

• Amazon SageMaker Model Parallel Library

See Ecosystem and historical integrations for additional community integrations and projects that have worked with Transformer Engine.

Contributing

We welcome contributions to Transformer Engine! To contribute to Transformer Engine and make pull requests, follow the guidelines outlined in the CONTRIBUTING.rst guide.

Technical deep dives

• Low precision training guide — FP8, MXFP8, NVFP4, scaling recipes, and performance considerations.

• Using FP8 and FP4 with Transformer Engine — a practical introduction with code examples.

• FP8 Formats for Deep Learning — the foundational paper describing the FP8 formats used for deep learning.

• Stable and Scalable FP8 Deep Learning Training on Blackwell — a technical GTC 2025 session on training numerics and scale.

See Resources for the complete collection of papers and recorded talks.