NVIDIA/TransformerEngine

NVIDIA Transformer Engine – 大型 Transformer 的快速低精度训练

是什么 – 一个 C++/CUDA 库（带 Python 绑定），为在 NVIDIA GPU 上训练 Transformer 类模型提供高度优化的构建模块。它为 PyTorch 和 JAX 提供融合内核和自动混合精度风格的 API，处理 FP8、MXFP8、NVFP4、BF16 和 FP16 训练所需的簿记工作。

为何重要 – 现代大语言模型（LLM）及其他大规模模型受限于训练所需的计算量和内存。通过将大部分算术运算迁移到更新的低精度格式（特别是 Hopper/Blackwell/Rubin GPU 上的 FP8），并融合多个操作（注意力、GEMM、归一化、MoE 路由等），Transformer Engine 可显著提升吞吐量，同时保持模型质量。它还抽象了低精度的细节，使您无需更改现有训练代码即可使用。

核心功能

• 在 NVIDIA Hopper、Ada、Blackwell 和 Rubin GPU 上支持 FP8 / MXFP8 / NVFP4 训练（支持这些格式的最新硬件）。
• 针对注意力、线性层、层归一化、MoE 路由和通信（张量、序列、上下文和专家并行）的优化融合内核。
• 自动缩放因子管理 – 库会跟踪 amax 历史和量化元数据，提供类似 PyTorch 原生 AMP 的 autocast 风格上下文。
• 框架支持 – 为 PyTorch（torch.nn.Module 包装器）和 JAX（Flax 兼容层）提供轻量、符合习惯的 API。
• 分布式训练就绪 – 与 NCCL 集成，支持通信重叠、零拷贝专家并行和 CUDA 图加速。
• 广泛的生态系统集成 – 与 Hugging Face Accelerate、PyTorch Lightning、NVIDIA Megatron‑LM、NVIDIA NeMo、Amazon SageMaker Model Parallel Library 等开箱即用。

典型工作流程

1. 安装您框架的包：

   pip install "transformer_engine[pytorch]"   # 或 [jax]
   

2. 用 TE 对应的版本包装模型的层（例如 te.Linear、te.LayerNorm、te.MultiHeadAttention）。
3. 将前向传播包裹在 autocast 上下文中，以自动启用 FP8/FP16/BF16：

   with transformer_engine.autocast(dtype=torch.float8):
       output = model(input)
   

4. 运行您通常的训练循环；TE 在后台处理缩放、量化和通信。

快速入门与资源

• 快速入门文档 – https://docs.nvidia.com/deeplearning/transformer-engine/getting_started/index.html
• API 参考 – PyTorch: https://docs.nvidia.com/deeplearning/transformer-engine/api/pytorch.html ; JAX: https://docs.nvidia.com/deeplearning/transformer-engine/api/jax.html
• 示例 – 仓库 examples/ 文件夹中的一组可直接运行的笔记本和脚本。
• 深入指南 – 低精度训练指南、FP8 入门、原始 FP8 论文（arXiv:2209.05433）。
• 发布说明 – 最新 v2.19 增加了 Rubin GPU 支持、混合量化和扩展的 FP8 注意力。

谁应该使用它

• 拥有最新 NVIDIA GPU 并训练大型语言模型、视觉 Transformer 或 MoE 模型的研究人员和工程师。
• 希望在不重写 PyTorch/JAX 代码库的情况下缩短训练时间并减少 GPU 内存使用量的团队。
• 在 Megatron‑LM、Hugging Face、NeMo 等主流 LLM 堆栈上构建并希望利用 FP8 加速的任何人。

许可证 – Apache 2.0（开源）。