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 构建。

资源

Transformer Engine 是一个由 NVIDIA 维护的开源项目，专注于利用尖端低精度硬件加速现代 Transformer 工作负载。

相关