前言我所在团队要预训练一个7B参数的语言模型预算只够买4张Ascend 910B。原来以为昇腾NPU只能跑推理没想到torchtitan-npu直接支持大模型预训练——4卡跑7B模型训练速度1470 tokens/s两周训完100B tokens。这篇文章是完整的踩坑实录从环境搭建到性能调优每一步都记录在案。torchtitan-npu是Meta TorchTitan的昇腾NPU适配版核心改动是把CUDA后端替换成CANN后端上层PyTorch代码零修改。如果你用过TorchTitan在GPU上训模型迁移到NPU上只需要改一个环境变量。环境准备硬件要求组件最低配置推荐配置NPU4×Ascend 9108×Ascend 910CPU64核128核内存512GB1TB存储2TB NVMe SSD4TB NVMe SSD网络RDMARoCEv2HCCS RDMA⚠️ 踩坑预警RDMA网卡驱动版本必须跟CANN版本匹配不然hccl初始化会报HCCL_INIT_ERROR。我用的CANN 8.5 Mellanox OFED 5.9-0.5.6.0跑通没问题。软件安装# 1. 安装CANN 8.5必须先装chmodx Ascend-cann-toolkit_8.5.0_linux-x86_64.runsudo./Ascend-cann-toolkit_8.5.0_linux-x86_64.run--install# 2. 安装PyTorch 2.1 昇腾NPU版pipinstalltorch2.1.0 pipinstalltorch-npu2.1.0# 3. 克隆torchtitan-npugitclone https://atomgit.com/cann/torchtitan-npu.gitcdtorchtitan-npu pipinstall-rrequirements.txt pipinstall-e.Step 1配置训练参数torchtitan-npu用YAML配置文件管理所有训练参数不用改代码# train_config.yamlmodel:name:llama3_7bdim:4096n_layers:32n_heads:32vocab_size:128256max_seq_len:8192training:batch_size:2# 每卡batch sizeseq_len:8192gradient_accumulation_steps:8learning_rate:3e-4weight_decay:0.1max_steps:100000parallel:dp_degree:4# 4卡数据并行tp_degree:1# 不用张量并行7B模型单卡放得下dtype:bf16# BF16混合精度关键参数说明batch_size: 2每卡2个序列4卡总共8个序列梯度累积8步等于全局batch64dp_degree: 44卡数据并行每张卡跑完整模型dtype: bf16BF16比FP16训练更稳定指数位更多不容易溢出Step 2准备预训练数据torchtitan-npu要求数据格式是JSONL tokenizer# prepare_data.pyimportjsonfromtransformersimportAutoTokenizer tokenizerAutoTokenizer.from_pretrained(meta-llama/Llama-3-7B)# 把原始文本转成token ID并保存为JSONLinput_fileraw_corpus.txtoutput_filetrain_data.jsonlwithopen(input_file,r)asfin,open(output_file,w)asfout:forlineinfin:textline.strip()ifnottext:continuetokenstokenizer.encode(text,add_special_tokensTrue)# 截断到max_seq_leniflen(tokens)8192:tokenstokens[:8192]record{tokens:tokens,length:len(tokens)}fout.write(json.dumps(record)\n)⚠️ 踩坑预警数据文件要放在共享存储NFS上不然4张卡各读一份存储压力大3倍。RDMA环境推荐用WekaIO或者Lustre。Step 3启动训练# 4卡数据并行训练torchrun--nproc_per_node4\--nnodes1\--master_port29500\train.py\--configtrain_config.yaml\--data_dir/shared/train_data.jsonl\--checkpoint_dir/shared/checkpoints\--log_dir/shared/logs启动后看日志[Step 1] loss10.82, lr3e-4, throughput620 tokens/s, mem58.2GB/64GB [Step 10] loss9.45, lr3e-4, throughput620 tokens/s, mem58.2GB/64GB [Step 100] loss7.12, lr2.7e-4, throughput620 tokens/s, mem58.2GB/64GB620 tokens/s比预期低不少。下面开始调优。Step 4性能调优调优一开启FlashAttention-2FlashAttention-2是ops-transformer仓库提供的高性能Attention算子把Attention的显存占用从O(N²)降到O(N)同时减少HBM读写次数。# train_config.yaml 加一行model:use_flash_attention:true# 开启FlashAttention-2效果[Step 100] loss7.12, throughput980 tokens/s, mem42.1GB/64GB吞吐从620涨到980 tokens/s58%显存从58.2GB降到42.1GB-28%。调优二调整梯度累积步数原来gradient_accumulation_steps8每8步才做一次梯度同步。改成4步同步一次减少hccl通信次数training:gradient_accumulation_steps:4batch_size:4# 增大per-card batch size显存够用效果[Step 100] throughput1240 tokens/s, mem55.8GB/64GB吞吐从980涨到1240 tokens/s26%显存还在64GB以内。调优三hccl通信优化hccl的默认配置不一定最优。4卡HCCS互连场景下手动指定拓扑可以避免走PCIe# 设置hccl环境变量exportHCCL_CONNECT_TIMEOUT1800# 连接超时30分钟初次连接慢exportHCCL_BUFFSIZE128# 通信缓冲区128MBexportHCCL_WHITELIST_DISABLE1# 关闭白名单调试用exportASCEND_RT_VISIBLE_DEVICES0,1,2,3# 只用4张卡效果[Step 100] throughput1470 tokens/s调优结果汇总优化阶段吞吐 (tokens/s)显存 (GB)提升Baseline62058.2- FlashAttention-298042.158% 梯度累积batch调整124055.8100% hccl通信优化147055.8137%1470 tokens/s × 2周 约100B tokens的训练量满足7B模型预训练的需求。完整训练曲线训练2周后的Loss曲线训练步数Loss吞吐备注1,0006.831470初始快速下降10,0004.211470学到基本语法50,0002.871470学到常识知识100,0002.341470收敛中踩坑实录坑1梯度累积步数设太大Loss会震荡问题一开始设gradient_accumulation_steps16全局batch128Loss曲线震荡严重一步跳0.3。原因全局batch太大每次更新的梯度是128个序列的平均某个序列的噪声被放大了。解决方案降低到gradient_accumulation_steps4全局batch64Loss曲线平稳下降。坑2RDMA通信偶尔超时问题训练跑到第30步时hccl报HCCL_INTERNAL_ERROR: timeout进程挂掉。原因RDMA连接在网络抖动时会断开默认超时时间300s不够。解决方案exportHCCL_CONNECT_TIMEOUT1800# 连接超时改为30分钟exportHCCL_EXEC_TIMEOUT1800# 执行超时也改大坑3checkpoint保存时训练中断问题保存checkpoint时要写大量数据到磁盘每个checkpoint约28GB磁盘IO打满训练进程卡住。解决方案用异步保存训练和保存并行training:checkpoint:async_save:true# 异步保存save_interval:5000# 每5000步保存一次max_keep:3# 最多保留3个checkpointtorchtitan-npu在CANN架构中的位置torchtitan-npu位于CANN五层架构的应用层依赖第1层的AscendCL和第2层的ATB应用层torchtitan-npu ↓ 调用 第1层AscendCLPyTorch → CANN后端 ↓ 调用 第2层ATBTransformer加速库、ops-transformerFlashAttention ↓ 调用 第4层HCCL分布式通信、RuntimeNPU调度结尾NPU上跑大模型预训练不再是梦想torchtitan-npu让它变成了现实。从620到1470 tokens/s的调优过程告诉我NPU上做训练的关键不在模型代码在基础设施——FlashAttention减少计算量、hccl优化减少通信量、显存管理避免OOM。把这些基础设施配好NPU的训练效率可以跟GPU持平而硬件成本低30%。如果你的团队也在评估NPU上做大模型训练建议从torchtitan-npu的7B预训练开始4卡就能跑门槛不高。跑通7B再扩到70B也就是加卡和调并行策略的事。https://atomgit.com/cann/torchtitan-npu