.. 版权所有 © 2022-2026，NVIDIA公司及其关联公司。保留所有权利。

有关许可信息，请参阅 LICENSE。

|许可证|

Transformer Engine

快速入门 <#examples>_ | 安装 <#installation>_ | 用户指南 <https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/index.html>_ | 示例 <https://github.com/NVIDIA/TransformerEngine/tree/main/examples>_ | 收敛性 <#convergence>_ | 集成 <#integrations>_ | 发布说明 <https://docs.nvidia.com/deeplearning/transformer-engine/documentation-archive.html>_

最新消息

• [12/2025] NVIDIA Nemotron 3：高效开放的智能 <https://arxiv.org/abs/2512.20856>_ - 使用 Transformer Engine 上的 NVFP4 训练
• [11/2025] NVIDIA Blackwell 架构横扫 MLPerf 训练 v5.1 基准测试 <https://developer.nvidia.com/blog/nvidia-blackwell-architecture-sweeps-mlperf-training-v5-1-benchmarks/>_
• [11/2025] 使用 PyTorch 和 NVIDIA BioNeMo 配方扩展生物学 Transformer 模型 <https://developer.nvidia.com/blog/scale-biology-transformer-models-with-pytorch-and-nvidia-bionemo-recipes/>_
• [11/2025] 大规模 RL 模型的 FP8 训练 <https://lmsys.org/blog/2025-11-25-fp8-rl/>_
• [09/2025] 使用 NVFP4 预训练大型语言模型 <https://www.arxiv.org/pdf/2509.25149>_
• [09/2025] Ling 2.0 的原生 FP8 混合精度训练已开源！<https://huggingface.co/blog/im0qianqian/ling-mini-2-fp8-mixed-precision-training-solution>_
• [09/2025] 借助 NVIDIA NeMo，在 FP8 精度下实现更快的训练吞吐量 <https://developer.nvidia.com/blog/faster-training-throughput-in-fp8-precision-with-nvidia-nemo/>_
• [08/2025] 我们如何使用 FP8 进行训练和推理来构建 DeepL 的下一代 LLM <https://www.deepl.com/en/blog/tech/next-generation-llm-fp8-training>_
• [08/2025] NVFP4 以 16 位精度进行训练，同时兼具 4 位的高速与高效 <https://developer.nvidia.com/blog/nvfp4-trains-with-precision-of-16-bit-and-speed-and-efficiency-of-4-bit/>_

往期新闻 <#previous-news>_

什么是 Transformer Engine？

.. overview-begin-marker-do-not-remove

Transformer Engine (TE) 是一个用于在 NVIDIA GPU 上加速 Transformer 模型的库，包括在 Hopper、Ada 和 Blackwell GPU 上使用 8 位浮点数 (FP8) 精度，从而在训练和推理中提供更好的性能以及更低的内存占用。在 Blackwell GPU 上，TE 还支持 MXFP8（微缩放 FP8）和 NVFP4 格式，以实现更高的效率。TE 提供了一系列针对流行 Transformer 架构的高度优化的构建模块，以及一个类似自动混合精度的 API，可以与您特定框架的代码无缝集成。此外，TE 还包含一个框架无关的 C++ API，可与其他深度学习库集成，以启用 Transformer 的 FP8 支持。

随着 Transformer 模型参数数量的持续增长，诸如 BERT、GPT 和 T5 等架构的训练和推理变得非常消耗内存和计算资源。大多数深度学习框架默认使用 FP32 进行训练。然而，对于许多深度学习模型来说，这并不是实现完全准确性的必要条件。通过混合精度训练，在训练模型时将单精度 (FP32) 与较低精度格式（例如 FP16）相结合，可以在与 FP32 训练相比几乎没有精度损失的情况下显著提高速度。随着 Hopper GPU 架构的推出，FP8 精度被引入，它在不降低精度的情况下提供了比 FP16 更好的性能。尽管所有主要的深度学习框架都支持 FP16，但目前框架本身并不原生支持 FP8。

TE 通过提供与流行的大型语言模型 (LLM) 库集成的 API 来解决 FP8 支持的问题。它提供了一个由模块组成的 Python API，可以轻松构建 Transformer 层；同时还提供了一个框架无关的 C++ 库，其中包含用于 FP8 支持所需的结构体和内核。TE 提供的模块会在内部维护 FP8 训练所需的缩放因子和其他值，从而大大简化了用户的混合精度训练过程。

亮点

• 易于使用的模块，用于构建支持 FP8 的 Transformer 层
• 针对 Transformer 模型的优化（例如融合内核）
• 支持 NVIDIA Hopper、Ada 和 Blackwell GPU 上的 FP8
• 支持 NVIDIA Blackwell GPU 上的 MXFP8 和 NVFP4
• 支持 NVIDIA Ampere GPU 架构及更高版本上所有精度（FP16、BF16）的优化

示例

PyTorch ^^^^^^^

.. code-block:: python

import torch import transformer_engine.pytorch as te from transformer_engine.common import recipe

设置维度。

in_features = 768 out_features = 3072 hidden_size = 2048

初始化模型和输入。

model = te.Linear(in_features, out_features, bias=True) inp = torch.randn(hidden_size, in_features, device="cuda")

创建 FP8 配方。注意：所有输入参数都是可选的。

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

.. code-block:: python

  import flax
  import jax
  import jax.numpy as jnp
  import transformer_engine.jax as te
  import transformer_engine.jax.flax as te_flax
  from transformer_engine.common import recipe

  BATCH = 32
  SEQLEN = 128
  HIDDEN = 1024

  # 初始化随机数生成器和输入。
  rng = jax.random.PRNGKey(0)
  init_rng，data_rng = jax.random.split(rng)
  inp = jax.random.normal(data_rng，[BATCH，SEQLEN，HIDDEN], jnp.float32)

  # 创建 FP8 配方。注意：所有输入参数都是可选的。
  fp8_recipe = recipe.DelayedScaling(margin=0，fp8_format=recipe.Format.HYBRID)

  # 启用前向传播的自动转换
  with te.autocast(enabled=True，recipe=fp8_recipe):
      model = te_flax.DenseGeneral(features=HIDDEN)

      def loss_fn(params，other_vars，inp):
        out = model.apply({'params':params，**other_vars}, inp)
        return jnp.mean(out)

      # 初始化模型。
      variables = model.init(init_rng，inp)
      other_variables，params = flax.core.pop(variables，'params')

      # 构建前向和反向函数
      fwd_bwd_fn = jax.value_and_grad(loss_fn，argnums=(0，1))

      for _ in range(10):
        loss，(param_grads，other_grads) = fwd_bwd_fn(params，other_variables，inp)

如需更全面的教程，请查看我们的《入门指南》：<https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/getting_started.html>。

.. 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) 目录 <https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch>`_ 上使用 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 编译器）的 CUDA 工具包。

要使用 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-forge 安装最新稳定版：

.. code-block:: bash

    # 用于 PyTorch 集成
    conda install -c conda-forge transformer-engine-torch
    
    # JAX 集成（即将推出）

源码安装
^^^^^^^^^^^^^^^^^^^

`请参阅安装指南 <https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/installation.html#installation-from-source>`_

环境变量
^^^^^^^^^^^^^^^^^^^^^
这些环境变量可以在安装前设置，以自定义构建过程：

* **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-3。

您可以通过设置以下环境变量来验证正在使用的 FlashAttention 版本：

.. code-block:: bash

    NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=1 python your_script.py

众所周知，FlashAttention-2 的编译非常消耗资源，需要大量内存（参见 `bug <https://github.com/Dao-AILab/flash-attention/issues/358>`_），这可能导致在安装 Transformer Engine 时出现内存不足错误。请尝试在环境中设置 **MAX_JOBS=1** 来规避此问题。

.. troubleshooting-begin-marker-do-not-remove

故障排除
^^^^^^^^^^^^^^^

**常见问题及解决方案**：

1. **ABI 兼容性问题**：

   * **症状**：导入 transformer_engine 时出现未定义符号的 ``ImportError``
   * **解决方案**：确保 PyTorch 和 Transformer Engine 使用相同的 C++ ABI 设置进行构建。请使用匹配的 ABI 从源码重新构建 PyTorch。
   * **背景**：如果您使用的是与系统默认 C++ ABI 不同的 PyTorch 构建版本，则可能会遇到这些未定义符号错误。这种情况在通过 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 <te_pypi_package_or_wheel_or_source_dir>``，而不是常规的 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. **构建轮子包：**

   * **症状：** 使用常规方法安装 TE 可以正常工作，但使用 UV 构建的轮子包在运行时会失败。
   * **解决方案：** 确保在构建轮子包以及通过 pip 安装该轮子包时都使用 ``uv build --wheel --no-build-isolation -v``。使用 ``-v`` 参数以获取详细输出，从而验证 TE 是否引入了与 UV 环境版本不匹配的 PyTorch 或 JAX 版本。

**JAX 特有的常见问题及解决方案：**

1. **FFI 问题：**

   * **症状：** ``未为平台 CUDA 注册 <some_te_ffi> 自定义调用的实现``
   * **解决方案：** 确保在安装过程中使用 ``--no-build-isolation``。如果预先构建轮子包，确保轮子包的构建和安装都使用 ``--no-build-isolation``。如果使用 UV，请参阅上文“使用 UV 或虚拟环境时的问题”。

.. troubleshooting-end-marker-do-not-remove

重大变更
================

v1.7：PyTorch 中的填充掩码定义
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
为了在 Transformer Engine 的三个框架中统一注意力掩码的定义和使用方式，我们的 PyTorch 实现中，填充掩码的含义已从之前的 `True` 表示包含相应位置进行注意力计算，改为 `True` 表示屏蔽相应位置。自 v1.7 起，所有类型的注意力掩码都遵循相同的定义：`True` 表示屏蔽相应位置，`False` 表示包含该位置参与注意力计算。

例如，

.. code-block:: bash

    # 对于一个包含 3 个序列的批次，其中 `a`、`b` 和 `c` 是有效标记，而 `0` 是填充标记，
    [a, a, a, 0, 0,
     b, b, 0, 0, 0,
     c, c, c, c, 0]
    # 在 v1.7 之前，该批次的填充掩码是：
    [ 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.mosaicml.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 <https://github.com/deepspeedai/DeepSpeed>`_
* `Hugging Face Accelerate <https://huggingface.co/docs/accelerate/main/en/usage_guides/low_precision_training#configuring-transformersengine>`_
* `Lightning <https://lightning.ai/docs/pytorch/stable/common/precision.html>`_
* `MosaicML Composer <https://github.com/mosaicml/composer/releases/tag/v0.13.1>`_
* `NVIDIA JAX Toolbox <https://github.com/NVIDIA/JAX-Toolbox>`_
* `NVIDIA Megatron-LM <https://github.com/NVIDIA/Megatron-LM>`_
* `NVIDIA NeMo Megatron Bridge <https://github.com/NVIDIA-NeMo/Megatron-Bridge>`_
* `Amazon SageMaker Model Parallel Library <https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-core-features-v2-tensor-parallelism.html>`_
* `Levanter <https://github.com/stanford-crfm/levanter>`_
* `GPT-NeoX <https://github.com/EleutherAI/gpt-neox>`_
* `Hugging Face Nanotron <https://github.com/huggingface/nanotron>`_

贡献
============

我们欢迎对 Transformer Engine 的贡献！要为 Transformer Engine 做出贡献并提交拉取请求，请遵循 `<CONTRIBUTING.rst>`_ 指南中的说明。

论文
======

* `注意力机制原始论文 <https://proceedings.neurips.cc/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf>`_
* `Megatron-LM 张量并行 <https://arxiv.org/pdf/1909.08053.pdf>`_
* `Megatron-LM 序列并行 <https://arxiv.org/pdf/2205.05198.pdf>`_
* `用于深度学习的 FP8 格式 <https://arxiv.org/abs/2209.05433>`_

视频
======

* `Blackwell 上稳定且可扩展的 FP8 深度学习训练 | GTC 2025 <https://www.nvidia.com/en-us/on-demand/session/gtc25-s72778/>`_
* `Blackwell 中的 AI 数值计算 | GTC 2025 <https://www.nvidia.com/en-us/on-demand/session/gtc25-s72458/>`_
* `构建 LLM：使用 FP8 精度加速基础模型的预训练 | GTC 2025 <https://www.nvidia.com/gtc/session-catalog/?regcode=no-ncid&ncid=no-ncid&tab.catalogallsessionstab=16566177511100015Kus&search=zoho#/session/1726152813607001vnYK>`_
* `从 FP8 LLM 训练到推理：大规模语言 AI | GTC 2025 <https://www.nvidia.com/en-us/on-demand/session/gtc25-s72799/>`_
* `Transformer Engine 和 FP8 训练的新特性 | GTC 2024 <https://www.nvidia.com/en-us/on-demand/session/gtc24-s62457/>`_
* `使用 Transformer Engine 进行 FP8 训练 | GTC 2023 <https://www.nvidia.com/en-us/on-demand/session/gtcspring23-s51393>`_
* `FP8 在深度学习中的应用 | GTC 2023 <https://www.nvidia.com/en-us/on-demand/session/gtcspring23-s52166/>`_
* `Hopper 架构揭秘 | GTC 2022 <https://www.nvidia.com/en-us/on-demand/session/gtcspring22-s42663/>`_

.. |License| image:: https://img.shields.io/badge/License-Apache%202.0-blue.svg
   :target: https://opensource.org/licenses/Apache-2.0

往期新闻
=============

* [06/2025] `浮点数 8：高效、低精度 AI 训练简介 <https://developer.nvidia.com/blog/floating-point-8-an-introduction-to-efficient-lower-precision-ai-training/>`_
* [05/2025] `NVIDIA Grace Hopper 上 LLM 训练的高级优化策略 <https://developer.nvidia.com/blog/advanced-optimization-strategies-for-llm-training-on-nvidia-grace-hopper/>`_
* [03/2025] `Blackwell 上稳定且可扩展的 FP8 深度学习训练 | GTC 2025 <https://www.nvidia.com/en-us/on-demand/session/gtc25-s72778/>`_
* [03/2025] `使用 NVIDIA DGX Cloud 基准测试工具衡量并提升 AI 工作负载性能 <https://developer.nvidia.com/blog/measure-and-improve-ai-workload-performance-with-nvidia-dgx-cloud-benchmarking/>`_

.. image:: docs/examples/comparison-fp8-bf16-training-nvidia-dgx-cloud-benchmarking-performance-explorer.jpg
  :width: 600
  :alt: NVIDIA DGX Cloud 基准测试性能浏览器中显示的 FP8 与 BF16 训练对比图

* [02/2025] `借助 Evo 2，在全新尺度上理解生命生物分子在进化中的语言 <https://developer.nvidia.com/blog/understanding-the-language-of-lifes-biomolecules-across-evolution-at-a-new-scale-with-evo-2/>`_
* [02/2025] `NVIDIA DGX Cloud 推出即用型模板，用于基准测试 AI 平台性能 <https://developer.nvidia.com/blog/nvidia-dgx-cloud-introduces-ready-to-use-templates-to-benchmark-ai-platform-performance/>`_
* [01/2025] `利用 iGenius 和 NVIDIA DGX Cloud 继续预训练最先进的 LLM，以满足主权 AI 和受监管行业的需求 <https://developer.nvidia.com/blog/continued-pretraining-of-state-of-the-art-llms-for-sovereign-ai-and-regulated-industries-with-igenius-and-nvidia-dgx-cloud/>`_
* [11/2024] `使用 NVIDIA Megatron-LM 开发具有强大日语能力的 172B LLM <https://developer.nvidia.com/blog/developing-a-172b-llm-with-strong-japanese-capabilities-using-nvidia-megatron-lm/>`_
* [11/2024] `FP8 如何使 Amazon SageMaker P5 实例上的 LLM 训练速度提升 18% <https://aws.amazon.com/blogs/machine-learning/how-fp8-boosts-llm-training-by-18-on-amazon-sagemaker-p5-instances/>`_
* [11/2024] `使用 Amazon SageMaker 模型并行高效训练具有长序列长度的模型 <https://aws.amazon.com/blogs/machine-learning/efficiently-train-models-with-large-sequence-lengths-using-amazon-sagemaker-model-parallel/>`_
* [09/2024] `通过 FP8 混合精度训练升级，仅需一行代码即可将 AI 大模型训练成本降低 30% <https://company.hpc-ai.com/blog/reducing-ai-large-model-training-costs-by-30-requires-just-a-single-line-of-code-from-fp8-mixed-precision-training-upgrades>`_
* [05/2024] `使用 NVIDIA cuDNN 9 加速 Transformer 模型 <https://developer.nvidia.com/blog/accelerating-transformers-with-nvidia-cudnn-9/>`_
* [03/2024] `涡轮增压式训练：使用 FP8 优化 Databricks Mosaic AI 堆栈 <https://www.databricks.com/blog/turbocharged-training-optimizing-databricks-mosaic-ai-stack-fp8>`_
* [03/2024] `SageMaker 模型并行库支持 FP8 训练 <https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-release-notes.html>`_
* [12/2023] `NVIDIA NeMo 框架新功能及 NVIDIA H200 <https://developer.nvidia.com/blog/new-nvidia-nemo-framework-features-and-nvidia-h200-supercharge-llm-training-performance-and-versatility/>`_

.. image:: docs/examples/H200-NeMo-performance.png
  :width: 600
  :alt: H200

* [11/2023] `Inflection-2：下一步的飞跃 <https://inflection.ai/inflection-2>`_
* [11/2023] `借助 NVIDIA Transformer Engine 释放 Transformer 的强大威力 <https://lambdalabs.com/blog/unleashing-the-power-of-transformers-with-nvidia-transformer-engine>`_
* [11/2023] `使用 FP8 加速 PyTorch 训练工作loads <https://towardsdatascience.com/accelerating-pytorch-training-workloads-with-fp8-5a5123aec7d7>`_
* [09/2023] `AWS DL 容器中已添加用于 PyTorch 训练的 Transformer Engine <https://github.com/aws/deep-learning-containers/pull/3315>`_
* [06/2023] `凭借 NVIDIA H100 GPU 打破 MLPerf 训练记录 <https://developer.nvidia.com/blog/breaking-mlperf-training-records-with-nvidia-h100-gpus/>`_
* [04/2023] `CoreWeave 使用 NVIDIA H100 GPU 对大型语言模型进行基准测试（第一部分） <https://www.mosaicml.com/blog/coreweave-nvidia-h100-part-1>`_

Transformer Engine 快速上手指南

Transformer Engine (TE) 是 NVIDIA 推出的加速库，专为在 NVIDIA GPU（Hopper、Ada、Blackwell 架构）上加速 Transformer 模型设计。它支持 FP8、MXFP8 和 NVFP4 等低精度格式，能在几乎不损失精度的情况下显著提升训练和推理速度，并降低显存占用。

1. 环境准备

在开始之前，请确保您的系统满足以下硬件和软件要求：

硬件要求

• GPU 架构：NVIDIA Blackwell, Hopper, Grace Hopper/Blackwell, Ada, 或 Ampere。
  • 注意：FP8 特性需要 Compute Capability 8.9+ (Ada/Hopper/Blackwell)。

软件依赖

• 操作系统：Linux (官方支持), WSL2 (有限支持)。
• CUDA Toolkit：
  • CUDA 12.1+ (适用于 Hopper/Ada/Ampere)
  • CUDA 12.8+ (适用于 Blackwell)
  • 需搭配兼容的 NVIDIA 驱动。
• cuDNN：版本 9.3+。
• 编译器：GCC 9+ 或 Clang 10+ (需支持 C++17)。
• Python：推荐 3.12。
• 构建工具 (若从源码安装)：CMake 3.18+, Ninja, Git 2.17+, pybind11 2.6.0+。

2. 安装步骤

推荐使用 Docker 方式以获得最稳定的环境，也可通过 pip 直接安装。

方法一：使用 Docker (推荐)

NVIDIA GPU Cloud (NGC) 提供了预装好所有依赖的容器镜像，这是最快捷的上手方式。

PyTorch 环境：

docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:26.01-py3

JAX 环境：

docker run --gpus all -it --rm nvcr.io/nvidia/jax:26.01-py3

注：26.01 代表 2026 年 1 月发布的版本，建议前往 NGC Catalog 获取最新镜像。库已预装在容器内的 /opt/transformerengine 目录。

方法二：使用 pip 安装

如果您希望在现有环境中安装，请确保已安装兼容的 C++ 编译器和 CUDA Toolkit。

安装 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]

国内加速提示：如果下载速度较慢，可添加国内镜像源参数，例如： pip install --no-build-isolation transformer_engine[pytorch] -i https://pypi.tuna.tsinghua.edu.cn/simple

方法三：从 GitHub 源码安装

pip install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@stable

若需指定框架，可使用环境变量：

NVTE_FRAMEWORK=pytorch,jax pip install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@stable

3. 基本使用

Transformer Engine 提供了类似自动混合精度的 API，可无缝集成到现有代码中。以下是基于 PyTorch 的最简示例。

PyTorch 示例

此示例展示了如何创建一个支持 FP8 的线性层，并使用 autocast 上下文管理器启用低精度计算。

import torch
import transformer_engine.pytorch as te
from transformer_engine.common import recipe

# 设置维度
in_features = 768
out_features = 3072
hidden_size = 2048

# 初始化模型和输入数据
model = te.Linear(in_features, out_features, bias=True)
inp = torch.randn(hidden_size, in_features, device="cuda")

# 创建 FP8 策略 (所有参数均为可选)
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) 示例简述

对于 JAX 用户，用法类似，只需导入 transformer_engine.jax 模块并使用 te_flax.DenseGeneral 等组件，同样通过 te.autocast 上下文控制精度。详细用法请参考官方文档中的 JAX 部分。