TransformerEngine

TransformerEngine 项目介绍

项目背景

TransformerEngine（简称TE）由NVIDIA开发，是一个用于在NVIDIA GPU上加速Transformer模型的库。特别是在Hopper GPU上使用8位浮点精度（FP8），不仅可以提高性能，还可以降低训练和推理时的内存消耗。随着各种Transformer架构的参数数量不断增长，例如BERT、GPT和T5，这些模型在训练和推理阶段需要消耗大量的内存和计算资源。

核心功能

TransformerEngine 提供了一组高度优化的构建模块，这些模块适用于流行的Transformer架构。此外，它还带有一个类似自动混合精度的API，用户可以无缝地将其与特定深度学习框架的代码结合使用。

自动化FP8支持：TE包括一个框架无关的C++ API，能够与其他深度学习库集成，为Transformers提供FP8支持。
高效的模块化设计：TE提供的模块可以帮助用户在内部维护FP8训练所需的缩放因子和其他值，简化了混合精度训练。
对多种精度的优化支持：不仅支持FP8，还对FP16和BF16等精度提供优化，尤其针对NVIDIA Ampere GPU架构及更高版本。

使用指南

PyTorch 示例

在PyTorch使用TE进行FP8训练非常简单，只需以下几步：

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)

# 启用FP8自动类型转换进行前向计算
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
    out = model(inp)

loss = out.sum()
loss.backward()

安装方法

通过Docker快速入门：可以利用NVIDIA GPU Cloud（NGC）上的容器快速开始。
通过pip安装：使用命令pip install transformer_engine[pytorch]可以获取最新稳定版，支持多种深度学习框架。
从源码编译：如果需要精细化控制，可以选择从源码安装。

整合支持

TransformerEngine已集成于多个流行的大型语言模型（LLM）框架中，如DeepSpeed、Hugging Face Accelerate、NVIDIA Megatron-LM、Amazon SageMaker Model Parallel Library等。

FP8收敛性验证

FP8在不同模型架构和配置下进行了广泛测试，与BF16相比，无论训练损失曲线还是下游任务的准确性，均未表现出显著差异。

结论

TransformerEngine不仅为Transformer模型提供了显著的性能提升，还保持了对准确性的保证，特别是在资源消耗高的应用情境下。此外，它持续支持最新的NVIDIA硬件架构，为用户提供了高效的深度学习解决方案。通过与流行框架的整合，它为开发者提供了丰富的工具来优化大规模语言模型的训练和运行性能。