训练一个大模型需要同时管理三件事：把模型切成多块分到不同 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=8 train.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 训练循环的核心结构
for batch in dataloader:
    # 前向
    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）替开发者做了封装。

import torch.distributed as dist
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 的实测吞吐：

并行策略	每卡 Batch	吞吐 (tokens/s)	显存占用/卡
DP=8	4	8,500	31GB
TP=8	4	12,200	9.5GB
TP=4, PP=2, DP=1	4	11,800	12GB
TP=4, PP=2, DP=2	4	22,500	12GB

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=32 train.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