NVIDIA/TransformerEngine
GitHub: NVIDIA/TransformerEngine
NVIDIA 推出的用于在 Hopper、Ada 和 Blackwell GPU 上通过 FP8/MXFP8/NVFP4 等低精度格式加速 Transformer 模型训练与推理的专用库。
Stars: 3435 | Forks: 776
..
版权所有 (c) 2022-2026，NVIDIA CORPORATION & AFFILIATES. 保留所有权利。
```
See LICENSE for license information.
```
|License|
# Transformer Engine
`快速入门 <#examples>`_ | `安装 <#installation>`_ | `用户指南 `_ | `示例 `_ | `收敛性 <#convergence>`_ | `集成 <#integrations>`_ | `发布说明 `_
# 最新新闻
* [12/2025] `NVIDIA Nemotron 3: 高效且开放的智能 `_ - 使用 NVFP4 在 Transformer Engine 上训练
* [11/2025] `NVIDIA Blackwell 架构横扫 MLPerf Training v5.1 基准测试 `_
* [09/2025] `Ling 2.0 的原生 FP8 混合精度训练，已开源！ `_
* [09/2025] `使用 NVIDIA NeMo 在 FP8 精度下实现更快的训练吞吐量 `_
* [08/2025] `NVFP4 以 16 位的精度和 4 位的速度与效率进行训练 `_。
.. overview-end-marker-do-not-remove
# 安装说明
## 系统要求
* **硬件：** Blackwell, Hopper, Grace Hopper/Blackwell, Ada, Ampere
* **操作系统：** Linux（官方支持），WSL2（有限支持）
* **软件：**
* CUDA：12.1+ (Hopper/Ada/Ampere)，12.8+ (Blackwell) 以及兼容的 NVIDIA 驱动程序
* cuDNN：9.3+
* 编译器：支持 C++17 的 GCC 9+ 或 Clang 10+
* Python：推荐 3.12
* **源码构建要求：** CMake 3.18+，Ninja，Git 2.17+，pybind11 2.6.0+
* **注意：** FP8 功能要求计算能力在 8.9 以上 (Ada/Hopper/Blackwell)
## 安装方法
## Docker (推荐)
开始使用 Transformer Engine 最快的方法是使用 `NVIDIA GPU Cloud (NGC) Catalog `_ 上的 Docker 镜像。
例如，要以交互方式使用 NGC PyTorch 容器，
.. code-block:: bash
```
docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:26.01-py3
```
例如，要以交互方式使用 NGC JAX 容器，
.. code-block:: bash
```
docker run --gpus all -it --rm nvcr.io/nvidia/jax:26.01-py3
```
这里的 26.01（对应于 2026 年 1 月发布）是容器版本。
我们建议更新到此处提供的最新 NGC 容器：
* https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch
* https://catalog.ngc.nvidia.com/orgs/nvidia/containers/jax
如果你运行任何示例，请确保使用匹配版本的 TransformerEngine。TransformerEngine 已预构建并打包在容器内，示例可在 ``/opt/transformerengine`` 或 ``/opt/transformer-engine`` 中找到。
**使用 NGC 容器的好处：**
* 所有依赖项均已预装，且版本兼容并经过优化配置
## pip 安装
**pip 安装的前提条件：**
* 兼容的 C++ 编译器
* 带有 cuDNN 和 NVCC (NVIDIA CUDA Compiler) 的 CUDA Toolkit（如果从源码安装）。
使用 pip 安装最新的稳定版本：
.. code-block:: bash
```
# 对于 PyTorch 集成
pip install --no-build-isolation transformer_engine[pytorch]
# 对于 JAX 集成
pip install --no-build-isolation transformer_engine[jax]
# 对于两种框架
pip install --no-build-isolation transformer_engine[pytorch,jax]
```
或者，直接从 GitHub 仓库安装：
.. code-block:: bash
```
pip install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@stable
```
从 GitHub 安装时，你可以使用环境变量显式指定框架：
.. code-block:: bash
```
NVTE_FRAMEWORK=pytorch,jax pip install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@stable
```
## conda 安装
使用 conda 从 conda-forge 安装最新的稳定版本：
.. code-block:: bash
```
# 对于 PyTorch 集成
conda install -c conda-forge transformer-engine-torch
# JAX 集成（即将推出）
```
## 源码安装
`查看安装指南 `_
## 环境变量
在安装之前可以设置这些环境变量以自定义构建过程：
* **CUDA_PATH**：CUDA 安装路径
* **CUDNN_PATH**：cuDNN 安装路径
* **CXX**：C++ 编译器路径
* **NVTE_FRAMEWORK**：要为其构建的框架的逗号分隔列表（例如 ``pytorch,jax``）
* **MAX_JOBS**：限制并行构建作业的数量（默认值因系统而异）
* **NVTE_BUILD_THREADS_PER_JOB**：控制每个构建作业的线程数
* **NVTE_CUDA_ARCHS**：要为其编译的 CUDA 计算架构的分号分隔列表（例如 ``80;90`` 对应 A100 和 H100）。如果未设置，将根据 CUDA 版本自动确定。设置此项可以显著减少构建时间和二进制文件大小。
## 使用 FlashAttention 编译
Transformer Engine 在 PyTorch 中同时支持 FlashAttention-2 和 FlashAttention-3 以提升性能。FlashAttention-3 在 v1.11 版本中加入，并在环境中同时存在两者时优先于 FlashAttention-2 使用。
你可以通过设置这些环境变量来验证正在使用的 FlashAttention 版本：
.. code-block:: bash
```
NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=1 python your_script.py
```
众所周知，FlashAttention-2 编译非常消耗资源，并且需要大量 RAM（参见 `bug `_），这可能导致在安装 Transformer Engine 期间出现内存溢出错误。请尝试在环境中设置 **MAX_JOBS=1** 以规避此问题。
.. troubleshooting-begin-marker-do-not-remove
## 故障排除
**常见问题与解决方案：**
1. **ABI 兼容性问题：**
* **症状：** 导入 transformer_engine 时出现带有未定义符号的 ``ImportError``
* **解决方案：** 确保 PyTorch 和 Transformer Engine 使用相同的 C++ ABI 设置构建。从源码重新构建 PyTorch 以匹配 ABI。
* **背景：** 如果你使用的 PyTorch 构建所用的 C++ ABI 与系统默认值不同，你可能会遇到这些未定义符号错误。这在容器外使用 pip 安装的 PyTorch 中尤为常见。
2. **缺少头文件或库：**
* **症状：** CMake 关于缺少头文件（``cudnn.h``, ``cublas_v2.h``, ``filesystem`` 等）的错误
* **解决方案：** 安装缺少的开发包或设置环境变量以指向正确的位置：
.. code-block:: bash
export CUDA_PATH=/path/to/cuda
export CUDNN_PATH=/path/to/cudnn
* 如果 CMake 找不到 C++ 编译器，请设置 ``CXX`` 环境变量。
* 确保在安装前正确设置了所有路径。
3. **构建资源问题：**
* **症状：** 编译挂起、系统卡死或内存溢出错误
* **解决方案：** 限制并行构建：
.. code-block:: bash
MAX_JOBS=1 NVTE_BUILD_THREADS_PER_JOB=1 pip install ...
4. **详细构建日志：**
* 获取详细的构建日志以帮助诊断问题：
.. code-block:: bash
cd transformer_engine
pip install -v -v -v --no-build-isolation .
**使用 UV 或虚拟环境的问题：**
1. **导入错误：**
* **症状：** 无法导入 ``transformer_engine``
* **解决方案：** 确保你的 UV 环境已激活，并且你使用了 ``uv pip install --no-build-isolation ``，而不是常规的 pip 安装到你的系统环境。
2. **cuDNN 子库加载失败：**
* **症状：** 运行时出现 ``CUDNN_STATUS_SUBLIBRARY_LOADING_FAILED`` 错误
* **解决方案：** 当 TE 是基于容器的系统级 cuDNN 安装进行构建，但虚拟环境内的 pip 包引入了 ``nvidia-cudnn-cu12/cu13`` 的 pip 包时，可能会发生这种情况。要解决此问题，从源码构建 TE 时，请指定以下环境变量以指向你虚拟环境中的 cuDNN。
.. code-block:: bash
export CUDNN_PATH=$(pwd)/.venv/lib/python3.12/site-packages/nvidia/cudnn
export CUDNN_HOME=$CUDNN_PATH
export LD_LIBRARY_PATH=$CUDNN_PATH/lib:$LD_LIBRARY_PATH
3. **构建 Wheels：**
* **症状：** 常规 TE 安装正常，但 UV wheel 构建在运行时失败。
* **解决方案：** 确保在 wheel 构建以及该 wheel 的 pip 安装过程中都使用了 ``uv build --wheel --no-build-isolation -v``。使用 ``-v`` 获取详细输出，以验证 TE 没有引入与 UV 环境版本不匹配的 PyTorch 或 JAX 版本。
**特定于 JAX 的常见问题与解决方案：**
1. **FFI 问题：**
* **症状：** ``No registered implementation for custom call to for platform CUDA``
* **解决方案：** 确保在安装期间使用了 ``--no-build-isolation``。如果预构建 wheels，请确保该 wheel 的构建和安装都使用了 ``--no-build-isolation``。如果使用 UV，请参阅上文的“使用 UV 或虚拟环境的问题”。
.. troubleshooting-end-marker-do-not-remove
# 重大变更
## v1.7：PyTorch 的 Padding mask 定义
为了统一 Transformer Engine 中所有三个框架的 attention mask 定义和用法，在我们的 PyTorch 实现中，padding mask 已从 `True` 表示包含对应位置于 attention 中，更改为表示排除该位置。自 v1.7 起，所有 attention mask 类型均遵循相同的定义：`True` 表示屏蔽对应位置，而 `False` 表示在 attention 计算中包含该位置。
此更改的示例为，
.. code-block:: bash
```
# 对于一个包含 3 个 sequences 的 batch，其中 `a`、`b` 和 `c` 是有用的 tokens
# 并且 `0` 是 padding tokens，
[a, a, a, 0, 0,
b, b, 0, 0, 0,
c, c, c, c, 0]
# 此 batch 在 v1.7 之前的 padding mask 为，
[ True, True, True, False, False,
True, True, False, False, False,
True, True, True, True, False]
# 而对于 v1.7 及更高版本，它应该是，
[False, False, False, True, True,
False, False, True, True, True,
False, False, False, False, True]
```
# 收敛性
FP8 和 MXFP8 已在不同的模型架构和配置中进行了广泛测试，我们发现 FP8/MXFP8 和 BF16 训练损失曲线之间**没有显著差异**。FP8 和 MXFP8 的准确性也在下游 LLM 任务（例如 LAMBADA 和 WikiText）中得到了验证。以下是跨不同框架测试收敛性的模型示例。
+------------+------------------+---------------------------------------------------------------------------------------------------------+
| 模型 | 框架 | 来源 |
+============+==================+=========================================================================================================+
| MPT-1.3B | Mosaic Composer | https://www.databricks.com/blog/coreweave-nvidia-h100-part-1 |
+------------+------------------+---------------------------------------------------------------------------------------------------------+
| LLama2-7B | Alibaba Pai | https://mp.weixin.qq.com/s/NQT0uKXLbXyh5031zBdeBQ |
+------------+------------------+---------------------------------------------------------------------------------------------------------+
| LLM-8B | Megatron Core | https://arxiv.org/abs/2506.08027 |
+------------+------------------+---------------------------------------------------------------------------------------------------------+
| MPT-13B | Mosaic Composer | https://www.databricks.com/blog/turbocharged-training-optimizing-databricks-mosaic-ai-stack-fp8 |
+------------+------------------+---------------------------------------------------------------------------------------------------------+
| MoE-16B | Megatron Core | https://arxiv.org/abs/2506.08027 |
+------------+------------------+---------------------------------------------------------------------------------------------------------+
| LLama2-70B | Alibaba Pai | https://mp.weixin.qq.com/s/NQT0uKXLbXyh5031zBdeBQ |
+------------+------------------+---------------------------------------------------------------------------------------------------------+
# 集成
Transformer Engine 已与流行的 LLM 框架集成，例如：
* `DeepSpeed `_
* `Hugging Face Accelerate `_
* `Lightning `_
* `MosaicML Composer `_
* `NVIDIA JAX Toolbox `_
* `NVIDIA Megatron-LM `_
* `NVIDIA NeMo Megatron Bridge `_
* `Amazon SageMaker Model Parallel Library `_
* `Levanter `_
* `GPT-NeoX `_
* `Hugging Face Nanotron `_
# 贡献
我们欢迎对 Transformer Engine 的贡献！要为 Transformer Engine 做出贡献并提交 pull request，
请遵循 ``_ 指南中概述的准则。
# 论文
* `Attention 原始论文 `_
* `Megatron-LM 张量并行 `_
* `Megatron-LM 序列并行 `_
* `用于深度学习的 FP8 格式 `_
# 视频
* `Blackwell 上稳定且可扩展的 FP8 深度学习训练 | GTC 2025 `_
* `从 FP8 LLM 训练到推理：大规模语言 AI | GTC 2025 `_
* `用于深度学习的 FP8 | GTC 2023 `_
* [05/2024] `使用 NVIDIA cuDNN 9 加速 Transformer `_
* [03/2024] `SageMaker 模型并行库中的 FP8 训练支持 `_
* [12/2023] `全新 NVIDIA NeMo 框架特性与 NVIDIA H200 `_
* [11/2023] `使用 NVIDIA Transformer Engine 释放 Transformer 的力量 `_
* [11/2023] `使用 FP8 加速 PyTorch 训练工作负载 `_
* [09/2023] `Transformer Engine 已添加到用于 PyTorch 训练的 AWS DL Container 中 `_
* [06/2023] `使用 NVIDIA H100 GPU 打破 MLPerf 训练记录 `_
标签：FP8, NVIDIA GPU, Transformer, Vectored Exception Handling, 凭据扫描, 推理优化, 模型加速, 深度学习, 混合精度训练, 逆向工具