安装

先决条件

  1. Linux x86_64

  2. CUDA 11.8

  3. 支持 CUDA 11.8 或更高版本的NVIDIA 驱动程序。

  4. cuDNN 8.1 或更高版本。

  5. 对于 FP8/FP16/BF16 融合注意力，需要 CUDA 12.1 或更高版本、支持 CUDA 12.1 或更高版本的NVIDIA 驱动程序以及 cuDNN 8.9.1 或更高版本。

如果在标准安装路径（例如 CUDA_HOME）中运行时无法找到 CUDA Toolkit 头文件，请在环境中设置 NVTE_CUDA_INCLUDE_PATH。

NGC 容器中的 Transformer Engine

Transformer Engine 库已预安装在 NVIDIA GPU Cloud 上 22.09 及更高版本的 PyTorch 容器中。

pip - 从 GitHub

附加先决条件

  1. [针对 PyTorch 支持] 带有 GPU 支持的 PyTorch。

  2. [针对 JAX 支持] 带有 GPU 支持的 JAX，版本 >= 0.4.7。

安装 (稳定版本)

执行以下命令安装最新稳定版本的 Transformer Engine

pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable

这将自动检测是否安装了任何支持的深度学习框架，并为其构建 Transformer Engine 支持。要显式指定框架，请将环境变量 NVTE_FRAMEWORK 设置为逗号分隔列表（例如 NVTE_FRAMEWORK=jax,pytorch）。

安装 (开发版本)

警告

Transformer Engine 的开发版本可能包含官方版本尚未提供的新功能，但它不受支持，因此不建议用于一般用途。

执行以下命令安装最新开发版本的 Transformer Engine

pip install git+https://github.com/NVIDIA/TransformerEngine.git@main

这将自动检测是否安装了任何支持的深度学习框架，并为其构建 Transformer Engine 支持。要显式指定框架，请将环境变量 NVTE_FRAMEWORK 设置为逗号分隔列表（例如 NVTE_FRAMEWORK=jax,pytorch）。

要安装特定的 PR，请将 NNN 更改为 PR 编号后执行

pip install git+https://github.com/NVIDIA/TransformerEngine.git@refs/pull/NNN/merge

安装 (从源代码)

执行以下命令从源代码安装 Transformer Engine

# Clone repository, checkout stable branch, clone submodules
git clone --branch stable --recursive https://github.com/NVIDIA/TransformerEngine.git

cd TransformerEngine
export NVTE_FRAMEWORK=pytorch   # Optionally set framework
pip install .                   # Build and install

如果已经克隆了 Git 仓库，请确保同时也克隆子模块

git submodule update --init --recursive

可以通过设置“test”选项来安装额外的测试依赖项

pip install .[test]

要使用调试符号构建 C++ 扩展（例如使用 -g 标志）

pip install . --global-option=--debug