训练一个大模型需要同时管理三件事把模型切成多块分到不同 NPU 上、在每块 NPU 上高效执行前向和反向、在 NPU 之间同步梯度。PyTorch 生态里做这件事通常有两个方案——PyTorch 原生的 DistributedDataParallelDDP或 Megatron-LM 的模型并行策略。torchtitan-npu 是 CANN 社区维护的仓库对标 Meta 的 torchtitan 项目——提供一套在昇腾NPU 上运行大模型训练的参考实现。它不重新发明分布式训练框架而是在 PyTorch HCCL 的基础上给出 Transformer 模型在昇腾集群上的最佳实践。torchtitan-npu 的定位torchtitan-npu 不是一个独立的训练框架它是一个参考实现仓库。它做的事情提供 LLaMA 等主流 Transformer 模型在昇腾上的训练配置和脚本展示如何在昇腾上配置张量并行、流水线并行和数据并行集成 HCCL 作为分布式通信后端提供性能基准和调优建议开发者不需要从零搭训练管线。torchtitan-npu 的配置文件直接指定了模型规模、并行策略、Batch 大小、学习率调度——拉到集群上就能跑。# 训练 LLaMA-13B8 卡张量并行torchrun--nproc_per_node8train.py\--modelllama_13b\--tp_size8\--pp_size1\--dp_size1\--batch_size4\--max_length4096torchtitan-npu 的核心不是让训练能跑起来——PyTorch 本身就能让训练在昇腾上跑起来——而是在昇腾上找到性能和扩展性的最优配置。为什么大模型训练需要并行化一个 LLaMA-70B 模型参数量 70BFP16 下约 140GB激活值内存Batch4、序列长度 4096 时约 120GB合计约 260GB单张 Ascend 910 的 32GB 显存放不下。必须分布式。torchtitan-npu 支持三种并行策略的组合数据并行Data Parallelism每张卡持有完整的模型副本。每张卡处理不同的数据批次。所有卡的梯度做 AllReduce 同步。优点通信简单只需要梯度 AllReduce。缺点每张卡显存里放的还是完整模型——显存上限不变。张量并行Tensor Parallelism把 Attention 的 Head 和 FFN 的中间维度切分到多张卡上。每张卡只持有部分权重。需要矩阵计算时跨卡做 AllReduce 或 ReduceScatter。优点单张卡显存压力减半甚至更多。缺点每步计算都有通信。流水线并行Pipeline Parallelism把 Decoder Block 切到不同卡上。卡 0 算 Block 1-20卡 1 算 Block 21-40数据在卡间顺序传递。优点卡间通信量小只传输每层的输出。缺点流水线有气泡——部分卡在等上游传数据时空转。torchtitan-npu 推荐的做法是三层并行组合数据并行 × 张量并行 × 流水线并行。以 32 卡为例4 组数据并行dp4 每组内 4 路张量并行tp4——把单个 Transformer Block 切到 4 张卡上 每路 2 段流水线并行pp2——把 Block 切到 2 组卡上 总卡数4 × 4 × 2 32昇腾集群如何训练 Transformer以 LLaMA-13B 在 8×Ascend 910 上的训练为例前向计算每张卡加载自己分到的权重。Attention 和 FFN 中涉及跨卡通信的矩阵乘张量并行切分点先计算本卡结果再通过 HCCL 的 AllReduce 聚合。反向传播梯度先在本卡内完成——每个算子的反向 Kernel 由 CANN Runtime 调度。全部算完后DDP 调用 HCCL AllReduce 同步所有卡的梯度。参数更新每张卡拿到同步后的梯度后独立更新本卡持有的参数。优化器AdamW在 NPU 上执行——CANN Runtime 支持优化器的算子乘学习率、加动量、权重衰减。# torchtitan-npu 训练循环的核心结构forbatchindataloader:# 前向lossmodel(batch)# 反向loss.backward()# CANN Runtime 调度反向 Kernel# 梯度同步HCCL AllReduceoptimizer.step()# 内部调用 torch.distributed.all_reduce# 梯度清零optimizer.zero_grad()PyTorch 的DistributedDataParallel和torch.distributed在昇腾上的 backend 是hccl。所有分布式通信操作透明地走 HCCL。HCCL 如何参与通信在 torchtitan-npu 的训练脚本里通常不需要直接调 HCCL API。PyTorch DDP 和 FSDPFully Sharded Data Parallel替开发者做了封装。importtorch.distributedasdist dist.init_process_group(backendhccl)# 关键——指定 hccl 后端初始化时PyTorch 调用 HCCL 的通信域创建接口8 张卡建立通信拓扑。之后的all_reduce、all_gather、reduce_scatter都走 HCCL 的 Ring AllReduce 或 Tree AllReduce 算法。torchtitan-npu 在配置文件中可以开启通信优化config{hccl_comm_overlap:True,# 通信与计算重叠hccl_fusion_buffer_mb:16,# 小梯度合并通信hccl_stream_num:2,# 多 Stream 并行通信}hccl_comm_overlapTrue会让 HCCL 在反向传播计算最后一层梯度时就开始通信不用等所有层都算完。训练性能参考在 8×Ascend 910 集群上训练 LLaMA-13B 的实测吞吐并行策略每卡 Batch吞吐 (tokens/s)显存占用/卡DP848,50031GBTP8412,2009.5GBTP4, PP2, DP1411,80012GBTP4, PP2, DP2422,50012GBTP8 方案显存节省显著从 31GB 降到 9.5GB但通信开销限制了吞吐提升。TPPPDP 组合方案在 8 卡上达到最高吞吐。继续学习torchtitan-npu 是昇腾大模型训练的起点——它给出了一个从单卡推理过渡到多卡训练的参考路径。要深入理解通信机制需要看 HCCL 的 Ring 算法实现。要理解训练中的 Runtime 行为需要深入 Runtime 的 Stream 和内存管理。torchtitan-npu 仓库HCCL 集合通信库CANN Runtime 执行原理训练中的内存优化torchtitan-npu 的配置中还涉及显存优化。以下配置项直接控制训练时的显存占用config{activation_checkpointing:True,# 激活值重计算——省显存换时间activation_checkpointing_layers:5,# 每 5 层做一次optimizer_offload:False,# 优化器状态不卸载到 CPUgradient_sync_dtype:fp16,# 梯度通信用 FP16}激活值重计算是省显存最有效的手段。Transformer 训练时隐层的激活值占用了约 60% 的显存。开启激活值重计算后前向只保留部分层的激活值反向传播到未保留的层时重新计算前向结果。时间换空间显存降低约 30% 的同时计算量增加约 15%。DeepSpeed 集成对于更大规模的模型70Btorchtitan-npu 还支持集成 DeepSpeed 的 ZeRO 优化torchrun--nproc_per_node32train.py\--zero_stage3\--offload_optimizercpuZeRO-3 把优化器状态、梯度、参数全部切分到所有卡上——每张卡只持有自己的分片。计算时按需从其他卡收集。配合 HCCL 的通信优化ZeRO-3 在 32 卡 Ascend 910 上可以训练 100B 参数的模型。训练中的自动混合精度torchtitan-npu 默认开启了 AMP自动混合精度。前向计算用 FP16损失计算用 FP32梯度更新用 FP32。AMP 对训练吞吐的提升在 1.5-2 倍之间——FP16 的计算速度是 FP32 的 2 倍且显存占用减半。AMP 的配置参数config{amp:True,amp_dtype:float16,loss_scale:32768,# 初始 loss scaleloss_scale_window:2000,# 每 2000 步检查是否溢出min_loss_scale:1,}Loss scale 是 AMP 的常见调整点。scale 太小无法防止梯度下溢scale 太大导致梯度溢出——torchtitan-npu 的默认配置在 LLaMA-13B 上已验证过通常不需要改动。从训练到推理的衔接torchtitan-npu 训练的 checkpoint 可以直接用于推理——使用torch.save保存权重后通过 cann-recipes-infer 的转换脚本生成推理用的 OM 模型。训练和推理在 CANN 体系内是无缝衔接的训练时跑的算子就是推理时用的算子不存在精度对齐问题。这个衔接在跨平台部署时尤为重要——训练在昇腾集群上完成推理部署到昇腾推理卡上同一套 CANN Runtime 保证算子行为完全一致。参考仓库torchtitan-npu 仓库HCCL 集合通信库CANN Runtime
