Quickstart | Installation | User Guide | Examples | FP8 Convergence | Integrations | Release notes
Transformer Engine (TE) is a library for accelerating Transformer models on NVIDIA GPUs, including using 8-bit floating point (FP8) precision on Hopper, Ada, and Blackwell GPUs, to provide better performance with lower memory utilization in both training and inference. TE provides a collection of highly optimized building blocks for popular Transformer architectures and an automatic mixed precision-like API that can be used seamlessly with your framework-specific code. TE also includes a framework agnostic C++ API that can be integrated with other deep learning libraries to enable FP8 support for Transformers.
As the number of parameters in Transformer models continues to grow, training and inference for architectures such as BERT, GPT and T5 become very memory and compute-intensive. Most deep learning frameworks train with FP32 by default. This is not essential, however, to achieve full accuracy for many deep learning models. Using mixed-precision training, which combines single-precision (FP32) with lower precision (e.g. FP16) format when training a model, results in significant speedups with minimal differences in accuracy as compared to FP32 training. With Hopper GPU architecture FP8 precision was introduced, which offers improved performance over FP16 with no degradation in accuracy. Although all major deep learning frameworks support FP16, FP8 support is not available natively in frameworks today.
TE addresses the problem of FP8 support by providing APIs that integrate with popular Large Language Model (LLM) libraries. It provides a Python API consisting of modules to easily build a Transformer layer as well as a framework-agnostic library in C++ including structs and kernels needed for FP8 support. Modules provided by TE internally maintain scaling factors and other values needed for FP8 training, greatly simplifying mixed precision training for users.
import torch
import transformer_engine.pytorch as te
from transformer_engine.common import recipe
# Set dimensions.
in_features = 768
out_features = 3072
hidden_size = 2048
# Initialize model and inputs.
model = te.Linear(in_features, out_features, bias=True)
inp = torch.randn(hidden_size, in_features, device="cuda")
# Create an FP8 recipe. Note: All input args are optional.
fp8_recipe = recipe.DelayedScaling(margin=0, fp8_format=recipe.Format.E4M3)
# Enable autocasting for the forward pass
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
out = model(inp)
loss = out.sum()
loss.backward()
import flax
import jax
import jax.numpy as jnp
import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax
from transformer_engine.common import recipe
BATCH = 32
SEQLEN = 128
HIDDEN = 1024
# Initialize RNG and inputs.
rng = jax.random.PRNGKey(0)
init_rng, data_rng = jax.random.split(rng)
inp = jax.random.normal(data_rng, [BATCH, SEQLEN, HIDDEN], jnp.float32)
# Create an FP8 recipe. Note: All input args are optional.
fp8_recipe = recipe.DelayedScaling(margin=0, fp8_format=recipe.Format.HYBRID)
# Enable autocasting for the forward pass
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
model = te_flax.DenseGeneral(features=HIDDEN)
def loss_fn(params, other_vars, inp):
out = model.apply({'params':params, **other_vars}, inp)
return jnp.mean(out)
# Initialize models.
variables = model.init(init_rng, inp)
other_variables, params = flax.core.pop(variables, 'params')
# Construct the forward and backward function
fwd_bwd_fn = jax.value_and_grad(loss_fn, argnums=(0, 1))
for _ in range(10):
loss, (param_grads, other_grads) = fwd_bwd_fn(params, other_variables, inp)
For a more comprehensive tutorial, check out our Quickstart Notebook.
The quickest way to get started with Transformer Engine is by using Docker images on NVIDIA GPU Cloud (NGC) Catalog.
For example to use the NGC PyTorch container interactively,
docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:25.04-py3
For example to use the NGC JAX container interactively,
docker run --gpus all -it --rm nvcr.io/nvidia/jax:25.04-py3
Where 25.04 (corresponding to April 2025 release) is the container version.
Benefits of using NGC containers:
Prerequisites for pip installation:
To install the latest stable version with pip:
# For PyTorch integration pip install --no-build-isolation transformer_engine[pytorch] # For JAX integration pip install --no-build-isolation transformer_engine[jax] # For both frameworks pip install --no-build-isolation transformer_engine[pytorch,jax]
Alternatively, install directly from the GitHub repository:
pip install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@stable
When installing from GitHub, you can explicitly specify frameworks using the environment variable:
NVTE_FRAMEWORK=pytorch,jax pip install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@stable
To install the latest stable version with conda from conda-forge:
# For PyTorch integration conda install -c conda-forge transformer-engine-torch # JAX integration (coming soon)
These environment variables can be set before installation to customize the build process:
pytorch,jax)Transformer Engine supports both FlashAttention-2 and FlashAttention-3 in PyTorch for improved performance. FlashAttention-3 was added in release v1.11 and is prioritized over FlashAttention-2 when both are present in the environment.
You can verify which FlashAttention version is being used by setting these environment variables:
NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=1 python your_script.py
It is a known issue that FlashAttention-2 compilation is resource-intensive and requires a large amount of RAM (see bug), which may lead to out of memory errors during the installation of Transformer Engine. Please try setting MAX_JOBS=1 in the environment to circumvent the issue.
Common Issues and Solutions:
ABI Compatibility Issues:
ImportError with undefined symbols when importing transformer_engineMissing Headers or Libraries:
Symptoms: CMake errors about missing headers (cudnn.h, cublas_v2.h, filesystem, etc.)
Solution: Install missing development packages or set environment variables to point to correct locations:
export CUDA_PATH=/path/to/cuda export CUDNN_PATH=/path/to/cudnn
If CMake can't find a C++ compiler, set the CXX environment variable.
Ensure all paths are correctly set before installation.
Build Resource Issues:
Symptoms: Compilation hangs, system freezes, or out-of-memory errors
Solution: Limit parallel builds:
MAX_JOBS=1 NVTE_BUILD_THREADS_PER_JOB=1 pip install ...
Verbose Build Logging:
For detailed build logs to help diagnose issues:
cd transformer_engine pip install -v -v -v --no-build-isolation .
In an effort to unify the definition and usage of the attention mask across all three frameworks in Transformer Engine, the padding mask has changed from True meaning inclusion of the corresponding position in attention to exclusion of that position in our PyTorch implementation. Since v1.7, all attention mask types follow the same definition where True means masking out the corresponding position and False means including that position in attention calculation.
An example of this change is,
# for a batch of 3 sequences where `a`s, `b`s and `c`s are the useful tokens # and `0`s are the padding tokens, [a, a, a, 0, 0, b, b, 0, 0, 0, c, c, c, c, 0] # the padding mask for this batch before v1.7 is, [ True, True, True, False, False, True, True, False, False, False, True, True, True, True, False] # and for v1.7 onwards it should be, [False, False, False, True, True, False, False, True, True, True, False, False, False, False, True]
FP8 has been tested extensively across different model architectures and configurations and we found no significant difference between FP8 and BF16 training loss curves. FP8 has also been validated for accuracy on downstream LLM tasks (e.g. LAMBADA and WikiText). Below are examples of models tested for convergence across different frameworks.
Transformer Engine has been integrated with popular LLM frameworks such as:
We welcome contributions to Transformer Engine! To contribute to Transformer Engine and make pull requests, follow the guidelines outlined in the CONTRIBUTING.rst guide.
RetroSearch is an open source project built by @garambo | Open a GitHub Issue
Search and Browse the WWW like it's 1997 | Search results from DuckDuckGo
HTML:
3.2
| Encoding:
UTF-8
| Version:
0.7.4
