训练环境与分布式启动

标签

AI/infra/训练工程

字数

2397 字

阅读时间

10 分钟

02-分布式训练总论与显存账本讲的是「为什么要并行」，这一篇讲怎么把多卡真正跑起来 —— 进程模型、编号体系、启动器，以及卡住时怎么排查。

这里的每一个概念都在后面的 3D 并行编排里被反复用到，不是可跳过的操作步骤。

三种规模与带宽分层

规模	典型配置	互联	场景
单机单卡	1 GPU	—	调试、原型
单机多卡	8 GPU / 节点	NVLink / NVSwitch	中等模型、张量并行
多机多卡	N 节点 × 8 GPU	InfiniBand / RoCE	百亿–千亿参数

关键认知是带宽分层：机内 NVLink 数百 GB/s，跨机 IB 只有机内的几分之一到十分之一（具体数字见 03-多卡互联与集群网络）。

一栋楼里的同事可以走过去当面讨论，跨城市的团队只能开视频会议 —— 能当面说的事尽量别开远程会，这是并行策略布局的基本直觉。

一卡一进程

PyTorch 推荐的模型是一块 GPU 对应一个独立操作系统进程。

为什么不用单进程多线程管 8 块卡？ 两个原因：

原因	说明
绕开 GIL	Python 全局解释器锁让同进程内的多线程无法真正并行执行字节码。用多线程驱动多卡，调度逻辑会被 GIL 串行化，多卡很难喂饱。多进程各有各的解释器
独占 CUDA 上下文	每个进程绑定一块 GPU，拥有干净的 CUDA 上下文与显存空间，崩溃只影响单个进程

早期的 DataParallel 正是栽在 GIL 与主卡负载不均上（见 06-数据并行：DP、DDP 与 FSDP），所以现代训练全部转向「一卡一进程」的 DDP。

三个编号：最容易搞错的点

Node 0 (node_rank=0)          Node 1 (node_rank=1)
  进程 rank=0, local_rank=0 → GPU 0    进程 rank=2, local_rank=0 → GPU 0
  进程 rank=1, local_rank=1 → GPU 1    进程 rank=3, local_rank=1 → GPU 1

概念	定义
`rank`	全局唯一编号， $0 \sim world_size - 1$ 。`rank == 0` 常用来指定主进程（打印日志、存模型）
`local_rank`	进程在所在节点内部的编号 —— 它决定绑定哪块 GPU
`world_size`	总进程数 = 节点数 × 每节点进程数

rank = node_rank \times G + local_rank, world_size = nnodes \times G

（ $G$ = 每节点进程数。）

最常见的错误：用 `rank` 去 `set_device`

上图里 Node 1 的两个进程全局 rank 是 2、3，但 local_rank 重新从 0 开始数。用 rank=2 去 set_device，在第二台机器上根本没有「第 2 块卡」—— 轻则报错，重则绑错设备。

决定 GPU 绑定的永远是 local_rank。

两个节点、每节点 2 进程时的编号体系：

   Node 0（node_rank = 0）             Node 1（node_rank = 1）
   ┌────────────────────────────┐      ┌────────────────────────────┐
   │ rank 0   local_rank 0 → GPU0│      │ rank 2   local_rank 0 → GPU0│
   │ rank 1   local_rank 1 → GPU1│      │ rank 3   local_rank 1 → GPU1│
   └────────────────────────────┘      └────────────────────────────┘
                world_size = 2 节点 × 2 进程 = 4

  ⇒ rank 全局唯一（0 … world_size−1），rank 0 常用来干「只需要一份」的事
  ⇒ local_rank 在节点内重新从 0 开始数 —— 它才决定绑哪块 GPU
  ⇒ 拿 rank = 2 去 set_device，在第二台机器上根本没有「第 2 块卡」

进程组

有了编号，进程之间还需要一条约定好的通信通道 —— 进程组（Process Group）。任何分布式通信之前都必须先 init_process_group。

python

def setup():
    rank       = int(os.environ["RANK"])         # torchrun 注入
    local_rank = int(os.environ["LOCAL_RANK"])
    world_size = int(os.environ["WORLD_SIZE"])

    torch.cuda.set_device(local_rank)            # 用 local_rank 绑定！
    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)

后端选择：

后端	用于
`nccl`	GPU 训练（走 NVLink / IB）
`gloo`	CPU 通信、调试

子通信组：为 3D 并行铺路

dist.new_group(ranks=[...]) 创建子通信组。这不是可选的高级技巧 —— 每个并行维度都对应一组 NCCL communicator：

TP 组（机内 8 卡）、PP 组、DP 组各自是独立的通信域
world_size = TP × PP × DP 的约束，在实现上就是「怎么把 rank 切分成这些组」（见 13-3D 并行与混合并行策略）

world_size = TP × PP × DP 在实现上就是「把 rank 切成这些组」：

  以 TP=2、PP=2、DP=2（共 8 个 rank）为例：

       rank:    0     1     2     3     4     5     6     7
                │     │     │     │     │     │     │     │
   TP 组:       └──0──┘     └──1──┘     └──2──┘     └──3──┘      ← 机内，走 NVLink
   PP 组:       └─────────0─────────┘     └─────────1─────────┘  ← 跨机，走 IB
   DP 组:       └───────────────0───────────────┘  └────1────┘   ← 跨节点，要求最低

  ⇒ 每个并行维度对应一组独立的 NCCL communicator（dist.new_group）
  ⇒ 维度越多，创建的通信组越多，初始化开销也越大

为什么需要专门的启动器

多进程训练要解决一堆麻烦：谁先启动、怎么互相找到、端口怎么分配、某个进程挂了怎么办。torchrun 就是干这件事的。

它相比早期的 torch.distributed.launch，多了弹性与容错能力（Rendezvous 机制：节点可以中途加入或重启，靠 c10d 重新会合）。

torchrun 注入的环境变量：

变量	含义
`RANK`	全局 rank
`LOCAL_RANK`	节点内 rank
`WORLD_SIZE`	总进程数
`MASTER_ADDR` / `MASTER_PORT`	主节点地址与端口（所有进程靠它互相找到）

bash

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

# 多机：每台机器各跑一次，参数一致
torchrun --nnodes=2 --nproc_per_node=8 \
         --node_rank=0 --master_addr=10.0.0.1 --master_port=29500 \
         train.py

MASTER_ADDR / MASTER_PORT 是所有进程的会合点 —— 多机训练起不来，第一件事就是查这两个值在各节点上是否一致、端口是否被占、防火墙是否放行。

NCCL 调优与排障

NCCL 的传输路径与拓扑感知机制在 03-集合通信与 NCCL 有完整说明。这里只讲卡住时怎么办。

排查第一招

bash

NCCL_DEBUG=INFO python train.py

它会把拓扑识别结果与算法选择打出来 —— 多机训练起不来时，第一件事就是开它。

指定网卡与 IB 设备

变量	作用
`NCCL_SOCKET_IFNAME`	指定 TCP/IP 接口（如 `eth0`；`^lo` 表示排除回环）
`NCCL_IB_HCA`	指定 IB 网卡（如 `mlx5_0`）
`NCCL_NET_GDR_LEVEL`	GPUDirect RDMA 的拓扑级别

机器上有多张网卡时，NCCL 选错网卡会让性能掉到原理论值的一个零头。 这两个变量是多机训练最常见的必设项，也是「配置都对、就是慢」的头号原因。

二分排障开关

变量	关掉什么
`NCCL_P2P_DISABLE=1`	禁用 P2P（NVLink / PCIe 直连）
`NCCL_SHM_DISABLE=1`	禁用共享内存传输

「二分」的意思是：关掉某一路径后如果能跑通，说明问题出在那条路径上（比如 NVLink 拓扑异常或 /dev/shm 太小）。这是把「不知道哪里坏」缩小到「某一层坏」的最快手段。

设备绑定

CUDA_VISIBLE_DEVICES 决定进程能看到哪些卡。多进程训练里由 torchrun 与 set_device 共同决定绑定关系 —— 手动再设 CUDA_VISIBLE_DEVICES 容易与框架的编号体系冲突。

最小 DDP 脚本的五段式

python

# ① init
dist.init_process_group(backend="nccl")
torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))

# ② Sampler：让每卡只读自己那部分数据
sampler = DistributedSampler(dataset, shuffle=True)
loader  = DataLoader(dataset, batch_size=..., sampler=sampler)

# ③ DDP 包裹
model = DDP(model.to(local_rank), device_ids=[local_rank])

# ④ 训练循环
for epoch in range(epochs):
    sampler.set_epoch(epoch)          # 每轮换切分方式，否则各卡看到的数据完全相同
    for batch in loader:
        ...

# ⑤ destroy
dist.destroy_process_group()

三个容易被忽略的点：

点	为什么
`sampler.set_epoch(epoch)`	`DistributedSampler` 默认按 epoch 决定切分方式。不调用它，每个 epoch 各卡看到的都是同一份分片 —— 数据打乱的意义全没了
`sampler=` 而不是 `shuffle=True`	用了 `shuffle=True` 会与 sampler 冲突。数据并行下每卡必须只读自己那份
仅 rank 0 干的事要加判断	日志、存 checkpoint、TensorBoard 写盘 —— 不加 `if rank == 0` 会导致 8 个进程同时写同一个文件

rank == 0 的惯例不是形式主义：8 个进程同时写同一个 checkpoint 会互相覆盖甚至损坏文件；同时打日志会把输出刷成 8 倍。这类问题在训练跑起来之后才发现，但排查成本很高。

常见踩坑

症状	优先查
多机起不来 / 卡在 init	`MASTER_ADDR` / `MASTER_PORT` 一致性、端口占用、防火墙
各卡读到相同数据	忘了 `sampler.set_epoch(epoch)`
卡住不报错	`NCCL_DEBUG=INFO` 看卡在哪一步；检查是否所有 rank 都到了同一点（`__syncthreads` 类的死锁同理）
性能远低于预期	`NCCL_SOCKET_IFNAME` / `NCCL_IB_HCA` 是否选对网卡
主卡 OOM 其他卡空闲	用了 `DataParallel` 而不是 DDP
checkpoint 损坏	多个 rank 同时写盘

参考

贡献者

文件历史

最后编辑于大约 18 小时前查看完整历史

语法基础

gpu

cuda

train

infer

训练环境与分布式启动

三种规模与带宽分层

一卡一进程

三个编号：最容易搞错的点

进程组

子通信组：为 3D 并行铺路

为什么需要专门的启动器

NCCL 调优与排障

排查第一招

指定网卡与 IB 设备

二分排障开关

设备绑定

最小 DDP 脚本的五段式

常见踩坑

相关

参考

贡献者

文件历史

训练环境与分布式启动 ​

三种规模与带宽分层 ​

一卡一进程 ​

三个编号：最容易搞错的点 ​

进程组 ​

子通信组：为 3D 并行铺路 ​

为什么需要专门的启动器 ​

NCCL 调优与排障 ​

排查第一招 ​

指定网卡与 IB 设备 ​

二分排障开关 ​

设备绑定 ​

最小 DDP 脚本的五段式 ​

常见踩坑 ​

相关 ​

参考 ​

贡献者 ​

文件历史 ​

训练环境与分布式启动

三种规模与带宽分层

一卡一进程

三个编号：最容易搞错的点

进程组

子通信组：为 3D 并行铺路

为什么需要专门的启动器

NCCL 调优与排障

排查第一招

指定网卡与 IB 设备

二分排障开关

设备绑定

最小 DDP 脚本的五段式

常见踩坑

相关

参考

贡献者

文件历史