0
代码
代码
Issues
Pull Requests
讨论
Wiki
分析
0
cann-robotcann-robot[feat]Add VLM NPU training support
6139dd1f创建于 1 天前历史提交

torchtitan-npu

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

Documentation license contributing SIG pypi zread

简介

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

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

社群

SIG

SIG 例会：sig-framework-adapter

最新消息

  • [May. 2026]: 🚀 DeepSeek-V4-Pro 模型续训练支持:基于纯FSDP + 大EP极简切分，使能AutoFuse特性，达成训练入图。
  • [May. 2026]: ⚠️ 配置系统重构: master 分支对齐 torchtitan main 的 config_registry.py / ConfigManager 机制，模型训练使用 --module 和 --config 启动，不再通过 --job.config_file 加载 TOML。
  • [Apr. 2026]: 🚀 DeepSeek-V4-Flash 续训练 0day 支持:基于纯FSDP + 大EP极简切分，使能AutoFuse特性，达成训练入图，开箱即优。
  • [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 .

详情参见 安装教程 。

快速上手

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

特性支持概览

场景 特性名称 原生支持 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 融合算子适配 ❌ ✅
多模态模型 VLM debug model ✅ ✅

项目结构

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

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

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

性能基准

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 1056 27.67%
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 文件。

🤝联系我们

本项目功能和文档正在持续更新和完善中，欢迎您关注最新版本。

  • 问题反馈：通过GitCode【Issues】提交问题。
  • 社区互动：通过GitCode【讨论】参与交流。
  • 经验分享：通过GitCode【Wiki】分享经验总结。
  • 加入交流群：通过扫描下方微信二维码添加torchtitan‑npu小助手微信，加入微信群与我们进一步交流。

contact us