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.

What it solves

Transformer Engine (TE) addresses the high memory and compute demands of scaling Transformer models (such as LLMs, MoE, and multimodal models) to hundreds of billions of parameters. It reduces memory utilization and increases throughput during both training and inference by leveraging low-precision numerical formats.

How it works

TE provides a library of highly optimized building blocks and fused kernels for Transformer architectures. It implements an automatic mixed-precision API that allows users to seamlessly integrate low-precision formats into their existing framework code.

Key technical capabilities include:

  • Low-Precision Support: Native support for 8-bit floating point (FP8) on Hopper, Ada, and Blackwell GPUs, and further efficiency gains via MXFP8 and NVFP4 on Blackwell GPUs.
  • Framework Integration: Provides Python APIs for PyTorch and JAX (Flax), as well as a framework-agnostic C++ API for integration with other deep learning libraries.
  • Automated Scaling: Internally manages scaling factors required for FP8 training, simplifying the mixed-precision process for the user.
  • Optimized Kernels: Includes fused operations and optimizations for advanced features like Mixture-of-Experts (MoE) and various parallelism strategies (tensor, sequence, and context).

Who it’s for

It is designed for researchers and engineers building and training large-scale Transformer models who use NVIDIA GPUs (Ampere architecture or later) and want to maximize hardware efficiency and training speed.

Highlights

  • Broad Precision Support: Supports FP8, MXFP8, and NVFP4 on the latest NVIDIA GPUs, and FP16/BF16 on Ampere and later.
  • Easy Integration: Offers ready-to-use modules for building Transformer layers with FP8 support.
  • C++ API: Includes a framework-agnostic library for deep learning library developers.
  • Performance Optimizations: Features fused kernels and support for FlashAttention-2 and FlashAttention-3 in PyTorch.