TransformerEngine

让 Transformer 在 NVIDIA GPU 上跑得更快更省显存

开源 🏷️完全免费 🔧 仅开源项目，无在线服务，需自行部署基础设施 🍴 复刻 885 🐛 待处理 378 👀 关注 37

访问官网 ↗ GitHub ↗

GitHub 星标 ★ 3597

维护状态活跃

是否开源开源

定价模式 🏷️完全免费开源免费（Apache 2.0），无付费版本

项目数据

分类基础设施

开发团队NVIDIA

官网地址https://docs.nvidia.com/deeplearning/transformer-engine

GitHubhttps://github.com/NVIDIA/TransformerEngine

定价模式🏷️完全免费

价格说明开源免费（Apache 2.0），无付费版本

访问状态🔧 仅开源项目，无在线服务，需自行部署

是否开源是

开源协议Apache-2.0

主要语言Python

技术栈/模型cuda,deep-learning,fp4,fp8,gpu,jax,machine-learning,python,pytorch

GitHub 星标★ 3597

上线时间2022-09-20 00:00:00

代码更新2026-10-30

仓库体积29.0 MB

本站数据更新2026-10-31

维护状态活跃

收录时间2026-10-31

浏览次数0

项目介绍

TransformerEngine 是 NVIDIA 推出的开源加速库，面向 Transformer 类模型在 NVIDIA GPU 上的训练与推理。它把低精度计算（FP8、FP4）与算子融合、自动混合精度管理结合起来，让开发者不必手写复杂的精度转换代码，就能在 Hopper、Ada、Blackwell 等新架构 GPU 上获得更高吞吐、更低显存占用。核心能力包括：对 PyTorch 和 JAX 提供统一的高层 API，自动管理 FP8/FP4 的缩放因子与精度切换；内置针对注意力、LayerNorm、线性层等 Transformer 关键模块的融合算子；支持训练与推理两种路径，并与 Megatron-LM、NeMo 等生态集成。它解决的是大模型训练显存吃紧、算力利用率不高、低精度改造门槛高的问题，适合需要压榨 GPU 性能的团队直接接入现有训练框架。

Star 增长趋势

本站每日记录该项目 GitHub 星标变化，观测越久曲线越完整

数据来源：GitHub API 逐日快照 · 收录于 2026-10-31 · 当前星标 ★ 3597

技术标签

cuda deep-learning gpu jax machine-learning

使用教程

⚠ 本教程由 AI 依据该项目 README 整理；原始文档不完整，部分步骤为通用流程推断，实际安装请以项目官方文档为准。

难度：进阶约 20 分钟部署方式：库/依赖 5 步

⚠ 该项目 README 未提供完整安装说明，以下步骤为通用流程，请以官方文档为准。

环境要求

需要 NVIDIA GPU（FP8 支持 Hopper、Ada、Blackwell、Rubin 架构；MXFP8/NVFP4 需要 Blackwell）
已安装可用的 CUDA 驱动与 CUDA 工具链
已安装 Python 环境
已安装 PyTorch 或 JAX 训练框架

安装与启动步骤

1确认GPU架构

先确认本机 GPU 属于 Hopper、Ada、Blackwell 或 Rubin，否则无法使用 FP8 低精度训练能力。
```
nvidia-smi
```
2准备CUDA环境

确认驱动与 CUDA 工具链版本可用，低精度算子编译依赖本地 CUDA 环境。
```
nvcc --version
```
3准备Python框架

在 Python 环境中装好 PyTorch 或 JAX，TransformerEngine 以库的形式接入现有训练栈。
```
python -c "import torch; print(torch.__version__)"
```
4获取项目源码

克隆官方仓库，可直接查看源码与 examples 目录下的示例脚本。
```
git clone https://github.com/NVIDIA/TransformerEngine.git
cd TransformerEngine
```
5查阅官方文档

安装方式与版本对应关系以官方 User guide 的 Quick start 为准，按文档安装后再接入代码。