torchtitan-npu：大模型训练框架快速上手实战

前言

去年帮一个高校实验室把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.5chmod+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

# 先装依赖pipinstalltorch torchvision torchaudio# 克隆仓库gitclone https://atomgit.com/cann/torchtitan-npu.gitcdtorchtitan-npu# 安装（建议用可编辑模式，方便改配置）pipinstall-e.

装完验证：

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

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

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

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

准备数据集

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

# 下载wikitext数据集（用于预训练）from datasetsimportload_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文件配置训练参数。先复制一个预置的配置模板：

cpconfigs/llama3-7b.yaml ./my_train_config.yaml

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

# 模型配置model:name:"llama3-7b"vocab_size:128256hidden_size:4096num_hidden_layers:32num_attention_heads:32# 数据配置data:path:"./data/wikitext"seq_length:4096# 序列长度，7B建议4096batch_size:4# 每卡batch_size，显存不够就调小# NPU配置npu:npu_count:8# 每节点NPU数量nccl_topology:"tree"# 通信拓扑，8卡内用tree，跨节点用mesh# 训练配置train:learning_rate:3e-4num_steps:10000warmup_steps:1000save_interval:500# 每500步存一次checkpointoutput_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）

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

改法（在YAML配置里）：

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

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

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

train:batch_size:1gradient_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（训练脚本里会自动调用，也可以手动调用）fromtorchtitan_npu.checkpointimportsave_checkpoint save_checkpoint(model=model,optimizer=optimizer,step=current_step,output_dir="./checkpoints")

# 从checkpoint恢复训练fromtorchtitan_npu.checkpointimportload_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