CANN torchtitan-npu昇腾训练适配插件-育师

【免费下载链接】torchtitan-npuAscend Extension for torchtitan项目地址: https://gitcode.com/cann/torchtitan-npu

torchtitan-npu

基于 torchtitan 的昇腾全流程大模型训练适配插件

简介

torchtitan-npu定位为torchtitan的昇腾（Ascend）后端扩展插件，通过即插即用的硬件亲和性优化，充分释放NPU算力，助力PyTorch native训练在昇腾平台无缝、高效、稳定地运行。

本插件基于社区ModelConverter拓展机制构建，已支持多维度训练优化，涵盖 NPU融合算子、图优化、图下沉、算子自动融合、显存管理、分布式并行以及调试维测能力等等。

社群

SIG 例会：sig-framework-adapter

Roadmap

当前季度的规划见torchtitan-npuRoadmap。欢迎访问。

安装

源码安装：

git clone https://gitcode.com/cann/torchtitan-npu.git cd torchtitan-npu pip install -e .

详情参见安装教程。

快速上手

快速启动大语言模型的训练任务，参见快速上手文档。

云开发平台（2 die）单机最小可运行样例，参见云平台开发指南。

特性支持概览

场景	特性名称	原生支持	NPU支持
并行能力	4D 并行 (FSDP2/TP/CP/PP)	✅	✅
	专家并行 (EP/ETP)	✅	✅
	自定义 CP (DeepSeek V3.2 CP/SDPA Ulysses CP)	❌	✅
torch.compile	torch.compile	✅	✅
训练精度	MxFP8 量化	✅	✅ (Ascend 950)
训练精度	HiF8 量化	❌	✅ (Ascend 950)
训练调试与监控	分布式 Checkpoint	✅	✅
训练调试与监控	调试工具	✅	✅
性能优化	Swap Optimizer	❌	✅
性能优化	NPU 融合算子适配	❌	✅

项目结构

torchtitan-npu 充分利用了 torchtitan 提供的 ModelConverter 插件化机制。该机制介入模型定义之后、并行策略（如 TP/FSDP）应用之前，支持以非侵入式的方式，通过注册机制对特定模块进行替换或重写。基于此方案，我们实现了融合算子优化、量化支持以及优化器增强等功能。见以下项目结构：

torchtitan-npu/ ├── torchtitan_npu/ # torchtitan_npu核心源代码 │ ├── config/ # 对Config的补丁 │ ├── converters/ # 基于torchtitan ModelConverter机制的补丁 │ ├── distributed/ # 自定义分布式代码 │ ├── models/ # 基于torchtitan-npu的模型 (如Deepseek-V3.2) │ ├── patches/ # 其他补丁 │ ├── tools/ # 工具补丁 │ ├── entry.py # 启动训练 │ ├── train.py # 训练主流程补丁 │ └── __init__.py # torchtitan-npu 插件修改注入点 ├── docs/ # 文档

上下游软件栈架构图如下:

性能基准

2026.04

System: Atlas 800T A3 | Model | Number of NPUs | Precision | GBS | Local BS | Sequence Length | FSDP | TP | PP | CP | EP | Throughput (tokens/p/s) | MFU | | :----------------- | :------------- | :-------- | :-- | :------- | :-------------- | :--- | :-- | :-- | :-- | :-- | :----------- | :-- | | DeepSeek-V4-Flash | 64 | BF16 | 1024 | 1 | 4096 | 128 | 1 | 1 | 1 | 128 | 1100 | 28.78% | | DeepSeek-V3.2-671B | 64 | BF16 | 128 | 1 | 32768 | 4 | 4 | 1 | 8 | 64 | 103 | / | | DeepSeek-V3.2-671B | 64 | BF16 | 512 | 1 | 4096 | 32 | 4 | 1 | 1 | 64 | 146 | / | | DeepSeek-V3-671B | 64 | BF16 | 1024 | 1 | 4096 | 32 | 4 | 1 | 1 | 128 | 546 | / | | DeepSeek-V3-671B + compile(Autofuse) | 64 | BF16 | 1024 | 1 | 4096 | 32 | 4 | 1 | 1 | 128 | 576 | / |

注：以上MoE模型的性能数据均开启负载均衡配置moe_force_load_balance=true。

免责声明

致 torchtitan‑npu 使用者

torchtitan‑npu 提供的所有内容仅供您用于非商业目的。
对于 torchtitan‑npu 测试用例以及示例文件中所涉及的各模型和数据集，平台仅用于功能测试，华为不提供任何模型权重和数据集。如您使用这些数据进行训练，请您特别注意应遵守对应模型和数据集的 License，如您因使用这些模型和数据集而产生侵权纠纷，华为不承担任何责任。
如您在使用 torchtitan‑npu 过程中，发现任何问题（包括但不限于功能问题、合规问题），请在 GitCode 提交 issue，我们将及时审视并解决。

torchtitan‑npu 功能依赖的 PyTorch 等第三方开源软件，均由第三方社区提供和维护，因第三方开源软件导致的问题的修复依赖相关社区的贡献和反馈。您应理解，torchtitan‑npu 仓库不保证对第三方开源软件本身的问题进行修复，也不保证会测试、纠正所有第三方开源软件的漏洞和错误。