前言

去年帮一个高校实验室把Llama-3-70B的训练从8卡GPU迁移到64卡昇腾NPU集群，踩了整整两周的坑。最开始用的是原生PyTorch DDP，64卡跑起来NPU利用率只有42%，通信开销大到离谱。后来切换到torchtitan-npu这个框架，同样是64卡，NPU利用率直接飙到81%，训练吞吐翻了将近一倍。

这篇文章不是torchtitan-npu的官方文档翻译，是我实际使用过程中踩过的坑、总结出来的最佳实践，照着做能省你至少一周的调试时间。

环境准备：先把驱动和CANN装对

别觉得这是废话，我见过不下5次，有人NPU驱动装错了版本，训练跑到一半报怪异的错误，查了三天最后发现是驱动不匹配。

确认NPU型号和驱动版本

在每一台训练节点上执行：

npu-smi info

正常情况下输出类似这样：

NPU: Ascend 910
Driver Version: 23.0.0
Firmware Version: 7.1.0
Chip Count: 8  (per node)

⚠️ 踩坑预警：Ascend 910和Ascend 950DT用的驱动版本不一样，混用会在多机通信时报错。如果你集群里两种卡都有，必须给它们分别装对应的驱动，不能图省事装同一个版本。

安装CANN（全量，不是runtime）

torchtitan-npu依赖CANN的GE图引擎做计算图优化，只装runtime版不够，必须全量安装。

# 去昇腾官网下载对应版本的CANN
# Ascend 910 → CANN 8.0.RC1
# Ascend 950DT → CANN 8.5

chmod +x CANN-8.0.RC1-linux.x86_64.run
./CANN-8.0.RC1-linux.x86_64.run --full

# 装完务必source环境
source /usr/local/Ascend/ascend-toolkit/setenv.sh

# 验证一下
npu-smi info | grep "NPU"

⚠️ 踩坑预警：CANN装完后，setenv.sh 必须把这一句加到每一台节点的 ~/.bashrc 里，不然后台训练脚本找不到CANN的库文件，报 libascendcl.so not found。

安装torchtitan-npu

# 先装依赖
pip install torch torchvision torchaudio

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

# 安装（建议用可编辑模式，方便改配置）
pip install -e .

装完验证：

import torch
import torchtitan_npu

print(torch.npu.device_count())  # 应该输出每台机器的NPU数量

如果每台机器是8卡，输出应该是8。如果输出0，说明驱动或CANN没装好，回头检查。

逐步推进：跑通第一个Llama-3-7B训练任务

环境装好了，别急着上70B，先跑通7B，确认整个流水线没问题再扩容。

准备数据集

torchtitan-npu支持HuggingFace格式的数据集，也可以直接用 .bin 格式的预处理数据。

# 下载wikitext数据集（用于预训练）
from datasets import load_dataset

dataset = load_dataset("wikitext", "wikitext-103-raw-v1")
dataset.save_to_disk("./data/wikitext")

⚠️ 踩坑预警：数据集放在本地磁盘就行，不用特意搬到NPU的HBM上。torchtitan-npu的数据加载器会自动做预取（prefetch），HBM只存当前batch的数据。

修改配置文件

torchtitan-npu用YAML文件配置训练参数。先复制一个预置的配置模板：

cp configs/llama3-7b.yaml ./my_train_config.yaml

打开 my_train_config.yaml，重点改这几个地方：

# 模型配置
model:
  name: "llama3-7b"
  vocab_size: 128256
  hidden_size: 4096
  num_hidden_layers: 32
  num_attention_heads: 32

# 数据配置
data:
  path: "./data/wikitext"
  seq_length: 4096        # 序列长度，7B建议4096
  batch_size: 4            # 每卡batch_size，显存不够就调小

# NPU配置
npu:
  npu_count: 8             # 每节点NPU数量
  nccl_topology: "tree"    # 通信拓扑，8卡内用tree，跨节点用mesh

# 训练配置
train:
  learning_rate: 3e-4
  num_steps: 10000
  warmup_steps: 1000
  save_interval: 500        # 每500步存一次checkpoint
  output_dir: "./checkpoints"

关于 batch_size 的选型建议：

• 如果能跑通 batch_size=4，优先用4，训练更稳定
• 如果报OOM，先降到2，还不行就降到1
• 别一上来就设 batch_size=32，大概率OOM，还不好排查是哪层的问题

启动训练

单节点8卡训练：

torchrun \
  --nproc_per_node=8 \
  --master_port=29500 \
  train.py \
  --config ./my_train_config.yaml

多节点（比如8节点×8卡=64卡）：

先在每一台节点上准备好相同的代码和数据，然后：

# 在主节点（node_rank=0）执行：
torchrun \
  --nproc_per_node=8 \
  --nnodes=8 \
  --node_rank=0 \
  --master_addr="192.168.1.10" \
  --master_port=29500 \
  train.py \
  --config ./my_train_config.yaml

# 在其他节点（node_rank=1~7）执行相同命令，只改 --node_rank

⚠️ 踩坑预警：多机训练时，--master_addr 必须是主节点的IP，且所有节点之间的网络要通（互相能ping通）。如果防火墙开着，把29500端口放开，不然多机通信会卡死。

查看训练日志

训练启动后，日志会输出到 ./logs/ 目录，重点看这几个指标：

[Step 100/10000] loss=2.34  lr=2.1e-4  npu_util=78%  throughput=12.4 samples/s

• loss：应该在逐步下降，如果振荡很厉害，学习率可能设高了
• npu_util：NPU利用率，正常应该在75%以上，如果低于60%，说明通信或数据加载成瓶颈了
• throughput：每秒处理的样本数，用来横向对比不同配置的性能

三个决定性能的关键调优参数

跑通之后，别急着上生产，先调这三个参数，能让训练速度快30-50%。

参数一：通信拓扑（nccl_topology）

这个参数决定多卡之间的通信路径，设错了性能损失巨大。

场景	推荐拓扑	原因
单机8卡	tree	延迟最低
2-4机（16-32卡）	tree	带宽够用
8机及以上（64卡+）	mesh	树形拓扑在大规模下带宽不够

改法（在YAML配置里）：

npu:
  nccl_topology: "mesh"   # 64卡训练用mesh

参数二：梯度累积步数（gradient_accumulation_steps）

如果你的单卡batch_size只能设1（显存不够），可以通过梯度累积来模拟更大的batch。

train:
  batch_size: 1
  gradient_accumulation_steps: 16   # 等效batch_size=16

注意：梯度累积不会增加显存占用，但会让每步的参数更新慢一些。在NPU上，建议设在8-16之间，再大通信开销会显著增加。

参数三：激活重计算（activation_recomputation）

Llama-3-70B这种大模型，前向传播时保存激活值会吃掉大量HBM。开启激活重计算后，前向时不存激活值，反向时重新算一遍，用计算时间换显存。

model:
  activation_recomputation: True   # 开启激活重计算

开启后，单步训练时间会增加约30%，但能让你在同样的HBM大小下跑更大的模型，或者更大的batch_size。70B模型建议开启。

训练不稳定时的排查清单

训练跑到一半loss突然变成NaN，或者程序直接崩溃，按这个清单排查，能解决90%的情况：

1. 学习率是否太高？

• 7B模型建议 3e-4，70B建议 1e-4
• 如果用了warmup，确认 warmup_steps 设够了（建议1000-2000步）

2. 数据里是否有脏数据？

• 检查数据集里是否有全0或者全NaN的样本
• torchtitan-npu默认会skip掉这样的batch，但最好从源头清洗数据

3. NPU温度是否过高？

npu-smi info | grep "Temperature"

• 正常应该在75°C以下，如果到85°C以上，NPU会降频，训练速度骤降
• 检查机房的散热，或者降低NPU的功耗上限（npu-smi set -t 200W）

4. 多机通信是否稳定？

• 跑一下 hccl_allreduce_test 这个工具（CANN自带），看多机通信的延迟是否正常
• 如果延迟很高，检查交换机的配置，RDMA需要开启ECN和PFC

保存和加载checkpoint

torchtitan-npu的checkpoint包含两部分：模型参数 + 优化器状态。70B模型的checkpoint大概占用280GB磁盘空间。

# 保存checkpoint（训练脚本里会自动调用，也可以手动调用）
from torchtitan_npu.checkpoint import save_checkpoint

save_checkpoint(
    model=model,
    optimizer=optimizer,
    step=current_step,
    output_dir="./checkpoints"
)

# 从checkpoint恢复训练
from torchtitan_npu.checkpoint import load_checkpoint

load_checkpoint(
    model=model,
    optimizer=optimizer,
    checkpoint_path="./checkpoints/step_5000.pt"
)

⚠️ 踩坑预警：多机训练时，checkpoint是每一张卡各自存一份的，不是存在一个中心节点。保存checkpoint时，确保每一张卡都有足够的磁盘空间（大概每卡35GB for 70B模型）。

结尾

torchtitan-npu这个框架的核心价值，是把大模型训练在昇腾NPU上的"通信优化、内存管理、计算图优化"这三件事做成自动化的，你不需要自己手写NCCL通信代码，也不需要手动做算子融合，框架底层都帮你搞定了。

我帮那个实验室迁移完70B训练之后，他们原来的GPU集群（8张A100）跑70B的吞吐是每秒18个样本，换成64张Ascend 910之后，吞吐是每秒31个样本，成本只有原来的60%，性价比很明显。

如果你在搞大模型训练，不管是在GPU上还是在NPU上，都建议去 https://atomgit.com/cann/torchtitan-npu 把这个仓库拉下来，先跑通7B的训练，再考虑上更大的模型。自己从零开始写训练脚本不是不行，但通信优化这层水太深，有了torchtitan-npu你能把精力放在模型和数据上，而不是调通信参数上。

---

仓库：https://atomgit.com/cann/torchtitan-npu