Ascend End-to-End Large Model Training Adaptation Framework Based on torchtitan
Project description
简介
torchtitan-npu定位为torchtitan的昇腾(Ascend)后端扩展插件,通过即插即用的硬件亲和性优化,充分释放NPU算力,助力PyTorch native训练在昇腾平台无缝、高效、稳定地运行。
本插件基于社区 ModelConverter 拓展机制构建,已支持多维度训练优化,涵盖 NPU融合算子、图优化、图下沉、算子自动融合、显存管理、分布式并行以及调试维测能力等等。
社群
SIG 例会:sig-framework-adapter
最新消息
- [Apr. 2026]: 🚀 【重要特性支持】算子自动融合:基于AscendC AutoFuse的能力,支持torch.compile + Inductor后端的算子自动融合。
- [Apr. 2026]: 🚀 torchtitan‑npu 正式开源:在 NPU 上支持 4D 并行等 torchtitan 原生特性,并引入 Swap Optimizer 等 NPU 亲和优化。
Roadmap
当前季度的规划见 torchtitan-npu Roadmap。欢迎访问。
安装
源码安装:
git clone https://gitcode.com/cann/torchtitan-npu.git
cd torchtitan-npu
pip install -e . ,
详情请参考 部署文档 安装torchtitan-npu及其依赖。
快速上手
指导开发者快速启动大语言模型的训练任务,具体的操作请参考: 快速入门(基于PyTorch框架)
特性支持概览
| 场景 | 特性名称 | 原生支持 | 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) |
|---|---|---|---|---|---|---|---|---|---|---|---|
| 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 仓库不保证对第三方开源软件本身的问题进行修复,也不保证会测试、纠正所有第三方开源软件的漏洞和错误。
License 声明
- torchtitan‑npu 产品的使用许可证,具体请参见 LICENSE。
- torchtitan‑npu 工具 docs 目录下的文档适用相应许可证,具体请参见文档目录下的 LICENSE 文件。
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file torchtitan_npu-0.2.2.post1.tar.gz.
File metadata
- Download URL: torchtitan_npu-0.2.2.post1.tar.gz
- Upload date:
- Size: 87.9 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.11.7
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
2c690dd613cfcd184edaaea63b07e72c063120ea1397ac543d28458444d8bc56
|
|
| MD5 |
2881f1005ab47616eb11f33a8cd06956
|
|
| BLAKE2b-256 |
d2e7100533f651a8318cdb1a2a9a5a52feba75232e381cacbd8600df364f5a09
|
File details
Details for the file torchtitan_npu-0.2.2.post1-py3-none-any.whl.
File metadata
- Download URL: torchtitan_npu-0.2.2.post1-py3-none-any.whl
- Upload date:
- Size: 115.1 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.11.7
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
fb889eb4d503b5330725ea123fe596a553c1206e5bb9a4403e15fcd6283f5e8b
|
|
| MD5 |
f1f553a917e148ee99670e1617096216
|
|
| BLAKE2b-256 |
df0f608ad8dfd9bb2a05b4a533f26373a7cc2cd80f82a497e7dd2c5e67358e24
|