Transformer Engine

Transformer Engine is a library for developers building and running Transformer models on NVIDIA GPUs. It reduces compute and memory demands in training and inference through low-precision arithmetic and optimized model components.

• Low-precision training provides PyTorch and JAX APIs that manage the scaling data needed for mixed-precision training within an existing training stack.
• Transformer components cover attention, linear layers, normalization and Mixture-of-Experts (MoE), with fused operations that reduce execution overhead.
• Distributed training supports tensor, sequence and context parallelism, plus communication overlap for workloads spread across GPUs.
• Framework integration includes Hugging Face Accelerate, Lightning and NVIDIA Megatron-LM; a C++ API lets other deep learning libraries use its FP8 capabilities.

GPU architecture determines which precision formats you can use. FP8 works on NVIDIA Hopper, Ada and Blackwell GPUs. Blackwell also supports MXFP8 and NVFP4, while FP16 and BF16 optimizations work on Ampere GPUs and later.

Transformer Engine is open source under the Apache 2.0 license. Additional integrations include NVIDIA BioNeMo Recipes, NVIDIA JAX Toolbox and NVIDIA NeMo Megatron Bridge.