多卡多机训练任务调度系统：从梯度平均到拓扑感知调度

一个任务要占 8 台机器 64 张卡，少一张都开不了；另一个模型单卡装不下，必须把一层矩阵劈成四份分别算。这两类需求背后是同一套东西：集合通信 和调度器的编排能力。本文把「为什么要这么切」「切完怎么同步」「调度器到底管什么、不管什么」一层层拆开，最后给出一套可落地的调度器设计。

---

核心结论速览

1. 梯度「取平均」不是调度器干的活，是 NCCL 的 AllReduce 干的活。调度器的任务是：保证这一组进程能被同时创建出来、彼此能连通、并且知道自己是几号。

2. 切数据（DP）和切模型（TP/PP）是两类完全不同的通信模式。前者每个 step 只同步一次、可以跨机；后者每层都要同步、必须待在同一个 NVLink 域里。调度器要把它们放到对的物理位置。

3. 「8 台机器少 1 张卡就全等」这件事叫 gang scheduling。不做这个，集群必然出现一半 worker 占着卡干等、另一半永远排不上的死锁。

4. 调度器与训练框架之间靠一份「契约」通信 ：注入 MASTER_ADDR / WORLD_SIZE / RANK / LOCAL_RANK 等环境变量，然后放手不管。

01 问题从哪来

先说清楚我们要解决的两个真实痛点，后面所有设计都是为它们服务的。

一张显卡上有两块资源是分开的：算力 和显存。训练会同时撞上这两堵墙，但撞墙的原因不一样，对应的解法也完全不同：

• 一张卡太慢（墙一：算力不够） ：模型能装下，但跑一个 epoch 要三天。解法是把同一份模型复制多份，喂不同的数据并行跑 ------ 也就是切数据。
• 一张卡装不下（墙二：显存不够） ：模型本身 + 梯度 + 优化器状态就超过 80GB。解法是把模型本身劈开分到多张卡上 ------ 也就是切模型。

你描述的两种场景正好一一对应：

• 「多个显卡同时训练一个任务，梯度同步用平均值」→ 这是数据并行的经典描述。
• 「一个模型需要很大的显存，需要多张卡来训练」→ 这是模型并行（张量并行 / 流水线并行）的经典描述。

而现实中，规模上去以后这两种会同时用。所以本文的脉络是：先讲清三种切法各自的脾气（第 2 节），再讲它们怎么同步（第 3 节），然后讲显存怎么省（第 4 节），最后------也是最重点的------讲调度系统如何把这些约束翻译成「哪几张卡、放在哪个机架、什么时候启动」（第 5 节往后）。

💡 先记住一个分界线

训练框架（PyTorch / Megatron / DeepSpeed）负责「怎么算」------怎么劈矩阵、怎么调 AllReduce、怎么存 checkpoint。

调度系统负责「在哪儿算」------给哪个任务分配哪几张物理卡、什么时候一起启动、失败了怎么重启。

两者之间只有一份很窄的接口。绝大多数调度系统写复杂了，都是因为越界去管了「怎么算」。

02 三种切法：数据并行、张量并行、流水线并行

用三个生活化类比先建立直觉，再看它们的通信特征差异。

2.1 数据并行（DP）：复印机模式

每张卡上都放一份完整的模型副本，把一批数据切成 N 份，每张卡吃一份。各算各的梯度，算完把梯度汇总取平均，再各自用同一个平均梯度更新参数。

关键点：省不了任何显存（每张卡都得装下整个模型），但通信结构最简单------每个训练 step 只同步一次梯度。

2.2 张量并行（TP）：把一层劈开

把模型某一层内部 的矩阵乘法拆到多张卡上。比如一个全连接层 Y = XW，把权重矩阵 W 按列切成 4 份，4 张卡各算 1/4，再把结果拼起来。

关键点：省显存效果立竿见影（单层切 4 份，每卡只存 1/4），但通信极其频繁------每过一个 Transformer 层就要 2 次 AllReduce（前向）+ 2 次（反向）。

2.3 流水线并行（PP）：工厂流水线

把模型的层按顺序切成几段，每台机器负责一段。数据像半成品一样从前一段流向后一段：第 1 段算完把激活值（中间结果）发给第 2 段，自己回头算下一批。

关键点：通信量最小（是点对点传激活值，不是全员 AllReduce），但会产生气泡------流水线切换时有的阶段在空转等上游。

三种并行策略：切数据、切单层权重、切模型深度。颜色差异对应通信特征的差异。

2.4 一张表看清差异

维度	数据并行 DP	张量并行 TP	流水线并行 PP
切什么	输入数据	单层内的矩阵	模型层（纵向切段）
每卡显存	完整模型，不省	权重 ÷ TP 度	参数 ÷ PP 度
通信原语	AllReduce	AllReduce / AllGather	点对点 Send/Recv
通信频率	每 step 1 次	每层 4 次	每段 1 次
带宽要求	中（可容忍 IB / 以太网）	极高（必须 NVLink）	中（可容忍 IB）
典型放置	跨机、跨机架	同机 8 卡（或 NVL72 机柜）	跨机
主要代价	显存冗余	通信开销	流水线气泡

⚠️ 为什么 TP 一旦跨机就是灾难

算一笔账：模型 hidden size 8192、batch 4、序列长度 8192、bf16 精度，单次 AllReduce 要搬 4 × 8192 × 8192 × 2B = 512 MB。80 层、每层 4 次，一个 iteration 就是 160 GB 的 TP 通信量。

同样的数据，走 NVLink 是 900 GB/s，走 NDR InfiniBand（400 Gb/s ≈ 50 GB/s）要慢 18 倍。这 160GB 在 NVLink 上大约 0.18 秒，在 IB 上要 3.2 秒 ------ 训练时间直接翻好几倍。

所以业界的共识是：TP 的度不要超过单个 NVLink 域的大小。传统 HGX H100 域内是 8 卡，所以 TP ≤ 8；GB200 NVL72 把域扩展到 72 卡，TP 才能做到 16 以上。

2.5 组合起来：3D 并行

规模上去后这三种一定是叠加使用的，业内叫「3D 并行」，配置方式是一句乘法：

总卡数 = TP 度 × PP 度 × DP 度

比如 175B 参数模型用 TP=8, PP=8, DP=16，一共需要 1024 张卡。它的物理含义是：先用 TP 把「一层」变小到能塞进单机 8 卡，再用 PP 把「整个模型」分成 8 段分到 8 台机器，最后用 DP 复制 16 份跑不同数据。

模型	卡数	TP	PP	DP	框架
GPT-3 175B	1024 × V100	8	8	16	Megatron-LM
Megatron-Turing NLG 530B	2240 × A100	8	35	8	Megatron-DeepSpeed
LLaMA 3 405B	16384 × H100	8	16	128	Meta 内部栈

💡 这对调度系统意味着什么

3D 并行把一个「要 1024 张卡」的抽象需求，翻译成了一个带形状的资源申请：不是「我要 1024 张卡」，而是「我要 128 组，每组 8 台机器，每台上 TP=8 的 8 张卡互连」。

调度器如果只知道「1024 张卡」这个数字，就会把 TP 组拆散到不同机架上，训练能跑起来但性能腰斩。形状信息必须传给调度器------这就是第 7 节拓扑感知调度要解决的问题。

03 梯度为什么是「平均值」

这是你问题里的核心，也是很多讲解含糊带过的地方。要讲清三件事：算法怎么做、为什么这么做是对的、为什么必须除以 N。

3.1 先回答：为什么要同步

数据并行时，每张卡看到的是一批不同 的数据，所以算出的是不同的梯度。如果各用各的梯度更新参数，训练几步之后这几个模型副本就会漂移成几个不同的模型------那你训练的不是「一个大模型」，而是「N 个小模型」。

所以必须每步把梯度汇总成一个统一的值，让所有副本做完全相同 的更新。这个「汇总」的动作，就是集合通信里的 AllReduce。

AllReduce 是什么

一句话：N 个进程各有一份长度相同的向量，AllReduce 把对应位置的元素做规约运算（求和 / 求最大值等），然后把结果发还给每一个进程。「All」的意思是每个人都拿到完整结果，不是只发给一个人。

3.2 Ring AllReduce：绕圈传为什么反而最快

最朴素的实现是「星型汇聚」：所有卡把梯度发给 0 号卡，0 号卡求和后再广播回去。问题是 0 号卡的网卡带宽成了瓶颈，卡越多越慢。

Ring AllReduce 的做法完全不同：把所有卡连成一个逻辑环，每张卡只跟左右邻居说话，数据切成 N 块在环上流转，分两个阶段完成。

Ring AllReduce 的两个阶段。阶段一「分散规约」让每张卡各自算完一块；阶段二「全收集」把这些块绕圈广播，最后人人持有完整结果。

为什么它是带宽最优的

关键在于「每张卡实际发了多少数据」这个数字。设梯度总量 S、卡数 N：

每张卡发送量 = 2 × (N-1)/N × S

N = 8   → 1.75 S
N = 100 → 1.98 S
N → ∞   → 2 S（理论上限）

看出重点了吗：每卡的通信量不随卡数增长。从 8 张卡加到 100 张卡，每卡要发的数据只从 1.75S 涨到 1.98S。这就是数据并行能扩展到上千卡的根本原因。Patarasuk 和 Yuan 证明了这个构造是带宽最优的------没有任何 AllReduce 算法能让每卡流量低于这个值。

💡 代价是延迟

Ring 的步数是 2(N-1)，随卡数线增长。所以小张量用它反而不划算------NCCL 会在消息很小的时候自动换成树形算法（延迟 O(log N)）。梯度是 GB 级大张量，带宽占主导，所以用 Ring。

3.3 从「求和」到「平均」：那个除号在哪

注意一个容易搞混的细节：NCCL 的 AllReduce 默认做的是 SUM（求和），不是 AVG（平均） 。你在框架层面看到的「梯度平均」，是框架在求和之后又除以了 N。

验证一下：两个进程，各自造一个标量，看 AllReduce 之后是什么

# check_allreduce.py ------ 用两张卡跑：torchrun --nproc-per-node=2 check_allreduce.py
import torch, torch.distributed as dist

def main():
    dist.init_process_group(backend="nccl")
    rank, world = dist.get_rank(), dist.get_world_size()

    g = torch.tensor([float(rank + 1)], device="cuda")   # GPU0 -> 1.0   GPU1 -> 2.0

    dist.all_reduce(g, op=dist.ReduceOp.SUM)         # 默认求和：两张卡都变成 3.0
    if rank == 0:
        print("SUM   ->", g.item())                   # 3.0

    g = g / world                                    # 手动除以 N，才是「平均」
    if rank == 0:
        print("AVG   ->", g.item())                   # 1.5

    dist.destroy_process_group()

PyTorch 的 DDP 内部就是这么做的：它把参数打包成若干个 bucket，对每个 bucket 做 AllReduce(SUM)，然后乘以 1/world_size。所以你在写业务代码时完全看不到这一步。

PyTorch DDP 的最小可用脚本（注意注释里标出的三个坑）

import os, torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DistributedSampler, DataLoader

def main():
    # torchrun 已经把这些塞进环境变量：RANK / LOCAL_RANK / WORLD_SIZE / MASTER_ADDR
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    # 建立通信组：每个进程在这里得知「我是几号」和「一共几个人」
    dist.init_process_group(backend="nccl")

    model = MyModel().cuda()
    model = DDP(model, device_ids=[local_rank])   # 包一层，反向时自动 AllReduce 求平均

    sampler = DistributedSampler(dataset, shuffle=True)
    loader = DataLoader(dataset, sampler=sampler, batch_size=B, drop_last=True)

    for epoch in range(EPOCHS):
        sampler.set_epoch(epoch)              # 坑1：不写这句，每个 epoch 的数据切分完全相同
        for batch in loader:
            loss = model(batch)
            loss.backward()                    # 反向结束 -> 自动 AllReduce(SUM) -> 除以 world_size
            optimizer.step()
            optimizer.zero_grad()

    dist.destroy_process_group()

3.4 为什么「取平均」在数学上是正确的

这一步很多人跳过，但它解释了为什么数据并行可以放心用。设单卡 batch 大小为 B，N 张卡。

单卡训练时，一个 batch 的梯度是：

g_single = (1/B) · Σ ∇l(x_i)

数据并行时，每张卡各自算自己那份的均值，再对 N 张卡取平均：

g_dp = (1/N) · Σ_k [ (1/B) · Σ_{i∈batch_k} ∇l(x_i) ]
     = (1/(N·B)) · Σ_all ∇l(x_i)

结果正好等于「用 N×B 大小的 batch 在单卡上训练」的梯度。所以数据并行在数学上等价于扩大了 batch size，这带来两个直接推论：

• **学习率要重标定。**batch 从 B 变成 16B，通常要相应放大学习率（linear scaling rule）并加 warmup，否则收敛会变差。
• **各卡 batch 必须等大。**上面推导里能约掉 B，正是因为每张卡的 batch 大小相同。如果第 3 张卡因为数据不够只吃了半批，那它对平均梯度的「投票权重」就少了一半------梯度是有偏的。

⚠️ 必须加 drop_last=True

最后一个不满的 batch 会让各卡样本数不等。PyTorch 的 DistributedSampler 默认用 padding 补齐，但如果你自己写 sampler，很容易漏掉。症状是：训练能跑，loss 曲线正常，但精度就是差一点点------非常难查。

3.5 梯度累积与 AllReduce 的交互

显存不够时常用梯度累积：把一个大批拆成若干 micro-batch，累积梯度，最后再更新一次。此时 AllReduce 的位置很关键。

# 正确写法：前 N-1 个 micro-step 不做 AllReduce，最后一个才做
# 否则每个 micro-step 都 AllReduce 一次，通信量翻 N 倍，而且梯度被多除了一次
for i, batch in enumerate(loader):
    is_last = (i + 1) % accum_steps == 0
    ctx = contextlib.nullcontext() if is_last else model.no_sync()
    with ctx:
        loss = model(batch) / accum_steps
        loss.backward()
    if is_last:
        optimizer.step()
        optimizer.zero_grad()

04 单卡装不下：显存到底被谁吃掉了

先看清显存的四块构成，再逐个对症下药。这一节的方法都是训练框架的活，但调度系统需要知道它们的存在，因为不同的方案对卡数、对通信的要求完全不同。

4.1 显存的四块构成

训练一个模型，显存被四部分吃掉（P = 参数量）：

构成	含义	FP16/BF16 训练下的开销
模型参数	权重本身	2P 字节
梯度	每个参数对应一个梯度	2P 字节
优化器状态	Adam 需要动量 + 方差，外加一份 FP32 主权重	12P 字节
激活值	前向过程中为反向传播保留的中间结果	随 batch × 序列长度 × 层数增长

加起来是 16P 字节（不含激活值）。代入具体数字感受一下：

7B  模型：16 × 7   = 112 GB   → 单张 A100 80GB 放不下
10B 模型：16 × 10  = 160 GB   → 需要至少 2 张卡
70B 模型：16 × 70  = 1120 GB  → 需要至少 14 张 80GB 卡
100B 模型：16 × 100 = 1600 GB → 需要至少 20 张 80GB 卡

注意优化器状态占了 12/16 = 75%。这是个很重要的直觉：显存最大的黑洞不是模型本身，而是「为了更新模型而额外存的那堆东西」。

同样一个 100B 模型，四张卡上分别要装多少静态显存。注意优化器状态（红色）始终是最大的一块。

4.2 ZeRO：把「冗余的副本」切掉

ZeRO（Zero Redundancy Optimizer）的核心洞察很朴素：**数据并行里每张卡都存了一份完全相同的模型状态，这是纯浪费。**既然每张卡都要做同样的更新，那让每张卡只负责更新其中 1/N，需要的时候再从别人那里取回来就行。

阶段	切分什么	每卡显存（N 张卡）	额外通信
ZeRO-1	只切优化器状态	2P + 2P + 12P/N	无（更新后广播参数）
ZeRO-2	再切梯度	2P + (2P+12P)/N	gradient 需要 Reduce-Scatter
ZeRO-3	再切模型参数	(2P+2P+12P)/N = 16P/N	通信量约为 DDP 的 1.5--2 倍

N=16 时，100B 模型的每卡静态显存从 1600GB 降到 100GB。代价是 ZeRO-3 每一步都要 AllGather 把参数临时拼回来，通信量显著上升。

💡 决策顺序：先省显存，再调并行

实践中不要把这两种手段混着乱用。清晰的顺序是：

① 单机优化优先 （混合精度、激活检查点、FlashAttention，零通信代价）→ ② 还不够就上 ZeRO/FSDP （只加通信不加拓扑约束，跨机很自由）→ ③ 还装不下才上 TP/PP（引入严格的拓扑与通信约束）。

因为 ZeRO 不改变模型的形状，对调度器来说是「透明的」；而 TP/PP 会引入强拓扑约束。能不用就不用。

4.3 激活值：另一个隐形杀手

前面算的都是静态内存。实际训练中，激活值经常是压死骆驼的最后一根稻草------它随 batch size 和序列长度线性增长，在长上下文场景下能吃掉几十上百 GB。

对应手段是激活检查点（Activation Checkpointing），也叫梯度检查点：前向时只保存少数几个「检查点」层的输出，其余中间结果全部丢掉；反向传播需要时，从最近的检查点前向重算一遍。

• 收益：激活显存可降低 60%--70%，甚至一个数量级。
• 代价：多一次前向计算，训练变慢 20%--30%。

这是典型的「用算力换显存」。在显存紧张时这笔买卖通常划算，因为省下的显存可以换来更大的 batch，反而提升了 GPU 利用率。

手段	省多少显存	付出什么	是否需要改拓扑
混合精度 (BF16/FP16)	约 40--50%	几乎无，需注意数值稳定性	否
激活检查点	60--70%（激活值部分）	训练慢 20--30%	否
FlashAttention	注意力的 O(n²) → O(n)	无（还更快）	否
ZeRO-1/2/3	分别 4× / 8× / N×	通信量增加	否（很自由）
张量并行 TP	权重 ÷ TP 度	每层 4 次 AllReduce	是（必须同 NVLink 域）
流水线并行 PP	参数 ÷ PP 度	流水线气泡	是（段间需点对点）
CPU Offload	可再降 2--4 倍	PCIe 传输开销，显著变慢	否

4.4 选择路径

• 模型能装进一张卡吗？（按 16P 字节估算静态显存） 能 → 用 DDP 就行，只加卡数，不加复杂度。这是最省心的路径。
• 装不下，但只差一点？ 差 1.5--4 倍 → ZeRO-1 / ZeRO-2 / FSDP（分片优化器状态和梯度）。不改模型形状，通信可跨机，调度器无感知。
• 差很多，单层都塞不进？ → 引入张量并行。此时必须保证 TP 组落在同一个 NVLink 域内，调度器要开始做拓扑感知。
• 节点数很多、想减少段间通信？ → 叠加流水线并行 。TP 组内、PP 跨机，形成 TP × PP × DP 的三维结构。
• 显存还是不够？ → CPU Offload / NVMe Offload。这是最后手段，性能损失大，只在别无选择时用。

05 调度器到底管什么：划清职责边界

这是本文最重要的一节。你问「需要怎么实现」，答案的第一句是：先明确调度器不做什么。

5.1 一句话定位

训练框架负责「怎么算」，调度系统负责「在哪儿算」。调度器不参与任何梯度计算，它只做四件事：

• 选卡与放置（职责一）：从集群里挑出满足需求的物理卡，并决定哪个进程落在哪张卡上（拓扑映射）。
• 原子启动（职责二）：gang scheduling：要多少给多少，一张都不能少，否则整个作业继续排队。
• 注入坐标（职责三）：把 MASTER_ADDR / WORLD_SIZE / RANK 等环境变量写进每个进程，然后放手。
• 看护生命周期（职责四）：失败检测、整体重启、断点续训、优雅退出，以及队列公平和抢占。

5.2 职责划分表

具体事情	谁负责	说明
决定用几张卡、TP/PP/DP 怎么配	用户声明，调度器校验	调度器要检查「TP=8 但机器只有 4 张卡」这类矛盾
决定 rank ↔ 物理 GPU 的映射	调度器	这是拓扑感知的核心产出
注入 MASTER_ADDR 等环境变量	调度器	训练脚本只读，不猜
挂载共享存储 / 指定 checkpoint 路径	调度器	保证重启后能读到上次的快照
建立 NCCL 通信组、发 AllReduce	训练框架	调度器完全不碰
计算梯度、求平均、更新参数	训练框架	调度器完全不碰
写 checkpoint 文件	训练框架	调度器只负责「路径可写、重启能读到」
决定何时重启、重启多少次	调度器 + torchrun 协同	调度器看到容器挂了，torchrun 看到 rank 挂了
队列配额、优先级、抢占	调度器	训练框架没有这个概念

⚠️ 最常见的越界：调度器去管 batch size

有些调度系统为了「自动调优」，会根据卡数去改训练脚本的 batch size 或学习率。这是错的------因为改 batch size 会改变训练的数学等价性（见 3.4 节），而框架并不知道调度器动了什么。

正确的做法是：调度器只提供事实（你有几张卡、卡在哪），由框架或用户决定怎么用。

5.3 必须注入的环境变量

这是调度器与训练框架之间最重要的那份「契约」。参照 torchrun 的约定（这是事实标准，直接跟它对齐最省事）：

变量	含义	怎么算
MASTER_ADDR	rank 0 进程所在机器的地址	调度器从 placement 里取第 0 个 worker 的 IP
MASTER_PORT	用于 rendezvous 的端口	调度器分配，避免同机多任务冲突
WORLD_SIZE	总进程数（= 总 GPU 数）	= replicas × GPUsPer
RANK	全局编号，0 到 WORLD_SIZE-1	由拓扑映射决定，不是随便排
LOCAL_RANK	本机内的编号	决定用本机第几张卡，等价于 CUDA_VISIBLE_DEVICES 的索引
LOCAL_WORLD_SIZE	本机进程数	= 每台机器的 GPU 数（= GPUsPer）
NCCL_SOCKET_IFNAME	NCCL 走哪块网卡	多网卡机器必须指定，否则可能挑错
NCCL_IB_HCA	用哪块 IB 网卡	多 HCA 场景下指定，避免交叉

5.4 一次典型调度的完整时序

一次调度从提交到开始训练的完整时序。注意第 7 步之后，调度器就退出了计算链路。

✅ 这个边界让系统可以简单很多

因为调度器不需要知道任何深度学习的东西，它就是一个「带形状约束的资源匹配器」。这也意味着你可以用任何语言写它（Go / Rust / Python 都行），并且它可以同时调度 PyTorch、TensorFlow、JAX 的任务------因为它们的环境变量契约是一样的。

06 调度器架构设计

给出一套可落地的组件划分、数据模型和核心算法。

6.1 组件划分

调度器分层架构。控制面无状态（可多副本），数据面每节点一个 Agent。

6.2 数据模型

核心是三个概念：**作业（Job）**描述「要什么」，**任务（Task）**是作业的一个进程实例，**放置方案（Placement）**描述「给了什么」。

Go 数据结构（简化）

// Job ------ 用户提交的资源申请，关键在于「形状」而非「数量」
type Job struct {
    ID        string
    Queue     string
    Priority  int

    Replicas  int    // 需要多少个容器（通常 = DP 度 × PP 度）
    GPUsPer   int    // 每个容器要几张卡

    // Topology 是「形状」，告诉调度器这些卡之间该怎么连
    TP        int    // 张量并行度：这 TP 张卡必须在同一个 NVLink 域
    PP        int    // 流水线并行度：相邻两段之间需要点对点链路
    DP        int    // 数据并行度：最松，可跨机架
}

// Placement ------ 调度器的产出：逻辑 rank 到物理卡的映射
type Placement struct {
    Bindings []Binding      // 长度为 WorldSize 的映射表
    Cost     float64      // 打分，越小越好
}

type Binding struct {
    GlobalRank int
    Node       string
    DeviceID   int      // 本机内的第几张卡 -> LOCAL_RANK
}

6.3 核心调度循环

调度的本质是一个「带约束的最优化问题」：在所有满足硬约束的放置方案里，找代价最小的那个。

调度器主循环（简化伪代码）

func (s *Scheduler) ScheduleOnce() {
    snapshot := s.cluster.Snapshot()       // 当前每台机器的空闲卡 + 拓扑位置

    for _, job := range s.queue.Pending() {  // 已按 优先级 + 公平份额 排好序

        need := job.Replicas * job.GPUsPer

        // ── 步骤 1：Gang 判定 ────────────────────────────
        // 卡不够就整个作业继续等，绝不部分启动。
        // 这是防止「一半 worker 占着卡干等」死锁的关键。
        if snapshot.FreeGPUs() < need {
            continue
        }

        // ── 步骤 2：枚举满足硬约束的候选放置 ──────────────
        // 硬约束 = TP 组必须落在同一个 NVLink 域（required 级）
        candidates := CandidatePlacements(job, snapshot)
        if len(candidates) == 0 {
            continue              // 有卡但不满足形状，一样不能起
        }

        // ── 步骤 3：打分，取代价最小的 ───────────────────
        best := ArgMin(candidates, func(p Placement) float64 {
            return p.Cost       // 打分函数见下一节
        })

        // ── 步骤 4：一次性绑定全部容器 ───────────────────
        // 注意是「全部」------绑定和启动都必须是原子的
        s.BindAtomically(job, best, snapshot)
    }
}

💡 为什么是「枚举 + 打分」，而不是贪心逐个分配

逐个分配（先给 rank 0 挑一台空闲机器，再给 rank 1 挑......）看似简单，但会陷入局部最优：前面几步占掉了本该留给 TP 组同域的位置，后面就拼不成正确的形状了。

正确做法是以「组」为单位做整体决策。工程上不需要真的穷举所有组合（那是指数级），只用按 NVLink 域做粗粒度分桶，再在桶内做细粒度排列即可，复杂度是可接受的。

6.4 别忘了：作业完成后要能干净地退出

这部分容易被忽略，但直接影响集群可用性。几个必须处理的点：

• **卡释放要幂等。**进程崩溃、被 OOM killer 杀掉、节点掉线，三条路径都要能正确回收卡。回收漏了就会出现「卡看着空闲但谁也申请不到」的幽灵占用。
• **僵尸进程检测。**训练进程有时会 fork 出子进程，主进程挂了子进程还在占显存。Agent 要按进程树清理，不能只杀主进程。
• **TTL 与超时。**设置 ttlSecondsAfterFinished 之类的机制，避免已完成的作业一直挂在队列里。
• **重启计数上限。**某些故障（比如代码 bug）会无限重启，必须有 max_restarts 并且区分「可重试的故障」和「不可重试的故障」。

07 拓扑感知：把对的卡配成对

这一节是性能的胜负手。同样的 64 张卡，排布方式不同，训练速度能差好几倍。

7.1 现代 GPU 集群是分层的

带宽不是均匀的，而是像洋葱一样分层：

层级	互联方式	典型带宽	适合承载哪种并行
卡内	HBM	数 TB/s	---
同机 8 卡	NVLink 4 / NVSwitch (H100)	900 GB/s / GPU	张量并行 TP
机柜内 72 卡	NVLink Switch (GB200 NVL72)	1.8 TB/s / GPU	TP 扩展到 16--72 度
机架间	InfiniBand NDR	约 50 GB/s / 端口	流水线并行 PP、数据并行 DP
跨机房	以太网	25--50 GB/s	只适合 DP（且性能会打折）

相邻两层之间大约有 18 倍 的带宽差。所以「把 TP 组放到同一个 NVLink 域」这件事不是优化，是正确性要求。

7.2 rank 排布规则

Megatron 系的惯例是 rank 按「TP 最内层、PP 中间、DP 最外层」排列：

global_rank = dp_idx × (PP × TP) + pp_idx × TP + tp_idx

这样排列的好处是：连续的一段 rank 天然构成一个 TP 组，再外面一段构成 PP 组。调度器只要保证「连续的 TP 张卡在同一台机器上」即可。

举例：TP=8, PP=4, DP=8，共 256 张卡（32 台 8 卡机）：

rank  0- 7   → 机器 A  的 8 张卡   （TP 组 0，PP 段 0，DP 副本 0）
rank  8-15   → 机器 B  的 8 张卡   （TP 组 1，PP 段 1，DP 副本 0）
rank 16-23   → 机器 C  的 8 张卡   （PP 段 2）
rank 24-31   → 机器 D  的 8 张卡   （PP 段 3）
rank 32-39   → 机器 E  的 8 张卡   （DP 副本 1，PP 段 0）
...

调度器的约束翻译成一句话：每 8 个连续 rank 必须落在同一台机器的 8 张卡上，且这 32 台机器之间的 IB 网络尽量少跨交换层。

7.3 拓扑感知的打分函数

把「通信代价」显式建模成一个可比较的分数。最基本的模型是「通信量 ÷ 可用带宽」：

cost(placement) = Σ_over_all_links  (该链路上要搬的字节数 / 该链路的带宽)

Go 实现（简化）

// 代价越小越好。核心思想：通信越频繁的组，越要放在带宽越高的域里。
func CommCost(p Placement, job Job) float64 {
    var cost float64

    // TP 组：每层 4 次 AllReduce，是最贵的一环，权重最高
    // 如果被打散到不同机器，带宽从 900GB/s 掉到 50GB/s，代价放大 18 倍
    for _, group := range p.TensorGroups(job.TP) {
        cost += float64(job.BytesPerTPAllReduce) / BandwidthWithin(group)
    }

    // PP 组：段间传激活值，每 micro-batch 一次
    for _, pair := range p.PipelinePairs(job.PP) {
        cost += float64(job.BytesPerStageHandoff) / BandwidthBetween(pair)
    }

    // DP 组：每 step 一次梯度 AllReduce，而且能和反向计算重叠，
    // 大部分时间被隐藏掉了，所以权重给得很低
    cost += 0.2 * float64(job.GradBytes) / BandwidthBetween(p.DataGroups())

    return cost
}

💡 权重的物理依据

那个 0.2 不是随手写的。它的依据是：DP 的梯度同步可以和反向传播重叠（第 8 节会讲），实际暴露在关键路径上的时间远小于理论通信量。而 TP 的 AllReduce 在当前实现里很难完全隐藏，几乎全部落在关键路径上。

工程上不要凭感觉调这个系数，要用 nccl-tests 实测带宽，再用真实训练的 step time 校准。

7.4 硬约束与软约束

约束类型	含义	不满足时	典型用法
required（硬）	绝对不能违反	该放置方案直接作废	TP 组必须同 NVLink 域
preferred（软）	尽量满足，可降级	计入 cost，但仍可调度	DP 副本尽量少跨机架

一个常见的策略分层：从最细的层级开始尝试放置，失败就上一级 。比如先试「同一台机器」→ 不行试「同一机架」→ 再不行试「同一交换机域」→ 都不行就报错（TP 场景）或降级接受（DP 场景）。这正是 NVIDIA Run:ai 的做法：If scheduling cannot occur at a lower level, it automatically moves up the hierarchy.

7.5 拓扑信息从哪来

调度器需要知道每台机器在物理拓扑中的位置。有三条获取路径：

1. 手工打标签（最简单，起步阶段够用） ：给节点打上 topology/rack=rack-03、nvlink-domain=d0 之类的标签。缺点是扩容改造时要人工维护。
2. 自研 Agent 探测 ：Agent 启动时调用 NVML 探测 nvmlDeviceGetTopologyCommonAncestor 拿到 GPU 之间的亲和关系（哪些卡走 NVLink、哪些走 PCIe），再从中推导出「NVLink 域」的划分。同时探测 IB 网卡的 HCA 列表与 NUMA 亲和。
3. 网络设备联动：接 NetQ / 交换机 LLDP 数据，自动推导机架与交换层级。规模大时值得做。

⚠️ 一个血泪教训：不要相信 hostname 的命名规律

很多集群的机器命名（如 node-01、node-02）看起来和物理位置相关，但扩容、搬迁、故障替换之后就不再成立了。拓扑信息必须显式声明或程序化探测，绝不能从主机名推断。

08 通信优化进阶

调度器把卡放对了，接下来是框架层面把通信藏起来。这一节的内容可以不做，但决定了你的集群是「能跑」还是「跑得快」。

8.1 让通信藏在计算后面

反向传播是从最后一层往前算的，所以越靠后的层，梯度越早算完。既然这样，就不用等所有梯度都算完再统一 AllReduce------算完一层就发一层。

PyTorch DDP 的实现是三个机制配合：

1. 梯度分桶 ：不为每个参数单独发一次 AllReduce（小消息的启动开销比传输本身还大），而是把参数打包成 bucket，默认上限 25 MiB（bucket_cap_mb 可调）。
2. 反向钩子：给每个参数注册 backward hook，bucket 一满就立刻发起 AllReduce，不等其他 bucket。
3. 独立 CUDA 流：AllReduce 跑在单独的 stream 上，和反向计算并行执行。

效果是：当反向传播算到第一层时，后面那些层的 AllReduce 早就做完了。通信几乎完全被隐藏。所以在 NVLink 上，2 卡 DDP 的实际加速比能达到 1.9 倍左右；即使在 PCIe 上也能到接近 1.9。这也是为什么 DP 的通信代价被打分函数给了低权重。

8.2 分层 AllReduce

面对「机内快、机间慢」的分层结构，最优做法是按层级分两步做，而不是让所有卡围成一个大环：

第一步：机内 Reduce-Scatter（用 NVLink，900 GB/s）
   → 每台机器内部先把 8 张卡的梯度聚合，得到 1/8 大小的部分结果

第二步：机间 AllReduce（用 IB，50 GB/s）
   → 只有这部分结果过慢速网络，通信量只有原来的 1/8

第三步：机内 AllGather（用 NVLink）
   → 把完整结果分发回本机每张卡

这个策略的意义在于：慢速的跨机链路只承担了 1/N 的流量 。如果让 64 张卡围成一个大环，每张卡的梯度都要绕一圈穿过 8 次跨机链路，性能会差得多。NCCL 会根据拓扑自动选择这个策略，但前提是它探测到的拓扑是对的------这就回到第 7 节。

8.3 其他手段

手段	原理	收益	风险
梯度压缩	fp16/bf16 传输或稀疏化后再通信	通信量减半或更多	可能影响收敛
SHARP / 网内计算	在 IB 交换机内部完成规约，不给端点 GPU 增加负担	AllReduce 延迟降低 2--3 倍	需要支持的交换机
GPUDirect RDMA	网卡直接读写显存，绕开 CPU 和系统内存	减少一次拷贝	需要正确的亲和配置
拓扑感知 rank 重排	让通信最密集的 rank 落在物理最近的卡上	直接消除跨层流量	需要调度器配合

8.4 常用 NCCL 调优变量

调度器在注入环境变量时，顺手带上这些是很划算的：

NCCL_SOCKET_IFNAME=eth0        # 多网卡机器必须指定，否则 NCCL 可能挑到管理网口
NCCL_IB_HCA=mlx5_0,mlx5_1      # 指定使用的 IB 网卡，避免走错
NCCL_IB_DISABLE=0              # 启用 IB
NCCL_NET_GDR_LEVEL=5           # 允许 GPUDirect RDMA
NCCL_DEBUG=INFO                # 排查时打开，会让它打印选中的算法与拓扑
NCCL_DEBUG_SUBSYS=INIT,GRAPH   # 只关心初始化与拓扑探测
NCCL_IB_TIMEOUT=22             # 大集群下适当放大，避免瞬时拥塞导致超时

⚠️ NCCL_DEBUG=INFO 不要长期开着

它会为每个通信组打印大量日志。在 512 卡的作业里，这能产生几十万行日志，既拖慢启动又冲爆日志系统。只在排查问题时按需打开。

09 容错与弹性训练

当训练要跑三个月、占用上千张卡时，「出错」不是异常，是常态。必须假设任何时刻都有卡会挂。

9.1 为什么要做，成本有多高

以 BLOOM-176B 为例：训练用了约 3.5 个月、接近 100 万 GPU 小时。按这个量级，即使硬件单卡年故障率只有 1%，上千张卡跑三个月也必然会遇到故障。如果每次故障都要从头开始，训练永远跑不完。

所以容错的核心问题是：失败之后，怎么只损失一点点进度？

9.2 torchrun 的弹性机制

PyTorch 的 torchrun 内置了容错与弹性，靠三个参数控制：

torchrun \
  --nnodes=4:8 \              # 最少 4 台，最多 8 台（弹性范围）
  --nproc-per-node=8 \        # 每台 8 个进程
  --max-restarts=3 \          # 最多重启 3 次
  --rdzv-id=job-12345 \       # 作业唯一 ID，所有节点必须一致
  --rdzv-backend=c10d \       # rendezvous 后端
  --rdzv-endpoint=node0:29500 # rendezvous 服务地址
  train.py

工作方式是：所有参与节点先在 rendezvous 服务处「会合」，由它统一分配 rank。当一个 worker 挂掉，torchrun 会终止所有进程，然后重新会合、重新分配 rank、重启全部进程。训练脚本需要做的配合只有一件事------重启后从上一个 checkpoint 恢复。

训练脚本的正确结构

def main():
    init_process_group()            # 每次重启都会重新执行
    model, optimizer, start_epoch = load_snapshot(CKPT_PATH)   # 关键：先恢复
    train(model, optimizer, start_epoch)

def train(model, optimizer, start_epoch):
    for epoch in range(start_epoch, MAX_EPOCHS):
        for step, batch in enumerate(loader):
            train_step(model, optimizer, batch)
            if step % CKPT_INTERVAL == 0:
                save_snapshot(model, optimizer, epoch, CKPT_PATH)

9.3 Checkpoint 的三个工程细节

① 谁来写

由 rank 0 写，其他 rank 在 barrier 处等待：

if dist.get_rank() == 0:
    torch.save(state_dict, path)
dist.barrier()      # 坑：忘了这个 barrier，其他 rank 会在 rank0 还在写时就开始下一轮，
                    #     随后读到一个不完整的文件而 hang 死

但如果是 TP/PP 切分过的模型，每张卡只持有一部分参数，就必须所有 rank 都参与保存（各写各的分片），恢复时再各自加载自己的分片。

② 写多频繁

这是个权衡：写得勤，挂了损失少，但写盘本身会阻塞训练；写得稀，浪费算力。工程上的解法是异步 checkpoint：先把显存里的状态拷贝到 CPU 内存（或直接落到本机 NVMe），立刻返回继续训练，后台线程再慢慢传到共享存储。这样阻塞时间从「写完整个文件」降到「一次显存到内存的拷贝」。

③ 存哪里

存到共享存储（NAS / 对象存储），不要只存本地盘------因为重启后任务很可能被调度到别的机器上，本地盘上的文件读不到。如果追求性能，可以采用「两级策略」：先写本地 NVMe（快，用于快速恢复），再异步同步到共享存储（持久，用于节点故障后恢复）。

9.4 弹性扩缩容的三个坑

弹性训练听起来很美------集群闲了就多给几张卡。但 world_size 变化会打破训练的数学前提：

1. 全局 batch size 变了。如果每卡 batch 固定，卡数从 32 涨到 64，全局 batch 就翻倍了，学习率需要重新标定（见 3.4 节）。
2. 数据切分要重做 。DistributedSampler 需要重新 set_epoch 来重新分片，否则数据会被重复或跳过。
3. 优化器的状态要和 world_size 对齐。用 ZeRO 时优化器状态是按 N 分片的，N 变了分片方式也得重算。

所以实践建议：弹性缩放只在 checkpoint 边界上做，并且显式地重新标定学习率和重做数据切分。不要试图在训练中途无缝改变 world_size。

9.5 调度器侧的容错职责

场景	调度器该做什么
单个 worker 进程崩溃	不单独重启它（会造成 rank 错位），而是整体重启整个作业
整个节点掉线	把该节点的所有容器标记为失败，整体重启；优先避开这台机器重新放置
GPU 出现 ECC 错误 / 硬件故障	把该 GPU 从资源池摘除（加 taint），只重启作业
反复重启仍失败	达到 max_restarts 后置为 Failed 并通知用户，不要无限重试烧卡
训练正常收敛完成	释放全部资源，清理进程树，保留日志与最终 checkpoint

10 落地路线与验收标准

不要一次性把所有东西做完。下面五个阶段，每一阶段都有独立的验收标准，做完一阶段就能稳定用。

阶段一：单机多卡跑通 DDP

做什么 ：在一台 8 卡机上用 torchrun --standalone --nproc-per-node=8 跑通 DDP 训练，验证梯度平均的数学正确性。

验收标准：

• 单卡 batch 8 × 8 卡跑出的 loss 曲线，与单卡 batch 64 的 loss 曲线基本重合（证明确实等价于大 batch）。
• 用第 3.3 节的验证脚本确认 AllReduce 行为符合预期。
• 初步测出 8 卡相对 1 卡的加速比，记下这个基线数字。

阶段二：多机跑通，手工注入环境变量

做什么 ：两台机器，手工在两台机器上分别执行 torchrun，手动指定 --node-rank、--rdzv-endpoint。

验收标准：

• 16 卡训练能稳定跑起来，不出现 NCCL 初始化超时。
• 知道本集群正确的 NCCL_SOCKET_IFNAME 取值（用 NCCL_DEBUG=INFO 确认它选中的网卡是不是高速网）。
• 记录 16 卡的加速比。如果明显低于预期，说明跨机带宽没被用上，先解决这个再往下走。

阶段三：接入调度器，实现 gang scheduling

做什么：写一个最小调度器，只做三件事------接收作业、判定资源是否够、一次性创建全部容器并注入环境变量。

验收标准：

• 提交一个需要 8 张卡的作业到只有 7 张空闲卡的集群，作业保持排队状态，不启动任何容器。
• 空闲卡凑齐到 8 张时，作业自动开始运行。
• 模拟杀死其中一个 worker，作业能被检测到并整体重启。

阶段四：拓扑感知

做什么：给节点打拓扑标签，在调度器里加入打分函数，优先把 TP 组放在同一台机器。

验收标准：

• 在调度器日志里能看到每个作业的 placement 决策结果（哪个 rank 去了哪张卡，cost 是多少）。
• 做一个对照实验：人为把 TP 组拆到两台机器上，实测 step time 的劣化幅度，验证你的打分函数方向是对的。
• 用 nccl-tests 的 all_reduce_perf 实测同机与跨机的带宽差，作为打分函数系数的依据。

阶段五：容错、弹性与可观测性

做什么：加上 checkpoint 协同、重启策略、以及监控看板。

验收标准：

• 在训练中途强制杀掉一个节点，作业能在 checkpoint 间隔内自动恢复，且 loss 曲线连续（没有从头发散）。
• 调度器看板上能看到：每个作业的 GPU 数量、运行时长、重启次数、失败原因。
• 能看到每个作业的 GPU 利用率与 NCCL 带宽（通过 DCGM 导出到 Prometheus）。

✅ 一个务实的提醒

很多团队一上来就想做「自动并行策略搜索」（让调度器根据模型大小自动决定 TP/PP/DP 怎么配）。这件事技术上可行，但收益远不如把前面五个阶段做扎实。

原因是：自动搜索需要精确的性能模型（预测每种配置的 step time），而这个模型的准确性严重依赖于你要测量的正是第 7 节的拓扑数据。先把拓扑数据采准，自动搜索才有基础。

11 踩坑清单

按症状查表。这些都是多卡多机训练里反复出现的真实问题。

症状	最可能的原因	怎么查 / 怎么修
卡在 init_process_group 不动，也不报错	MASTER_ADDR / MASTER_PORT 写错；防火墙拦了端口；有 rank 根本没起来	开 NCCL_DEBUG=INFO 看它连到哪个地址；确认所有 rank 的 MASTER_ADDR 是同一个且可达；检查 WORLD_SIZE 是否等于实际启动的进程数
报 NCCL error: unhandled system error	NCCL 挑到了错误的网卡	显式设置 NCCL_SOCKET_IFNAME 和 NCCL_IB_HCA
各卡 loss 不一致 / 训练发散	DistributedSampler.set_epoch() 漏了，导致每个 epoch 数据切分完全相同	在每个 epoch 开头调用 sampler.set_epoch(epoch)
梯度累积后 loss 变成 NaN	没有用 no_sync()，梯度被累积后又 AllReduce 了一次	前 N-1 个 micro-step 用 model.no_sync() 包住 backward
保存 checkpoint 后全体 hang 死	rank 0 保存完了，其他 rank 没等它，读到半个文件	rank 0 写完后加 dist.barrier()
作业起来一半，另一半一直 Pending	没有 gang scheduling，默认调度器逐个放置造成死锁	实现 all-or-nothing 的原子放置（PodGroup / minAvailable 或自研）
训练能跑但速度远低于预期	TP 组被拆到了不同机器上，走 IB 而不是 NVLink	查调度器的 placement 日志，确认连续 rank 是否落在同一台机；实测 all_reduce_perf 的同机 / 跨机带宽差
单卡 batch 不等导致精度下降	最后一个 batch 不满，各卡样本数不一样，梯度有偏	DataLoader 加 drop_last=True
重启后 loss 从头发散	只加载了模型权重，没加载优化器状态 / epoch 计数 / 数据加载器进度	checkpoint 里要存全：model + optimizer + scheduler + epoch + step + RNG 状态
卡显示空闲但申请不到	进程崩溃后 GPU 被僵尸进程占着，资源回收不干净	Agent 按进程树清理，不只杀主进程；加定期扫描与对账机制
小张量的 AllReduce 很慢	Ring 算法步数是 2(N-1)，小消息时延迟占主导	交给 NCCL 自动选（大消息走 Ring，小消息走 Tree）；不要手动强制算法

⚠️ 最值得先做的一件事

如果你只能做一件事，做这个：把 NCCL_DEBUG=INFO 打开，完整读一遍 NCCL 初始化日志。

它会告诉你：探测到了几条 NVLink、走了哪块网卡、机内和机间分别选了哪个算法。90% 的「能跑但慢」问题，答案都在那几十行日志里。很多人跳过这一步直接去调参，方向就错了。

---

参考资料

• PyTorch 官方文档 · torchrun (Elastic Launch) ------ 环境变量契约与弹性机制的权威定义
• PyTorch 官方教程 · Fault-tolerant Distributed Training with torchrun
• Hugging Face Blog · The Technology Behind BLOOM Training ------ 3D 并行的工程实践案例
• OpenAI · Techniques for training large neural networks ------ 显存优化的全景梳理
• Volcano 官方文档 · Introduction ------ gang scheduling、拓扑感知、队列管理的工业实现
• Kubeflow 文档 · KAI Scheduler ------ 训练作业的 gang scheduling 接入方式
• NVIDIA Run:ai 文档 · Network Topology-Aware Scheduling ------ 分层拓扑与降级策略
• NVIDIA OSMO 文档 · Topology-Aware Scheduling ------ 硬约束 / 软约束的建模方式
• NVIDIA · NCCL 官方仓库与 nccl-tests ------ 实测带宽的必备工具
• DeepSpeed · ZeRO 论文（Rajbhandari et al., 2020）------ 显存分片的理论基础
• Megatron-LM / Megatron-DeepSpeed · 张量与流水线并行的参考实现

本文为技术调研整理，用于理解与设计参考。文中数值为理论估算，实际性能请以你的硬件与网络环境实测为准。