提供适配昇腾的TransformerEngine加速库
简介
TransformerEngine是NVIDIA提供的一个用于加速Transformer模型训练和推理的库，而TransformerEngineNPU则是其在昇腾设备上的适配版本，旨在为昇腾设备上的大模型训练和推理提供高效的计算支持。
昇腾提供基于Megatron-LM + MegatronAdaptor + TransformerEngineNPU的基础训练能力方案，使用户能够快速在昇腾设备上构建起大模型训练环境，实现开箱可用。
· 适配Megatron-LM在昇腾设备上运行的适配库: MegatronAdaptor
最新消息
· [May 12, 2026]: 🚀MegatronAdaptor + TransformerEngineNPU 支持 Mcore0.17.0 版本
版本说明
当前版本推荐配套表如下：
| 组件 | 版本要求 |
|---|---|
| MegatronAdaptor | 0.17.0 |
| Mcore | 0.17.0 |
| TransformerEngineNPU | main |
| CANN版本 | 9.0.0 |
| PyTorch版本 | 2.7.1 |
| torch_npu版本 | 26.0.0 |
| Python版本 | Python3.10.x |
版本对应表：
| TransformerEngineNPU | TransformerEngine |
|---|---|
| main | main |
安装
拉取MegatronAdaptor及TransformerEngineNPU的源码进行安装：
git clone https://gitcode.com/Ascend/MegatronAdaptor.git
cd MegatronAdaptor
git checkout core_r0.17.0
cd ..
pip install -e MegatronAdaptor
git clone https://gitcode.com/Ascend/TransformerEngineNPU.git
pip install -e TransformerEngineNPU
获取并切换Megatron-LM版本至 core_r0.17.0 版本，参考：
git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
git checkout core_r0.17.0
快速开始
概述
要开始使用MegatronAdaptor + TransformerEngineNPU，只需增加一行代码，即可在昇腾设备上启用Megatron-LM。
操作方法
以GPT模型为例：在Megatron-LM目录下修改pretrain_gpt.py文件，在import torch下新增一行：import megatron_adaptor，即如下修改：
from functools import partial
from typing import List, Optional, Tuple
import torch
import megatron_adaptor # 新增此行代码
from gpt_builders import gpt_builder
注意事项
TransformerEngineNPU 和 原生TransformerEngine 的库模块名相同，为了避免冲突，请确保同一环境内不要同时安装两者，仅保留TransformerEngineNPU即可。
使用限制
· TransformerEngineNPU不支持NVFP4低精度量化
免责声明
致TransformerEngineNPU使用者
- TransformerEngineNPU提供的所有内容仅供您用于非商业目的。
- 对于TransformerEngineNPU测试用例以及示例文件中所涉及的各模型和数据集，平台仅用于功能测试，华为不提供任何模型权重和数据集，如您使用这些数据进行训练，请您特别注意应遵守对应模型和数据集的License，如您因使用这些模型和数据集而产生侵权纠纷，华为不承担任何责任。
- 如您在使用TransformerEngineNPU过程中，发现任何问题（包括但不限于功能问题、合规问题），请在GitCode提交issue，我们将及时审视并解决。
- TransformerEngineNPU功能依赖的torch等第三方开源软件，均由第三方社区提供和维护，因第三方开源软件导致的问题的修复依赖相关社区的贡献和反馈。您应理解，TransformerEngineNPU仓库不保证对第三方开源软件本身的问题进行修复，也不保证会测试、纠正所有第三方开源软件的漏洞和错误。