transformerengine

Transformer Engine (TE) 是 NVIDIA 推出的一个用于在 NVIDIA GPU 上加速 Transformer 模型的库。其核心特性是低精度计算：在 Hopper、Ada 和 Blackwell GPU 上支持 8 位浮点（FP8），并在 Blackwell 上支持 MXFP8 和 NVFP4 格式，旨在提升训练和推理的性能并减少内存占用。它还支持 Ampere 及更新架构上 FP16 和 BF16 的优化。 该库为常见的 Transformer 架构提供了高度优化的构建模块，以及一个类似自动混合精度的 API，可与特定框架的代码配合使用。它还包含一个与框架无关的 C++ API，使其他深度学习库能够为 Transformer 添加 FP8 支持。TE 模块内部维护 FP8 训练所需的缩放因子和相关数值，从而简化了用户的混合精度工作流。 项目列出的亮点包括：用于构建支持 FP8 的 Transformer 层的易用模块、融合内核优化、在 Hopper、Ada 和 Blackwell 上的 FP8 支持、在 Blackwell 上的 MXFP8 和 NVFP4 支持，以及 Ampere 及更新架构上 FP16/BF16 的优化。 文档给出了 PyTorch 和 JAX/Flax 的使用示例。在 PyTorch 中，用户导入 transformer_engine.pytorch，创建带有选定 FP8 格式的配方（例如 DelayedScaling），并将前向传播包裹在 te.autocast 中。在 JAX/Flax 中，展示了使用 te_flax 模块和配方的类似 autocast 用法。文档还链接了入门指南以提供更完整的教程。 安装选项包括 NGC Docker 容器（推荐）、带有 PyTorch 和/或 JAX 附加组件的 pip 包、用于 PyTorch 集成的 conda-forge 包，以及源码构建。系统要求提到 Blackwell、Hopper、Grace Hopper/Blackwell、Ada 和 Ampere 硬件；官方操作系统为 Linux，对 WSL2 提供有限支持；CUDA 12.1+（Blackwell 需要 12.8+）；cuDNN 9.12+；GCC 9+ 或 Clang 10+ 并支持 C++17；推荐 Python 3.12。FP8 功能需要计算能力 8.9 或更高。文档还介绍了与构建相关的环境变量，如 CUDA_PATH、CUDNN_PATH、CXX、NVTE_FRAMEWORK、MAX_JOBS 和 NVTE_CUDA_ARCHS。 README 涵盖了 PyTorch 中对 FlashAttention-2 和 FlashAttention-3 的支持，当两者同时存在时优先使用 FlashAttention-3，并指出 FlashAttention-2 的编译可能非常消耗内存。故障排除部分涉及 ABI 兼容性导入错误、缺失的头文件或库、构建资源问题、详细构建日志、UV/虚拟环境问题（包括 cuDNN 子库加载失败）以及 JAX FFI 注册错误。 文档记录了 v1.7 的一项破坏性变更：PyTorch 填充掩码的定义发生了变化，现在 True 表示屏蔽某个位置而非包含该位置，从而统一了各框架的掩码语义。收敛性说明指出，在测试配置中，FP8 和 MXFP8 与 BF16 的训练损失曲线没有显著差异，并已在下游 LLM 任务上进行了验证，列出了 MPT-1.3B、Llama2-7B、LLM-8B、MPT-13B、MoE-16B 和 Llama2-70B 等模型，涵盖 Mosaic Composer、Alibaba Pai 和 Megatron Core 等框架。 列出的集成包括 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。该项目采用 Apache-2.0 许可证，并欢迎通过其 CONTRIBUTING 指南贡献代码。