Skip to main content

Ascend End-to-End Large Model Training Adaptation Framework Based on torchtitan

Project description

torchtitan-npu

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

Documentation license contributing SIG contributing

简介


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

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

社群

SIG

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 使用者

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

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

License 声明


  • torchtitan‑npu 产品的使用许可证,具体请参见 LICENSE
  • torchtitan‑npu 工具 docs 目录下的文档适用相应许可证,具体请参见文档目录下的 LICENSE 文件。

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

torchtitan_npu-0.2.2.post1.tar.gz (87.9 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

torchtitan_npu-0.2.2.post1-py3-none-any.whl (115.1 kB view details)

Uploaded Python 3

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

Hashes for torchtitan_npu-0.2.2.post1.tar.gz
Algorithm Hash digest
SHA256 2c690dd613cfcd184edaaea63b07e72c063120ea1397ac543d28458444d8bc56
MD5 2881f1005ab47616eb11f33a8cd06956
BLAKE2b-256 d2e7100533f651a8318cdb1a2a9a5a52feba75232e381cacbd8600df364f5a09

See more details on using hashes here.

File details

Details for the file torchtitan_npu-0.2.2.post1-py3-none-any.whl.

File metadata

File hashes

Hashes for torchtitan_npu-0.2.2.post1-py3-none-any.whl
Algorithm Hash digest
SHA256 fb889eb4d503b5330725ea123fe596a553c1206e5bb9a4403e15fcd6283f5e8b
MD5 f1f553a917e148ee99670e1617096216
BLAKE2b-256 df0f608ad8dfd9bb2a05b4a533f26373a7cc2cd80f82a497e7dd2c5e67358e24

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page