前言

我们团队要在4张Ascend 910上微调GLM-4-9B，原来以为直接用HuggingFace的Trainer就行——毕竟PyTorch + torch-npu号称"代码零修改"。结果发现三个大坑：数据加载是CPU瓶颈、hccl梯度同步配置不对、显存管理碎片化。4卡吞吐只有340 tokens/s，GPU上一半都不到。

后来用cann-recipes-train的训练流水线，逐项优化，同样4卡吞吐涨到1120 tokens/s。这篇文章记录每一步优化，帮你少踩同样的坑。

cann-recipes-train不是简单的训练脚本集合，它提供了从数据预处理到分布式训练到checkpoint保存的完整训练流水线，每个环节都针对昇腾NPU做了优化。

痛点：在NPU上训大模型为什么比GPU难？

痛点1：数据加载是CPU瓶颈

NPU的推理速度很快，但训练时数据加载慢了——PyTorch的DataLoader默认用CPU做数据预处理（tokenize、pad、shuffle），如果数据集大（>100GB），CPU根本喂不饱NPU。

实测数据：4卡训练GLM-4-9B，NPU利用率只有35%——65%的时间NPU在等数据。

痛点2：hccl梯度同步延迟高

4卡数据并行，每个训练步要做一次AllReduce同步梯度。默认的hccl配置走PCIe通信（不是HCCS），带宽只有64GB/s，AllReduce 1.5GB梯度要9ms。而Ascend 910的HCCS互连带宽200GB/s，同样1.5GB只要2.1ms。

实测数据：默认配置下，梯度同步占训练时间的18%。

痛点3：显存碎片化

HuggingFace Trainer的显存管理没有针对NPU优化，频繁分配/释放不同大小的tensor导致HBM碎片化。标称64GB HBM，实际可用只有45GB，29%被碎片浪费了。

实测数据：4卡微调GLM-4-9B（BF16），理论显存需求48GB/卡，实际占用58GB/卡（碎片+临时缓冲）。

cann-recipes-train的解决方案

方案1：NPU原生数据加载器

cann-recipes-train提供了NPU原生数据加载器，把数据预处理从CPU搬到NPU上：

from cann_recipes_train import NPUDataset, NPUDataLoader

# 1. 创建NPU数据集（数据预处理在NPU上做）
dataset = NPUDataset(
    data_path="/shared/train_data.jsonl",
    tokenizer_path="THUDM/glm-4-9b",
    max_seq_len=8192,
    preprocess_on_npu=True,  # 关键：在NPU上做tokenize和pad
)

# 2. 创建NPU数据加载器
dataloader = NPUDataLoader(
    dataset,
    batch_size=4,
    shuffle=True,
    num_workers=2,          # 2个NPU数据预取worker
    prefetch_factor=4,      # 预取4个batch
    pin_memory=True,        # 锁页内存，加速Host→Device搬运
)

原理：NPUDataLoader在NPU上做tokenize和pad，避免CPU预处理瓶颈。同时用异步预取（prefetch），NPU计算当前batch时，下一个batch的数据已经在HBM里等着了。

效果：NPU利用率从35%提升到68%。

方案2：hccl优化配置

cann-recipes-train自动检测HCCS拓扑，配置最优通信路径：

# cann-recipes-train自动生成的hccl配置
export HCCL_CONNECT_TIMEOUT=1800
export HCCL_BUFFSIZE=128

# 关键：指定HCCS拓扑（4卡全互连）
# 生成rank_table.json（hccl用这个文件确定拓扑）
python -m cann_recipes_train.generate_rank_table \
    --npu_count 4 \
    --topology hccs_full_mesh

生成的rank_table.json告诉hccl：4张卡通过HCCS全互连，不要走PCIe。

效果：AllReduce 1.5GB梯度从9ms降到2.1ms，梯度同步时间占比从18%降到4%。

方案3：显存池化管理

cann-recipes-train用显存池替代PyTorch的默认显存分配器：

from cann_recipes_train import MemoryPoolConfig

# 启用显存池
config = MemoryPoolConfig(
    enable=True,
    pool_size=56,  # 预分配56GB，留8GB给临时缓冲
    fragmentation_threshold=0.05,  # 碎片率<5%
)

原理：预分配一大块HBM，所有tensor从池里分配，用完归还池子而不是释放给操作系统。这消除了碎片化，因为池子内部做紧凑化（compaction）。

效果：实际显存占用从58GB降到52GB，省了6GB可以加大batch_size。

实战：4卡微调GLM-4-9B

完整训练配置

from cann_recipes_train import Trainer, TrainingConfig

config = TrainingConfig(
    # 模型
    model_name="THUDM/glm-4-9b",
    dtype="bf16",
    
    # 数据
    data_path="/shared/train_data.jsonl",
    max_seq_len=8192,
    preprocess_on_npu=True,
    
    # 训练
    batch_size=4,                    # 每卡4个序列
    gradient_accumulation_steps=4,   # 梯度累积4步
    learning_rate=2e-5,
    weight_decay=0.01,
    max_steps=10000,
    warmup_steps=200,
    
    # 并行
    dp_degree=4,                     # 4卡数据并行
    
    # 优化（cann-recipes-train特有）
    use_flash_attention=True,        # 开启FlashAttention-2
    use_gradient_checkpointing=False, # 不开梯度检查点（4卡显存够）
    hccl_topology="hccs_full_mesh",  # HCCS全互连
    memory_pool_size=56,             # 显存池56GB
    
    # Checkpoint
    checkpoint_dir="/shared/checkpoints",
    save_interval=2000,
    async_save=True,                 # 异步保存
)

trainer = Trainer(config)
trainer.train()

性能优化过程

优化阶段	吞吐 (tokens/s)	NPU利用率	显存 (GB/卡)	梯度同步占比
Baseline（HuggingFace Trainer）	340	35%	58.2	18%
+ NPU数据加载器	580	68%	58.2	18%
+ hccl拓扑优化	870	78%	58.2	4%
+ 显存池化	1120	89%	52.1	4%

最终1120 tokens/s，比Baseline提升了3.3倍。

训练曲线

步数	Loss	吞吐	显存	备注
100	2.87	1120	52.1	初始下降
1000	1.94	1120	52.1	快速学习
5000	1.23	1120	52.1	收敛中
10000	0.98	1120	52.1	微调完成

踩坑实录

坑1：hccl默认走PCIe不走HCCS

问题：4卡AllReduce延迟9ms，远超预期2ms。

原因：hccl默认不检测HCCS拓扑，走PCIe通信（64GB/s vs HCCS 200GB/s）。

解决方案：用cann-recipes-train的generate_rank_table工具生成正确的拓扑文件：

# 检查当前hccl走的什么路径
python -c "import hccl; print(hccl.get_topology())"
# 如果输出PCIe而不是HCCS，需要生成rank_table

python -m cann_recipes_train.generate_rank_table --npu_count 4 --topology hccs_full_mesh
# 生成rank_table.json，hccl会自动读这个文件

坑2：梯度检查点反而降低吞吐

问题：开了use_gradient_checkpointing=True，吞吐反而降了15%。

原因：梯度检查点是"用计算换显存"——前向传播时不保存中间激活，反向传播时重新计算。在NPU上，重新计算的开销比GPU大（因为NPU的数据搬运延迟更高），得不偿失。

解决方案：4卡微调9B模型，显存够用（52GB/卡 < 64GB/卡），不需要梯度检查点。只有模型更大（>30B）或者序列更长（>32K）时才开。

坑3：数据预处理用CPU做是瓶颈

问题：用PyTorch的DataLoader + CPU tokenize，NPU利用率只有35%。

原因：CPU tokenize慢（~2000 tokens/s），NPU计算快（~10000 tokens/s），CPU喂不饱NPU。

解决方案：用cann-recipes-train的NPUDataLoader，tokenize在NPU上做：

# ❌ 慢：CPU tokenize
from torch.utils.data import DataLoader
dataloader = DataLoader(dataset, batch_size=4, num_workers=8)  # 8个CPU worker
# 吞吐：340 tokens/s

# ✅ 快：NPU tokenize
from cann_recipes_train import NPUDataLoader
dataloader = NPUDataLoader(dataset, batch_size=4, preprocess_on_npu=True)
# 吞吐：580 tokens/s

cann-recipes-train在CANN架构中的位置

cann-recipes-train位于CANN架构的应用层，依赖第1层AscendCL和第2层ATB/ops-transformer：

应用层：cann-recipes-train
  ↓ 调用
第1层：AscendCL（PyTorch → CANN后端）
  ↓ 调用
第2层：ATB（Transformer加速）、ops-transformer（FlashAttention）
  ↓ 调用
第4层：HCCL（分布式通信）、Runtime

结尾

cann-recipes-train把NPU上训大模型的坑都踩过了，跟着它的流水线走能省很多时间。三个最关键的优化：NPU数据加载（解决CPU瓶颈）、hccl拓扑配置（解决通信瓶颈）、显存池化（解决碎片化）。搞定这三项，NPU上的训练吞吐可以从340涨到1120 tokens/s，跟GPU持平。

如果你刚开始在NPU上训模型，不要直接用HuggingFace Trainer——先跑通cann-recipes-train的示例，理解每一步优化在做什么，然后再根据你的场景调整参数。踩过的坑比看过的文档有用。

https://atomgit.com/cann/cann-recipes-train