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

功能簡介 – 透過提供高度優化的核心與自動混合精度 API，加速 NVIDIA GPU 上的 Transformer 風格神經網路。它允許您使用 FP8、MXFP8 和 NVFP4 等低精度格式來訓練與執行大語言模型 (LLM)、混合專家 (MoE) 模型及多模態 Transformer，在保持與 FP16/BF16 相當的精度的同時，減少記憶體使用並提高吞吐量。

核心能力

FP8 優先支援：在 Hopper、Ada 及 Blackwell GPU 上提供支援，並在 Blackwell 上支援更新的 MXFP8/NVFP4 格式。
框架無關的 C++ 核心：為 PyTorch 及 JAX/Flax 提供輕量級 Python 綁定。
融合核心 (Fused kernels)：例如 FlashAttention-2/-3，將多個操作合併為單個 GPU 啟動以實現更高速度。
自動縮放因子處理：使用者只需啟用 te.autocast，即可讓函式庫管理 FP8 訓練所需的帳務工作。
主要 LLM 技術棧整合：支援 DeepSpeed、Hugging Face Accelerate、PyTorch Lightning、MosaicML Composer 等。
並行模式支援：支援張量並行 (tensor)、序列並行 (sequence)、上下文並行 (context) 以及 MoE 工作負載。

典型工作流程 (PyTorch 範例)

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()

在 JAX/Flax 中也有類似的模式，其中 te.autocast 包裹前向傳播。

安裝

Docker (建議): 拉取 NVIDIA NGC 容器 (nvcr.io/nvidia/pytorch:26.01-py3 或 nvcr.io/nvidia/jax:26.01-py3)。引擎已預裝在 /opt/transformerengine 中。
pip: pip install --no-build-isolation transformer_engine[pytorch] (或 [jax] / 兩者)。可以使用常規的 CUDA、cuDNN、C++17 工具鏈進行原始碼安裝。
conda: conda install -c conda-forge transformer-engine-torch (JAX 支援即將推出)。

使用時機

訓練原本會受到 GPU 記憶體或計算頻寬限制的 LLM 或 MoE 模型。
部署對延遲與記憶體佔用敏感的推論流水線，特別是在提供 FP8 硬體的 Hopper/Blackwell GPU 上。
已經在使用 PyTorch 或 JAX，並希望在不重寫模型程式碼的情況下直接升級到低精度格式的專案。

限制 / 注意事項

FP8 功能需要具備 8.9+ 計算能力的 GPU (Ada/Hopper/Blackwell)。
從原始碼建置可能會消耗大量記憶體 (FlashAttention-2 編譯) —— 如果遇到 OOM，請設定 MAX_JOBS=1。
PyTorch 與引擎之間的 ABI 不匹配可能會導致匯入錯誤；請確保兩者均使用相同的 C++ ABI 建置。

資源

完整使用者指南: https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/index.html
範例 Notebook: https://github.com/NVIDIA/TransformerEngine/tree/main/examples
最新新聞 (2026年6月): 關於 MoE 吞吐量、NVFP4 及 Nemotron-3 Ultra 模型。

Transformer Engine 是一個由 NVIDIA 維護的開源專案，專注於利用尖端低精度硬體加速現代 Transformer 工作負載。

NVIDIA/TransformerEngine

NVIDIA Transformer Engine

相關