torchtitan-npu：在昇腾集群上训练大模型

训练一个大模型需要同时管理三件事：把模型切成多块分到不同 NPU 上、在每块 NPU 上高效执行前向和反向、在 NPU 之间同步梯度。PyTorch 生态里做这件事通常有两个方案——PyTorch 原生的 DistributedDataParallel（DDP）或 Megatron-LM 的模型并行策略。

torchtitan-npu 是 CANN 社区维护的仓库，对标 Meta 的 torchtitan 项目——提供一套在昇腾NPU 上运行大模型训练的参考实现。它不重新发明分布式训练框架，而是在 PyTorch + HCCL 的基础上，给出 Transformer 模型在昇腾集群上的最佳实践。

torchtitan-npu 的定位

torchtitan-npu 不是一个独立的训练框架，它是一个参考实现仓库。它做的事情：

1. 提供 LLaMA 等主流 Transformer 模型在昇腾上的训练配置和脚本
2. 展示如何在昇腾上配置张量并行、流水线并行和数据并行
3. 集成 HCCL 作为分布式通信后端
4. 提供性能基准和调优建议

开发者不需要从零搭训练管线。torchtitan-npu 的配置文件直接指定了模型规模、并行策略、Batch 大小、学习率调度——拉到集群上就能跑。

# 训练 LLaMA-13B，8 卡张量并行torchrun--nproc_per_node=8train.py\--model=llama_13b\--tp_size=8\--pp_size=1\--dp_size=1\--batch_size=4\--max_length=4096

torchtitan-npu 的核心不是"让训练能跑起来"——PyTorch 本身就能让训练在昇腾上跑起来——而是"在昇腾上找到性能和扩展性的最优配置"。

为什么大模型训练需要并行化

一个 LLaMA-70B 模型：

• 参数量 70B，FP16 下约 140GB
• 激活值内存：Batch=4、序列长度 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 组数据并行（dp=4） 每组内 4 路张量并行（tp=4）——把单个 Transformer Block 切到 4 张卡上 每路 2 段流水线并行（pp=2）——把 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:# 前向loss=model(batch)# 反向loss.backward()# CANN Runtime 调度反向 Kernel# 梯度同步（HCCL AllReduce）optimizer.step()# 内部调用 torch.distributed.all_reduce# 梯度清零optimizer.zero_grad()

PyTorch 的DistributedDataParallel和torch.distributed在昇腾上的 backend 是hccl。所有分布式通信操作透明地走 HCCL。

HCCL 如何参与通信

在 torchtitan-npu 的训练脚本里，通常不需要直接调 HCCL API。PyTorch DDP 和 FSDP（Fully Sharded Data Parallel）替开发者做了封装。

importtorch.distributedasdist dist.init_process_group(backend="hccl")# 关键——指定 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_overlap=True会让 HCCL 在反向传播计算最后一层梯度时就开始通信，不用等所有层都算完。

训练性能参考

在 8×Ascend 910 集群上训练 LLaMA-13B 的实测吞吐：

TP=8 方案显存节省显著（从 31GB 降到 9.5GB），但通信开销限制了吞吐提升。TP+PP+DP 组合方案在 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,# 优化器状态不卸载到 CPU"gradient_sync_dtype":"fp16",# 梯度通信用 FP16}

激活值重计算是省显存最有效的手段。Transformer 训练时隐层的激活值占用了约 60% 的显存。开启激活值重计算后，前向只保留部分层的激活值，反向传播到未保留的层时重新计算前向结果。时间换空间，显存降低约 30% 的同时计算量增加约 15%。

DeepSpeed 集成

对于更大规模的模型（70B+），torchtitan-npu 还支持集成 DeepSpeed 的 ZeRO 优化：

torchrun--nproc_per_node=32train.py\--zero_stage=3\--offload_optimizer=cpu

ZeRO-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 scale"loss_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