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) 解決了將 Transformer 模型（例如 LLM、MoE 和多模態模型）擴展到數千億個參數時所帶來的高記憶體與運算需求。它透過利用低精度數值格式，在訓練與推論過程中皆能降低記憶體利用率並提高吞吐量。

How it works

TE 提供了一套針對 Transformer 架構高度優化的構建模組與融合核心（fused kernels）。它實作了自動混合精度 API，讓使用者能夠無縫地將低精度格式整合到現有的框架代碼中。

關鍵技術能力包括：

  • Low-Precision Support: 在 Hopper、Ada 和 Blackwell GPU 上原生支持 FP8，並透過 Blackwell GPU 上的 MXFP8 和 NVFP4 進一步提升效率。
  • Framework Integration: 提供 PyTorch 和 JAX (Flax) 的 Python API，以及用於與其他深度學習函式庫整合的框架無關的 C++ API。
  • Automated Scaling: 內部管理 FP8 訓練所需的縮放因子（scaling factors），簡化了使用者的混合精度流程。
  • Optimized Kernels: 包含針對 Mixture-of-Experts (MoE) 等進階功能以及各種並行策略（tensor、sequence 和 context）的融合操作與優化。

Who it’s for

它是為使用 NVIDIA GPU（Ampere 架構或更新版本）並希望最大化硬體效率與訓練速度的大規模 Transformer 模型開發者與工程師所設計的。

Highlights

  • Broad Precision Support: 在最新的 NVIDIA GPU 上支持 FP8、MXFP8 和 NVFP4，並在 Ampere 及之後的版本上支持 FP16/BF16。
  • Easy Integration: 提供用於構建具有 FP8 支持的 Transformer 層的現成模組。
  • C++ API: 包含一個用於深度學習函式庫開發者的框架無關的 C++ 函式庫。
  • Performance Optimizations: 具備融合核心與對 PyTorch 中 FlashAttention-2 和 FlashAttention-3 的支持。