NVIDIA/TransformerEngine

A library for accelerating Transformer models on NVIDIA GPUs, including using 8-bit and 4-bit floating point (FP8 and FP4) precision on Hopper, Ada and Blackwell GPUs, to provide better performance with lower memory utilization in both training and inference.

NVIDIA Transformer Engine

What it is – A library that speeds up Transformer‑style neural networks on NVIDIA GPUs by providing highly‑optimized kernels and an automatic mixed‑precision API. It lets you train and run large language models, mixture‑of‑experts (MoE) models, and multimodal Transformers using low‑precision formats such as FP8, MXFP8 and NVFP4, which reduce memory use and increase throughput while keeping accuracy comparable to FP16/BF16.

Key capabilities

  • FP8‑first support on Hopper, Ada and Blackwell GPUs, plus newer MXFP8/NVFP4 formats on Blackwell.
  • Framework‑agnostic C++ core plus thin Python bindings for PyTorch and JAX/Flax.
  • Fused kernels (e.g., FlashAttention‑2/‑3) that combine multiple operations into a single GPU launch for higher speed.
  • Automatic scaling‑factor handling so users can enable te.autocast and let the library manage the bookkeeping required for FP8 training.
  • Integration hooks for major LLM stacks: DeepSpeed, Hugging Face Accelerate, PyTorch Lightning, MosaicML Composer, etc.
  • Support for parallelism patterns (tensor, sequence, context) and MoE workloads.

Typical workflow (PyTorch example)

import torch, transformer_engine.pytorch as te
from transformer_engine.common import recipe

model = te.Linear(768, 3072, bias=True)
inp   = torch.randn(2048, 768, device='cuda')
fp8_recipe = recipe.DelayedScaling(margin=0, fp8_format=recipe.Format.E4M3)

with te.autocast(enabled=True, recipe=fp8_recipe):
    out = model(inp)
loss = out.sum(); loss.backward()

A similar pattern works in JAX/Flax, where te.autocast wraps the forward pass.

Installation

  • Docker (recommended): Pull the NVIDIA NGC containers (nvcr.io/nvidia/pytorch:26.01-py3 or nvcr.io/nvidia/jax:26.01-py3). The engine is pre‑installed inside /opt/transformerengine.
  • pip: pip install --no-build-isolation transformer_engine[pytorch] (or [jax] / both). A source install is possible with the usual CUDA, cuDNN, C++17 toolchain.
  • conda: conda install -c conda-forge transformer-engine-torch (JAX support forthcoming).

When to use it

  • Training LLMs or MoE models that would otherwise be limited by GPU memory or compute bandwidth.
  • Deploying inference pipelines where latency and memory footprint matter, especially on Hopper/Blackwell GPUs that expose FP8 hardware.
  • Projects already using PyTorch or JAX that want a drop‑in low‑precision upgrade without rewriting model code.

Limitations / gotchas

  • FP8 features require GPUs with compute capability 8.9+ (Ada/Hopper/Blackwell).
  • Building from source can be memory‑intensive (FlashAttention‑2 compilation) – set MAX_JOBS=1 if you hit OOM.
  • ABI mismatches between PyTorch and the engine can cause import errors; ensure both are built with the same C++ ABI.

Resources

Transformer Engine is an NVIDIA‑maintained open‑source project focused on accelerating modern Transformer workloads with cutting‑edge low‑precision hardware.

Related