07-16 下午：分布式大模型训练与推理

下午这节课讨论的是模型和 batch¹ 都很大时，一张 GPU 已经放不下或算不动，怎样把一次训练和一次推理拆到多张卡上。后面的 DP、ZeRO²、PP、TP、长上下文并行名字很多，拆法也各不同，但绕不开的三个问题始终是每张卡保存哪一部分状态、负责哪一段计算、数据何时要交给别的卡。

自回归语言模型推理

把一句 prompt³ 交给语言模型后，模型一次只预测下一个 token⁴，条件是已有全部历史；把刚预测出的 token 接回输入，再预测下一个，直到结束符；这种逐 token 生成方式叫自回归（autoregressive）⁵。一个 token 在模型中先被映射为长度为 \(H\) 的向量，形成形状约为 \([B,S,H]\) 的张量⁶：

• \(B\) 是 batch size，一次并行处理多少个样本；
• \(S\) 是 sequence length，即 prompt 的 token 数；
• \(H\) 是 hidden size，单个 token 的隐藏向量维度。

这三个维度的读法要固定下来。\([2,16,4096]\) 表示两个请求，每个请求当前有 16 个 token，每个 token 是 4096 个浮点数。若只看其中第 1 个请求，可以去掉 batch 维得到 \([16,4096]\)；再看其中第 3 个 token，又得到长度 4096 的向量。矩阵乘法中的许多形状错误，来自没有写明当前张量的行和列各代表什么。

一次前向传播可以还原成下面几类操作。embedding 查表把 token 编号变成 [B,S,H]；attention 的线性投影把最后一维从 H 变成 Q/K/V 所需维度；\(QK^T\) 在每个 head 中产生 [B,heads,S,S] 的注意力分数；softmax 后乘 V 仍回到 [B,heads,S,head_dim]；FFN 对每个 token 独立做 [B,S,H] × [H,4H] 和 [B,S,4H] × [4H,H] 的矩阵乘；最后输出投影得到 [B,S,V] 的 logits。

embedding 只给 token 一个可计算的向量表示，并不含词序。位置编码或位置嵌入再把它位于第几个位置补进去。随后张量经过一层又一层 Transformer block；每层由 attention 和 FFN/MLP 等部分组成，最后经输出投影和 softmax 得到词表上的概率分布。

图：因果（masked）注意力让每个位置只看它之前的 token；这是自回归语言模型能够逐 token 生成的关键结构。

训练时的过程比推理多了反向传播。前向传播得到预测分布，与目标 token 算交叉熵损失；反向传播从损失倒着求每个参数的梯度；优化器根据梯度更新参数。以后看到并行训练的代码，先在脑中把它还原成这条链：forward -> loss -> backward -> optimizer.step()。并行策略没有改变这件事，只是把每一步拆到多张卡上。

训练语言模型时，loss 并不是只在最后一个 token 上计算。输入序列 \([x_0,x_1,\ldots,x_{S-1}]\) 的目标是右移一位的 \([x_1,x_2,\ldots,x_S]\)。第 1 个位置的输入是 \(x_0\)，模型要预测 \(x_1\)；第 2 个位置的输入包含 \(x_0,x_1\)，要预测 \(x_2\)。因果 mask 保证每个位置只能用左侧信息。这样一次前向能同时提供 S 个训练样本，而不必把每个上下文长度单独跑一遍。

Prefill、Decode 与推理调度

prefill⁷ 对已有 prompt 的所有 token 做前向并建立 KV cache⁸，query-key 交互与矩阵乘法规模较大，通常更接近计算密集；decode⁹ 每步只生成一个新 token，却要读取所有历史 KV，常受显存带宽和单请求延迟限制。continuous batching 把不同请求的 decode step 合在一个 batch，并在请求结束时补入新请求。它必须维护每个序列的长度、采样状态和 KV block 映射，不能像训练那样把所有样本简单 padding 到固定长度。

用一个短例子区分 prefill 和 decode。用户输入“今天 天气 很”，prefill 一次处理这三个 token，并保存每一层对它们的 K、V。第一步 decode 只把新 token“好”送入模型，attention 查询它和前四个 token 的 K/V，输出“。”；第二步只把“。”送入，前面四个 token 的 K/V 不必重算。若没有 KV cache，每生成一个 token 都要把完整历史重新跑一遍，序列越长浪费越大。

continuous batching 解决的是另一个问题。请求 A 已经生成到第 100 个 token，请求 B 刚到达，请求 C 在第 8 个 token 结束。传统静态 batch 必须等最长请求结束；continuous batching 可以在每个 decode step 结束后移除 C、补入 B，让设备里的 token 数尽量稳定。它换来的是调度复杂度、变长张量处理和 KV 空间管理。

图：prefill 大而规则、偏计算密集；decode 每步只出一个 token 却反复读 KV，偏带宽/延迟受限。两者决定了推理调度的不同取舍。

KV cache 是按层保存的 K/V 张量，其字节数近似为 layers × tokens × 2 × kv_heads × head_dim × bytes_per_element。若采用 grouped-query attention，kv_heads 小于 query heads，缓存会显著下降；若上下文达到显存上限，调度器必须限制并发、换出、分页管理或拒绝请求。吞吐、首 token 延迟和每 token 延迟互相牵制，单报 token/s 无法描述服务行为。

系统通常同时关注首 token 时间（TTFT）、后续 token 间延迟（TPOT）和总体吞吐。为了吞吐积累很大的 batch，可能伤害交互请求的首 token 延迟；连续批处理、分页式 KV cache、分开调度 prefill 与 decode，都是在处理这个矛盾。

训练和推理也会用 TP、PP、序列并行，但目标不同：训练更在意吞吐、全局 batch 和模型状态；推理更在意 KV cache 能否容纳、请求到达的动态性与用户实际等待时间。分析一个方案之前，先分清正在解决哪一种负载，后面的优化才有明确含义。

分布式训练的约束

模型参数不断增长，而单卡算力和显存并不能无限按同样速度增长。即便参数本身勉强放得下，训练还要保存梯度、优化器状态和中间激活¹⁰；长序列的 attention 又会放大激活与 KV cache 的开销。于是模型是否放得下和每一步是否来得及算完是两个不同问题。

训练显存至少有四类对象。参数是模型文件里的权重；梯度是损失对各参数的偏导，形状通常和参数相同；优化器状态是 Adam 等算法额外保存的一阶动量、二阶动量和主权重；激活是前向传播中为了反向求导而保存的中间张量。推理通常只需要参数和 KV cache，训练则四类都会出现。看到“7B 模型 14GB”时，那是 BF16 权重的下限，并不等于训练所需显存。

分布式训练并不是简单地把一块 GPU 换成多块。若多张卡各自拿到不同数据、各自更新同一个模型副本，却从不通信，它们很快就会变成不同的模型。多卡训练成立的前提是：在该同步的地方同步，在可以独立计算的地方不通信，并尽量把必须发生的通信藏在计算时间里。

通信：点对点与集合操作

点对点通信（point-to-point, P2P）是一个 rank 向另一个 rank send，对方 recv。集合通信（collective communication）由一个通信组里的所有 rank 共同参加，例如 4 张 GPU 同时执行一次 AllReduce¹²。它最终仍由许多点对点传输组成，但由 NCCL、MPI、Gloo 等通信库选择拓扑和实现。

集合操作有一个容易忽略的约束：参与者必须按相同顺序进入同一个 collective。一个 rank 没有进入、卡住或退出，其他 rank 会一直等待，最后通常表现为 timeout。因此调试分布式程序时，不能只看报错的那张卡；要找第一个没有到达同步点的 rank。

常见集合通信

设每个 rank 上有一个局部张量，collective 的名称描述数据在组内怎样变化。

用四个 rank 的 AllReduce 看语义。设局部向量分别是

rank0: [1, 0]    rank1: [0, 2]
rank2: [3, 0]    rank3: [0, 4]

若做 sum AllReduce，四个 rank 结束后都得到 [4,6]。实现可以是 ring reduce-scatter 再 all-gather，也可以是树形规约；使用者关心的是语义和通信量，底层库关心的是拓扑和算法。AllGather 则不同，若 rank0 有 [1,0]、rank1 有 [3,4]，结束后两个 rank 都得到 [1,0,3,4]，中间没有把对应元素相加。

图：Scatter 改变数据所有权，Gather 恢复集中表示；每个箭头都对应实际通信量。

图：三种集合通信都围绕每个进程持有一块的数据展开，区别在于是拼回去还是规约后分发。

不要把 AllGather 和 AllReduce 当作只差一个字的函数。前者的结果是拼接后的原始数据；后者的结果是逐元素相加、平均等规约后的数据。训练中究竟该用哪一个，取决于你需要的是完整参数、完整激活，还是一致的全局梯度。

通信延迟与带宽

课上用一个常见的近似模型表示一次消息传输：

图：通信时间由固定延迟与随消息大小增长的传输时间组成；据此决定是合并小消息还是用 ring 处理大消息。

\(\alpha\) 是启动延迟，例如建立/调度一次传输的固定开销；\(n\) 是消息大小；\(\mathrm{BW}\) 是有效带宽，\(\beta=1/\mathrm{BW}\)。消息很小时，主要输在 \(\alpha\)，此时合并小消息、减少调用次数比盯着峰值带宽更重要。消息很大时，\(n\beta\) 主导，链路带宽、拓扑和传输算法才成为重点。

例如 AllGather 可以用 ring 实现。把 \(p\) 个 rank 排成环，每个 rank 先持有自己的第 \(i\) 块；每一轮把当前块传给下家，同时从上家接收一块。经过 \(p-1\) 轮，每个 rank 都收到了其余所有块。每轮只搬一块较小的数据，链路可同时工作，因而对大消息常有很好的带宽利用率。Ring AllReduce 通常可理解为一段 ReduceScatter 加一段 AllGather。

图：ring 让所有链路的带宽同时被利用，适合大消息；多轮交换是它的代价。

这也解释了为什么通信库会按消息大小选算法：小张量未必适合 ring 的多轮启动，大张量又不适合让一个根节点成为瓶颈。实际耗时还会受 PCIe、NVLink/NVSwitch、跨节点网络和物理拓扑影响，不能只用 GPU 数量推断。

图：通信库按消息大小在启动开销优先与带宽优先之间切换实现；这正对应 \(\alpha\) 与 \(n\beta\) 哪个主导的判断。

数据并行

数据并行（data parallelism, DP）沿 batch 维切分数据。假设有四张 GPU，GPU 0 处理 mini-batch 的第 0 份，GPU 1 处理第 1 份；每张卡上却都保留完整的模型参数 \(\theta\)。它们分别执行前向与反向，得到局部梯度 \(g_0,g_1,g_2,g_3\)。

图：每卡都有一份完整模型，输入按 batch 切分；计算是并行的，参数靠通信保持一致。

要让四份模型继续保持一致，更新前需要计算全局平均梯度：

AllReduce 可以同时完成求和和结果分发，所以每张 GPU 得到同一个 \(g\)，再各自执行同样的 optimizer step，最终参数仍相同。这里的通信发生在反向传播产生梯度之后；没有它，四张卡等价于各自训练四个互不一致的模型。

DP 的加速上限受通信约束。若单卡一个 step 的计算时间是 \(C\)，AllReduce 梯度时间是 \(M\)，理想多卡时间仍至少有 \(C+M\) 的一部分。GPU 数增加后，单卡 batch 变小、固定开销占比上升，通信量也会随梯度大小变化。因此吞吐曲线常常先近似线性增长，再逐渐弯曲，最后增加卡数不再划算。

较早的做法是 parameter server：多个 worker 把梯度交给服务器，服务器汇总、更新参数后再发回。它的逻辑直接，但服务器既是吞吐瓶颈又是单点故障。现代同步 DP 更常让 worker 直接集合通信，不再保留一个集中式参数服务器。

图：集中式参数服务器逻辑直观，但所有梯度都汇集到一点，既是吞吐瓶颈也是单点故障。

图：现代 DP 让 worker 直接通过 AllReduce 交换梯度，去掉了集中式参数服务器这个单点。

Global Batch 与梯度累积

若每卡 batch 为 \(b\)，数据并行大小为 \(p\)，没有梯度累积时全局 batch 是 \(B_{\text{global}}=pb\)。显存不够而又希望维持大 batch 时，可连续跑 \(k\) 个 micro-batch，只累加梯度、不更新参数，最后再同步/更新；此时全局 batch 为 \(pkb\)。

这不是单纯的工程细节。batch 变大，单步梯度的统计特性、每个 epoch 的更新次数和学习率设置都会变。吞吐提升不自动意味着同样的收敛效果，训练配置通常要和并行规模一起重新检查。

AllReduce 的归约语义

数据并行的每个 rank 持有同一组参数、处理不同 mini-batch。若本地梯度为 \(g_r\)，同步 SGD 需要得到

然后所有 rank 用同一个 \(g\) 更新同一份参数。AllReduce 完成的是求和，除以世界大小可能由框架或优化器完成；把平均做两次会让学习率实际缩小 \(P\) 倍。梯度累积时，多个 micro-batch 的局部梯度要在正确边界再同步，no_sync 一类接口的含义正是推迟 collective，而不是跳过它。

通信-计算重叠依赖梯度产生的顺序。反向传播从最后一层开始；当某个 bucket 的所有梯度准备好，就可在通信 stream 上启动 AllReduce，同时计算前面层的梯度。bucket 太小，collective 启动次数多；太大，首个通信启动太晚。真正的重叠需要 profiler 时间线证明，并要求网络、GPU DMA 和计算资源没有已经互相饱和。

训练显存

模型状态与混合精度

以 \(\Psi\) 表示参数量。若模型权重和梯度用 BF16，各自约需 \(2\Psi\) 字节；使用 Adam 时，通常还有 FP32 主权重、一级动量和二级动量，各占约 \(4\Psi\) 字节。于是仅模型状态的常见粗略估算是：

$$ 2\Psi\;\text{(BF16 参数)} + 2\Psi\;\text{(BF16 梯度)} + 4\Psi\;\text{(FP32 主参数)} + 4\Psi\;\text{(一阶动量)} + 4\Psi\;\text{(二阶动量)} =16\Psi \text{bytes}.

把这笔账落到 7B 模型上，可以逐项写出来。参数 14GB、梯度 14GB、FP32 主权重 28GB、一阶动量 28GB、二阶动量 28GB，合计约 112GB。若有 8 张 80GB A100，仅模型状态占用就低于总显存的一半，但 activation 会随 batch、序列长度和 checkpointing 策略继续增加。估算显存时宁可保守，因为 kernel 临时空间、通信 buffer、框架缓存和碎片都会消耗额外空间。 $$

对一个 7B 参数模型，这一部分已约为 112 GB，尚未计入 activation、临时 buffer、通信 buffer 和框架开销。具体数值会随混合精度、优化器、是否保存主权重而变，但课上要传达的是：训练显存里最占空间的常常是训练要附带保留的状态，而不只是模型文件有多大。

以一个 7B 参数模型为例，仅权重和训练附带状态就可能占据极大的显存。

图：参数量由各层的权重矩阵之和决定；它是显存账本的起点，却只是训练显存的一部分。

这里每个位置用什么精度是有讲究的。权重和梯度在计算过程中可以用 bf16 这类低精度来省显存、省带宽；但优化器状态和主权重往往需要更高的 fp32 精度。原因是梯度要反复累积、更新，低精度累加很容易把很小的更新量损失掉，导致数值不稳定。于是常见的混合精度做法是，前向/反向用 bf16 跑得快，梯度算完后以 fp32 精度更新主权重，优化器的动量、方差也按 fp32 保存。这也是为什么训练显存里会同时出现 bf16 计算副本和 fp32 精度副本，账本要按两种精度分别估算。

Activation 与完整显存账本

激活是另一笔账。反向传播要用到前向阶段的中间结果，因此常规训练会保存多层 activation；sequence length、batch、hidden size 和层数一大，activation 也可能反过来成为首要瓶颈。activation checkpointing 的思路是少保存一部分，反向时重算，以额外计算换显存；它与 ZeRO 解决的不是同一笔开销。

checkpointing 可以用一条简单的时间线理解。普通训练在每个矩阵乘后保留中间结果，反向直接用；checkpointing 只保留某些边界处的输入，反向到该段时重算中间结果并接着求梯度。它减少的是激活显存，增加的是前向计算次数，通常不会减少优化器状态。ZeRO/FSDP 减少的则是参数、梯度和优化器状态的副本或分片。

训练与推理的账本也不同：推理主要保存权重和 KV cache，训练还要额外保留梯度、优化器状态和 activation，两者量级不能混用。

图：训练显存大头在模型状态与 activation，推理大头在权重与 KV cache；估算先分清在解决哪种负载。

checkpointing（activation recomputation）只保存若干边界 activation，反向时重跑中间前向以换取显存。它减少的是保存量，增加的是计算和可能的通信；不能把它与持久化检查点混为一谈。后者是故障恢复用的模型/优化器快照。选择并行策略前应分别列出参数、梯度、优化器状态、activation 与通信 buffer，而不是只报一个显存占用。

ZeRO / FSDP：把模型状态分片

ZeRO 的三个阶段可以按切分对象记。ZeRO-1 切优化器状态，ZeRO-2 再切梯度，ZeRO-3 连参数也切。对于前述 Adam 的 16 bytes/参数，如果四张卡完全分片，每张卡的模型状态平均可降到约 4 bytes/参数；通信则从梯度 AllReduce 变为前向或反向按需 AllGather 参数，以及 ReduceScatter 梯度。

FSDP 可以理解为 PyTorch 中这类分片训练的工程实现。某一层计算前，各 rank 临时 AllGather 出该层完整参数；计算后释放完整副本，反向再按需收集，梯度用 ReduceScatter 分片写回。它和 ZeRO-3 的核心思想一致，把“每卡都复制完整模型状态”改成“每卡长期只保存一个 shard”。

普通 DP 的问题很明显：每张卡都存完整参数、完整梯度和完整优化器状态，计算按数据分了，状态却没有分。ZeRO（Zero Redundancy Optimizer）按阶段去掉这些冗余副本；FSDP（Fully Sharded Data Parallel）是围绕参数完全分片的一类实现。

• ZeRO-1：分片优化器状态。每个 rank 只长期保存约 \(1/p\) 的 Adam 状态。
• ZeRO-2：在 ZeRO-1 基础上再分片梯度。完成反向后用 ReduceScatter，让每张卡只留下自己负责的梯度块。
• ZeRO-3 / FSDP：连模型参数也分片。每张卡长期只持有参数、梯度和优化器状态的一个分片。

分得越彻底，单卡常驻状态越少，但运行时需要更多通信。ZeRO-3 在算某一层之前先 AllGather，把该层的完整参数临时凑齐；用完前向后释放不再需要的完整参数。反向到这一层时再 AllGather 参数，计算出梯度后用 ReduceScatter 把已规约的梯度留在对应 owner 上，由 owner 更新自己的参数分片。

图：FSDP 在每次用到某层参数前临时 AllGather，算完即释放；通信换取了常驻显存的大幅下降。

图：ZeRO 依次分片优化器状态、梯度和参数；常驻显存下降，运行时 AllGather/ReduceScatter 增加。

因此 ZeRO 并不等于免费省显存。它选择了用通信换内存：模型因为显存根本无法训练时，这个交换是必要的；网络带宽不足、分片太细或每层过小而频繁 AllGather 时，吞吐反而可能下降。看配置时应区分两类问题：参数状态太大时考虑 ZeRO/FSDP；activation 太大时还要考虑 checkpoint、序列切分或减少 micro-batch。

流水线并行

流水线并行¹⁴（pipeline parallelism, PP）沿网络深度切分。若模型有很多层，可以让 GPU 0 负责前几层，GPU 1 负责中间层，GPU 2 负责后几层。前向时 activation 从第一个 stage 依次传到后面；反向时梯度再反向传回。每张卡只放本 stage 的层，因而可容纳单卡放不下的深模型。

最朴素的排程会产生 pipeline bubble：GPU 0 先忙，其他 stage 在等 activation；前向全部结束后，反向又从最后一个 stage 才开始，前面的卡再次等待。卡的数量越多、每轮只有一个大 batch，空闲比例越明显。

图：朴素流水线里，一个 batch 串行穿过各 stage，前后期的 stage 大量空闲；这种空闲比例叫 bubble。

处理办法是把 batch 切为多个 micro-batch，形成填充流水线。GPU 0 算完 micro-batch 0 后就把它交给 GPU 1，自己立刻开始 micro-batch 1；GPU 1 接到 0 后开始计算，后续 stage 也依次被填满。等流水线稳定下来，多个 stage 可以同时处理不同 micro-batch 的前向或反向。bubble 没有消失，但被更多 micro-batch 摊薄。

图：micro-batch 把大 batch 切小，流水线各 stage 因而能同时处理不同 micro-batch，摊薄 bubble。

PP 的通信主要是相邻 stage 间的 activation/gradient P2P 传输，延迟和负载均衡都很重要。如果某个 stage 的层更慢，整条流水线会被它卡住；切层时不能只按层数平均，还要考虑每层计算量、激活大小和显存。

流水线 bubble 可以用一个极小调度表看清楚。假设 4 个 stage、2 个 micro-batch，F 表示前向，B 表示反向，横轴是时间。若没有 micro-batch 交错，GPU 0 在第 1 个前向完成后要等后续 stage 逐步处理，反向又从 GPU 3 开始返回。画在时间线上，两端 GPU 有明显空闲，这些空闲就是 bubble。

GPU0: F0 F1 ........ B1 B0
GPU1: .... F0 F1 ... B1 B0
GPU2: ....... F0 F1 B1 B0
GPU3: .......... F0 B0 F1 B1

增加 micro-batch 数量、让不同 micro-batch 的前向和反向交错，可以填掉一部分空档，但通信次数和调度复杂度也随之增加。

张量并行

PP 是按层切，张量并行（tensor parallelism, TP）则把一层内部的大矩阵乘法切到多张 GPU。Transformer 的 attention 投影和 FFN 都包含很大的矩阵乘法，这往往是计算最重的地方。

考虑线性层 \(Y=XW\)。若按 \(W\) 的列切分：

每张 GPU 可独立算一个 \(XW_i\)，得到 \(Y\) 的一部分列；若后续需要完整的 \(Y\)，再做 AllGather。若按行切分，则每张卡先算局部乘积，最后用 AllReduce 把各部分相加。不同的切法决定通信发生在算子前还是算子后。

用数值说明列切和行切的区别。设

完整结果是 \(XW=[5,7]\)。若按列切 \(W_0=[1,2]^T\)、\(W_1=[1,3]^T\)，GPU 0 算 \(XW_0=[5]\)，GPU 1 算 \(XW_1=[7]\)，需要时 AllGather 拼成 [5,7]。若按行切 \(W\)，两块分别是 [[1,1]] 和 [[2,3]]，GPU 0 得 [3]，GPU 1 得 [4]，必须相加才有 [7]。列切天然得到输出列分片，行切天然得到需要求和的部分和。

FFN 可以写成

其中 \(\phi\) 是逐元素激活函数。第一层把 hidden size 扩到较大的中间维度，第二层再投回去。逐元素激活不会混合不同列，因此可在各 GPU 的局部结果上直接执行；恰当安排 \(W_1\) 的列切分与 \(W_2\) 的行切分，可把本来两次的全量通信压缩为较少的同步点。这就是 TP 在代码里常见列并行线性层和行并行线性层配对的原因。

TP 的通信频率很高，因为几乎每层都要在组内交换中间结果。因此通常把一个 TP group 放在同一台机器、使用 NVLink/NVSwitch 等高速互连；把频繁的 TP 通信跨慢网络，会非常不划算。

图：张量并行把一层内的矩阵乘切开，卡间需要频繁 allreduce；因此 TP group 应放在 NVLink 等高速互连内。

多维并行与通信计算重叠

真实的大模型训练常把几种策略叠起来：节点内做 TP，按层做 PP，剩余维度做 DP，再用 ZeRO/FSDP 分片状态。这常被称为 3D 并行，但不必被名字吓住。读一组并行度配置时，把总 GPU 数写成各维度的乘积，再逐维问：

图：总 GPU 数等于各并行维度的乘积，例如 GPT-3 的配置 \(8\times 8\times 60=3840\) 块 A100；维度越多，切分越细，通信模式也越多样。

以 GPT-3 175B 为例，可把 8 块 A100 通过 NVLink 组成一个节点作为 TP 组，8 个节点按层做 PP，再复制 60 份做 DP，得到 3840 块 GPU 的集群配置。每个维度都对应一类通信：TP 在节点内频繁 AllReduce，PP 在相邻 stage 之间传 activation/gradient，DP 需要全局 AllReduce 梯度。多出来的 SP、CP 沿序列维切分，EP 则按专家切分并触发 All-to-All。面对一个具体配置，先算清这些 GPUs 是怎么分组的、组间传什么，比背维度名字有用。

检查并行配置时，把总卡数和乘积逐项写出，例如 3840 与 \(8\times8\times60\)，随后给每个维度标注通信对象；TP 是节点内高频集合通信，PP 是相邻 stage 的点对点，DP 是全局梯度同步。最后核对每个组的大小是否匹配物理拓扑。TP 放在慢网络两端是最常见的错误之一，因为每层都可能等待通信。

• DP 组里的卡拿不同数据，何时 AllReduce 梯度？
• TP 组里的卡共享一层计算，哪个线性层后要 AllGather 或 AllReduce？
• PP 的相邻 stage 之间传什么 activation 和 gradient？
• ZeRO 分片后，哪一层前需要临时 AllGather 参数？

通信不一定必须停在计算之后才开始。反向传播从最后一层往前走，某层梯度就绪后可以立刻启动对应 bucket 的 AllReduce，同时 GPU 继续计算更前面的层。若通信在计算结束前完成，它的耗时被隐藏；若网络仍未完成，GPU 才需要等待。bucket 的大小过小会让启动延迟累积，过大又会推迟通信启动，因而是一个需要实测的折中。

图：梯度按层/按 bucket 就绪即启动集合通信，不等整个反向结束；重叠是否成立要看网络是否在计算完成前清空。

手动把 TP/PP/DP/SP/CP/EP 组合成一套合理配置并不容易，于是有了自动并行。它把怎么切当作问题交给机器：输入模型计算图与硬件拓扑（若干节点、节点内 NVLink、节点间 IB），输出每个算子的切分方式、流水线 stage 划分和 micro-batch 大小。

图：自动并行的关键是分而治之——层间划分、层内切分、设备映射三者解耦，让每一层找局部最优，整体逼近手调结果。

这类工具（Alpa 之外，PyTorch DTensor、OneFlow SBP 也是同类思路）把设备网格抽象出来，统一描述张量在设备间如何分布，从而让同一份模型代码在不同并行度配置下运行。它的价值不在替代工程师手调，而在把可复现、可搜索的并行空间交给算法，减少每次换模型/换集群都要重写配置的负担。

Transformer 层的算力估算

要判断长序列到底难在哪，先把一层 Transformer 的算力分清楚。以一层的 MLP 与 attention 为例，以 batch \(B\)、序列长 \(S\)、隐宽 \(H\) 计，各部分计算量大致如下。

图：投影与 FFN 随 \(S\) 线性增长、随 \(H\) 平方增长；attention 的两个矩阵运算随 \(S\) 平方增长、只随 \(H\) 线性增长。

关键在 attention 的 \(S^2H\) 项。把这一层的总计算写成 \(8BSH^2+4BS^2H+6BSH^2\)，当 \(S\ll H\)（短序列）时，\(SH^2\) 项占主导，FFN 与投影是瓶颈；当 \(S>H\)（长序列）时，\(S^2H\) 项爆炸，attention 变成瓶颈。

这里的系数来自矩阵乘的形状。一次 \([m,k]\times[k,n]\) 矩阵乘大约需要 \(2mkn\) 次乘加。QKV 和输出投影的形状都围绕 \(S\times H\) 与 \(H\times H\)，合起来给出 \(BSH^2\) 量级；\(QK^T\) 是 \(S\times H\) 乘 \(H\times S\)，乘 \(V\) 是 \(S\times S\) 乘 \(S\times H\)，所以各自都含 \(S^2H\)。做实验时，把每个矩阵的 shape 代入这个公式，就能解释 profiler 里 GEMM 和 attention kernel 的相对耗时。

图：横轴是序列长 \(S\)，纵轴是计算量；attention 的二次增长让它在一段序列长度后反超 MLP，成为长上下文的主要成本。

这正是序列并行、FlashAttention 和 ring attention 要处理的对象：attention 的 \(S^2\) 项既带来巨大计算量，也让中间 attention 矩阵和 KV cache 无法完整物化或驻留单卡。把范围缩小到 \(S\ll H\) 或 \(S>H\) 两种情况，再决定投影/FFN 该切分还是 attention 该分块，比一刀切更有依据。

长上下文与序列并行

长上下文会让 attention 的中间量和 KV cache 迅速增大。即使参数已经由 TP 或 FSDP 分片，单张卡仍可能放不下完整的序列。序列并行（sequence parallelism）和上下文并行（context parallelism）于是沿 sequence 维分割 token 或 K/V 块。

以 ring attention 为例，每张卡保存一段 Q/K/V。它先拿本地 Q 与本地 K/V 算一块 attention，然后把 K/V 块沿环传给下一个 rank，同时接收上一个 rank 的块；经过若干轮，每个 Q 块都和完整上下文的 K/V 做过计算。这样不必把所有 K/V 同时复制到每张卡，峰值内存随序列分片下降，代价是多轮环形通信。

图：每张卡只持有部分 K/V 块，通过循环传递与本地 Q 逐步和全局 K/V 交互；通信次数随分片数上升，显存峰值随之下降。

这里有一个数值细节，softmax 不能简单把每一块各自归一化后再相加，因为全局分母依赖所有块。可维护运行中的最大值 \(m\)、归一化量 \(l\) 和加权结果，块到来时用新的最大值重新缩放旧累计量。这类 online softmax 保持数值稳定，使得 attention 可以按块流式计算；FlashAttention 等高效实现同样依赖分块计算、不物化巨大注意力矩阵的思想。

假设前一块 softmax 的分母是 \(l_1=2\)，累计加权和是 \(u_1\)；后一块的最大值更大，旧分数在全局 softmax 中必须整体缩小。online softmax 记录当前全局最大值 \(m\)，遇到新块时先算 \(e^{m_{old}-m_{new}}\)，用它缩放旧的 \(l\) 和 \(u\)，再累加新块。这样不需要先看到全部分数，也能得到和完整 softmax 相同的数学结果，同时避免大量 \(e^x\) 上溢或下溢。

MoE 与 Expert Parallel

Mixture of Experts¹⁵（MoE）层用 router 为每个 token 选择少数几个 expert，而不是让所有 token 经过同一个 FFN。总参数可以非常大，单个 token 的计算量却只与被选中的少数专家有关。这是 MoE 能扩大模型容量的原因。

代价落在路由和网络上：某张卡上的 token 可能需要送去别张卡拥有的 expert，计算后结果还要按原 token 顺序送回，通常形成 All-to-All。router 若总偏向少数专家，就会出现有的专家/卡很忙、有的很闲；因此需要容量限制和负载均衡损失。MoE 的瓶颈不一定是矩阵乘法，可能是 token 重排、all-to-all、专家不均衡或跨节点拓扑。

一个 token 在 MoE 层里的路径可以这样读。设每层有 64 个 expert，每个 token 选 top-2。输入仍是 [B,S,H]，router 对每行输出 64 个分数，softmax 后取最大两个，例如 expert 7 权重 0.6、expert 42 权重 0.4。系统把所有 token 按 expert 分组，expert 7 收到自己的 token 矩阵后执行 FFN，expert 42 同样执行；最后每个 token 的输出是 \(0.6E_7(x)+0.4E_{42}(x)\)，再放回原序列位置。参数规模来自 64 个 expert，单个 token 计算量却只对应 2 个 expert。

有用的话请给我个 star => 本站总浏览