跳到主要内容
昇腾AI芯片与算子开发

4.3 昇腾分布式训练实战

用 torch_npu（PyTorch）在昇腾 NPU 上跑通数据并行训练，并介绍张量并行/流水线并行在昇腾上的配置方式，构建昇腾分布式训练实战能力

分布式训练 torch_npu 数据并行 3D并行

原语和配置都讲完了，现在真正动手：让代码在 8 卡甚至更多卡上跑起来。这一章聚焦昇腾分布式训练的主流实践路径——PyTorch（torch_npu），并介绍 3D 并行在昇腾上的落地方式。

📑 目录

1. 昇腾分布式训练的主路径：torch_npu

昇腾上做分布式训练，当前的主流实践路径是 torch_npu（PyTorch）：存量 PyTorch 代码只需把通信后端换成 HCCL、设备换成 NPU 即可跑通，生态与工具链最完整（DDP/FSDP、Megatron 风格的 TP/PP 均可落地），也是实际使用最广的路线。

补充说明：华为自研的 MindSpore 框架也支持昇腾分布式训练（内置 HCCL 接口、声明式并行），但生态与代码迁移上仍以 torch_npu 为主线，故本节只展开 torch_npu 路径。

2. torch_npu 数据并行实战（DDP）

核心改动只有三处：初始化进程组、分配本地 rank、把模型和数据放到对应 NPU。

import torch
import torch.distributed as dist
import torch_npu  # 引入 npu 设备支持

def main():
    # 1. 初始化进程组：后端用 hccl（对应 NCCL）
    dist.init_process_group(backend='hccl', init_method='env://')

    local_rank = int(os.environ['LOCAL_RANK'])
    torch.npu.set_device(local_rank)          # 每进程绑定一张卡

    # 2. 模型包装成 DDP（device_ids 指向本地卡）
    model = TinyNet().to(f'npu:{local_rank}')
    ddp_model = DDP(model, device_ids=[local_rank])

    # 3. 数据加载：每进程取各自分片
    dataset = make_dataset()
    sampler = torch.utils.data.distributed.DistributedSampler(dataset)
    loader = DataLoader(dataset, sampler=sampler, batch_size=32)

    # 4. 训练循环与单卡几乎一样
    for epoch in range(3):
        loader.sampler.set_epoch(epoch)
        for x, y in loader:
            x, y = x.to(f'npu:{local_rank}'), y.to(f'npu:{local_rank}')
            out = ddp_model(x)
            loss = loss_fn(out, y)
            opt.zero_grad()
            loss.backward()
            opt.step()

if __name__ == '__main__':
    main()

启动命令：

# 单机 8 卡
torchrun --nproc_per_node=8 train_ddp.py

# 多机（每台机器上以对应 node_rank 启动）
torchrun --nproc_per_node=8 --nnodes=2 --node_rank=0 \
  --master_addr=MASTER_IP --master_port=29500 train_ddp.py

要点：backend='hccl' 是昇腾专属改动；其余 DDP 逻辑与 CUDA 环境完全一致。

3. 3D 并行在昇腾上的落地

大模型（百亿/千亿参数）训练需要 3D 并行 = 数据并行（DP）+ 张量并行（TP）+ 流水线并行（PP）：

维度昇腾上的实现通信模式
DP梯度 AllReduce（HCCL）每步一次，机内优先
TP层内切分，AllGather/ReduceScatter高频小通信，必须尽量放机内（HCCS）
PP层间切分，点对点传输低频大通信，可跨机

昇腾上落地 3D 并行：

  • torch_npu 上的实现：可以迁移 Megatron-LM 风格的 TP/PP 代码（把 nccl 换成 hccl 后端即可），社区也有适配方案。
  • 关键工程建议：把 TP 组放在同一台机器内——TP 的通信频次远高于 DP，跨机会拖垮整体性能。

FSDP / ZeRO 在昇腾：torch_npu 生态支持 FSDP（backend='hccl'），其梯度/参数分片通信对应 HCCL 的 ReduceScatter/AllGather，思路与 GPU 环境完全一致。

4. 分布式训练的性能检查清单

训练跑起来之后，用这份清单快速体检：

  1. 通信占比：训练一个 step 的总时间中，通信耗时占比多少？（用 profiler 查看）
  2. 机内/跨机流量：确认高频通信（TP/梯度同步）主要走 HCCS。
  3. 多卡利用率：各卡利用率是否均衡（npu-smi info 看 AI Core 利用率）。
  4. DataLoader 是否瓶颈：CPU 预处理是否跟不上（NPU 计算快时尤其明显）。
  5. 梯度归约开销：大数据并行组数多时，AllReduce 占比是否偏高（可考虑梯度压缩/延迟同步技巧）。

一句话原则：先用小规模验证正确性与性能模型，再放大到全集群。分布式训练的绝大多数坑，都是「先小后大」能避开的。

总结

  • 昇腾分布式训练主路径：torch_npu（backend=‘hccl’）。
  • torch_npu DDP 迁移成本极低：初始化进程组换后端 + 设备换成 npu 即可。
  • 3D 并行落地要点：TP 组必须尽量放机内（HCCS），PP 可跨机，DP 梯度同步走 HCCL。
  • 性能体检四件套：通信占比、机内/跨机流量、多卡利用率、DataLoader 吞吐。

延伸阅读