你试过在昇腾NPU上训练大模型吗？

我试过。用原生 PyTorch 跑，报了一堆奇怪的错。查文档，文档说"请参考 CANN 官方文档"。去翻 CANN 官方文档，文档说"请参考 PyTorch 官方文档"。

踢皮球踢了一下午，一个 epoch 都没跑出来。

后来才知道，有个东西叫 torchtitan-npu。它把昇腾NPU的训练适配全部做好了，你只需要改几行代码，就能把大模型跑起来。

这东西就是昇腾NPU上的"一键训练脚本"。

先说清楚：torchtitan-npu 是什么？

根据昇腾CANN开源社区的知识库：

torchtitan-npu 是基于 PyTorch Titan 移植的昇腾NPU适配版本，提供大模型分布式训练的全套解决方案。

你可以理解为：

• PyTorch Titan = NVIDIA 官方的大模型训练参考实现（支持 Megatron-LM、DeepSpeed 等）
• torchtitan-npu = 把 PyTorch Titan 适配到昇腾NPU上的版本

简单说，就是：NVIDIA 能跑的分布式训练，昇腾NPU 也能跑。

为什么需要 torchtitan-npu？

在昇腾NPU上跑大模型训练，有几个绕不过去的坎：

坎1：分布式训练框架适配

PyTorch 的分布式训练（DistributedDataParallel、FSDP）都是针对 NVIDIA 的 NCCL 通信库写的。

昇腾NPU用的是 HCCL，接口不一样。

直接用原生 PyTorch 跑，报错：

RuntimeError: Distributed package doesn't have NCCL built-in

torchtitan-npu 把这个坑填了，它把 NCCL 全部替换成 HCCL，你不用改一行代码。

坎2：算子适配

大模型训练涉及很多自定义算子（比如 FlashAttention、MoE 路由）。

这些算子在 NVIDIA 上有 CUDA 实现，在昇腾NPU上没有。

torchtitan-npu 提供了昇腾NPU上的算子实现，比如：

• FlashAttention（基于 ops-transformer）
• MoE 算子（基于自定义 Ascend C 实现）

坎3：显存优化

大模型训练最吃显存。昇腾NPU的显存比 NVIDIA A100 小（32GB vs 80GB）。

torchtitan-npu 内置了多种显存优化技术：

• Gradient Checkpointing（用计算换显存）
• ZeRO（分布式显存优化）
• Int8 量化训练

环境准备：装 torchtitan-npu 比想象中简单

第一步：确认硬件和软件环境

硬件：昇腾910（Ascend 910）× 8 卡
CANN版本：8.0 及以上
Python：3.8 - 3.11
PyTorch：2.1 及以上

第二步：安装 torchtitan-npu

# 克隆仓库
git clone https://atomgit.com/cann/torchtitan-npu.git
cd torchtitan-npu

# 安装依赖
pip install -r requirements.txt

# 安装 torchtitan-npu
pip install -e .

requirements.txt 里的关键依赖：

torch>=2.1.0
torch_npu>=2.1.0 # 昇腾 PyTorch 适配层
ascend-toolkit>=8.0.0 # CANN Toolkit
hccl-utils-npu>=1.0.0 # HCCL 工具

第三步：验证安装

import torch
import torch_npu
import torchtitan

print(f"PyTorch version: {torch.__version__}")
print(f"torchtitan-npu version: {torchtitan.__version__}")
print(f"NPU available: {torch.npu.is_available()}")
print(f"NPU device count: {torch.npu.device_count()}")

如果输出类似下面，说明安装成功：

PyTorch version: 2.1.0
torchtitan-npu version: 1.0.0
NPU available: True
NPU device count: 8

逐步推进：用 torchtitan-npu 训练 LLaMA-2-7B

Step 1：准备数据集

torchtitan-npu 支持常见的预训练数据集格式（Arrow、Parquet）。

以 Arrow 格式为例：

from datasets import load_dataset

# 加载数据集
dataset = load_dataset("arrow", data_dir="./data/my_dataset")

# 保存为 torchtitan 需要的格式
dataset.save_to_disk("./data/llama2_dataset")

Step 2：写训练配置

torchtitan-npu 用 YAML 文件配置训练参数：

# config_llama2_7b.yaml
model:
 name: llama2_7b
 vocab_size: 32000
 hidden_size: 4096
 num_attention_heads: 32
 num_layers: 32
 max_seq_len: 4096

training:
 batch_size: 4
 learning_rate: 3e-4
 num_epochs: 3
 warmup_steps: 100
 gradient_clip: 1.0
 weight_decay: 0.1

distributed:
 backend: hccl # ← 关键：这里写 hccl，不是 nccl
 world_size: 8 # ← 8 卡训练
 master_addr: "127.0.0.1"
 master_port: "29500"
 fp8_enabled: false # ← 昇腾910暂不支持 FP8

npu:
 memory_fraction: 0.9 # ← 显存占用上限
 checkpoint_enabled: true # ← 开启 Gradient Checkpointing
 zero_stage: 2 # ← ZeRO-2 优化

Step 3：启动训练

# 单机8卡训练
torchrun --nproc_per_node=8 \
 torchtitan/train.py \
 --config config_llama2_7b.yaml \
 --data-path ./data/llama2_dataset \
 --output-dir ./outputs/llama2_7b

训练日志会输出类似：

[2026-05-21 10:00:00] Training started
[2026-05-21 10:00:05] Model initialized on 8 NPUs
[2026-05-05] Step 100/1000 | Loss: 2.345 | LR: 1.0e-5 | Throughput: 45 tokens/s/gpu
[2026-05-21 10:01:00] Step 200/1000 | Loss: 1.987 | LR: 2.0e-5 | Throughput: 47 tokens/s/gpu
...
[2026-05-21 11:30:00] Training completed | Total time: 90 min | Avg throughput: 48 tokens/s/gpu

Step 4：查看训练结果

# 训练完成后，加载 checkpoint
from torchtitan.checkpoint import load_checkpoint

checkpoint_path = "./outputs/llama2_7b/checkpoints/step_1000"
model = load_checkpoint(checkpoint_path, model)

# 验证模型
from torchtitan.generate import generate

input_text = "The future of AI is"
output = generate(model, input_text, max_new_tokens=100)
print(output)

进阶：多机训练怎么配？

如果你有多台服务器，每台 8 张卡，要做多机训练：

Step 1：每台机器上的配置一样

# 所有机器上的 config.yaml 一样
distributed:
 backend: hccl
 world_size: 16 # ← 2台机器 × 8卡 = 16
 rank: 0 # ← 这台机器的 rank（0 或 1）
 master_addr: "10.0.0.1" # ← 主节点的 IP
 master_port: "29500"

Step 2：分别启动

# 机器 0（主节点）
torchrun --nproc_per_node=8 \
 --nnodes=2 --node_rank=0 \
 --master_addr="10.0.0.1" \
 --master_port="29500" \
 torchtitan/train.py \
 --config config.yaml

# 机器 1（从节点）
torchrun --nproc_per_node=8 \
 --nnodes=2 --node_rank=1 \
 --master_addr="10.0.0.1" \
 --master_port="29500" \
 torchtitan/train.py \
 --config config.yaml

Step 3：确保节点间网络互通

# 测试节点间网络
ping -c 3 10.0.0.2

# 检查 HCCL 通信
python -c "import torch_npu; print(torch_npu.nccl.version())"

性能调优：torchtitan-npu 有哪些调优技巧？

技巧1：开启 Gradient Checkpointing（用计算换显存）

显存不够时，开启 Gradient Checkpointing 能把显存占用减半：

# config.yaml
npu:
 checkpoint_enabled: true # ← 开启
 checkpoint_ratio: 0.5 # ← 每2层checkpoint一次

代价是训练速度慢 20-30%，但能跑更大的模型。

技巧2：ZeRO 分布式优化

ZeRO 把optimizer状态、梯度、参数分片到不同 GPU：

# config.yaml
npu:
 zero_stage: 2 # ← ZeRO-2（推荐）
 # zero_stage: 3 # ← ZeRO-3（显存更省，但通信更慢）

技巧3：开启混合精度训练

昇腾910支持FP16和BF16：

# config.yaml
training:
 dtype: bf16 # ← 推荐用 BF16，比 FP16 数值稳定性好

技巧4：调整 Batch Size

Batch Size 对训练吞吐影响很大。原则是：往大了调，直到显存不够。

# 尝试不同的 batch size
for bs in 4 8 16 32; do
 torchrun ... --config "batch_size=$bs" ...
done

极简总结

torchtitan-npu = 昇腾NPU上的"一键训练脚本"。
它把分布式训练、算子适配、显存优化全部做好，你只需要改配置、改数据，就能训练大模型。
支持单机8卡、多机多卡、ZeRO、Gradient Checkpointing、混合精度。
NVIDIA 能跑的训练，昇腾NPU 也能跑。

仓库链接：
https://atomgit.com/cann/torchtitan-npu
https://atomgit.com/cann/cann-recipes-train