Lecture 23: Parallel Deep Learning: Model and Pipeline Parallelism

目录 · ← l21 · l23 →

Lecture 23: Parallel Deep Learning: Model and Pipeline Parallelism

1. 章节标题与概述

Lecture 23: Parallel Deep Learning: Model and Pipeline Parallelism

• 本讲核心问题：当模型本身大到一张 GPU 装不下时（GPT-3 175B 参数的 FP16 权重就有 350 GB，而 A100 只有 40/80 GB 显存），”数据并行”这条路直接失效——它要求每张卡保存一份完整模型副本。本讲回答两个问题：(1) 如何把模型的参数/计算本身切开（Model Parallelism，尤其是层内的 Tensor Model Parallelism）；(2) 切开之后如何避免设备互相等待（Pipeline Model Parallelism 与 micro-batch 流水线），并给出通信量、气泡（bubble）与设备利用率的定量模型。
• 涉及的主要硬件/软件机制：跨 GPU 的集合通信原语（AllReduce / AllGather / ReduceScatter / 点对点 send-recv）与其运行的分层互连（节点内 NVLink/NVSwitch ≈ 300 GB/s 单向、节点间 InfiniBand HDR ≈ 25 GB/s、PCIe Gen4 ≈ 32 GB/s）；GPU 内部的分块 GEMM 执行模型（shared memory tiling、warp 级 SIMD 执行、tensor core）；软件层的张量并行算子分解（Megatron-LM 的 f/g 算子、多头注意力的头并行）、micro-batch 流水线调度（GPipe）、以及 DP×PP×TP 的三维并行（DeepSpeed 3D parallelism）。
• 在并行计算知识体系中的角色：本讲是”并行模式 → 真实大规模系统”链条的收口。它把前面各讲的四个基础工具——fork-join/SPMD 编程模型（数据并行、流水线级就是 SPMD 进程）、work-span 与 Amdahl 定律（气泡就是无法并行的串行段）、算术强度与 Roofline（micro-batch 变小会把 GEMM 推入带宽受限区）、通信延迟/带宽模型（α-β 模型与集合通信代价）——第一次真正用在”一个可运行的工业级系统”（Megatron-LM / GPipe / DeepSpeed）上。它同时给出了一条贯穿全课的设计准则：并行度不是免费的，每一层并行都对应一类通信，而通信的频率决定了它能用多慢的链路。
• 配套材料：
  • 本讲讲义（已公开）：https://www.cs.cmu.edu/~418/lectures/25-parallel_deep_learning_model_pipeline_parallel.pdf（HTTP 200 可直接匿名下载，2.9 MB，共 34 页，文件时间戳 2025-11-01）。讲义首页署名为客座讲师 Zhihao Jia (Stanford University)，页脚写 CMU 15-418/15-618, Fall 2025、编号 Lecture 25；而 Fall 2026 日程表中本讲编号为 Lecture 23（Oct 26）。这是课程”讲义沿用旧学期版本”的正常现象（同一份 deck 在不同学期编号不同），并非错误。
  • 本笔记的事实依据：extracted/25-parallel_deep_learning_model_pipeline_parallel.txt（上述 PDF 的逐页抽取文本，标记 ===== [slide i/34] =====）。
  • Fall 2026 日程表（已公开）：https://www.cs.cmu.edu/~418/schedule.html，第 23 讲 Oct 26 标题为 Parallel Deep Learning (model and pipeline parallelism)。该行的 slides/video 链接仍以 HTML 注释形式保留（注释原文：slide/video from a previous offering; uncomment when posted for Fall 2026），因此 Fall 2026 本讲的 slides 与录像在日程表上尚未发布；注释中记录的历史公开录像为 YouTube 上的”模型并行”与”流水线并行”两段。Fall 2026 录像在日程表中被注释隐藏 → 属未发布。
  • 补充材料（已公开）：cs149_supp/dnninference.txt = Stanford CS149 Fall 2025 Lecture 9 Efficiently Evaluating DNNs（公开课程材料，75 页）。本笔记用它补充 GEMM 分块/向量化、算术强度与 Roofline、算子融合（loop fusion / flash-attention）、显式 vs 隐式 GEMM 等底层性能背景——这些正是张量并行的”每一格方块”要落到硬件上执行的东西。
  • 需登录 / 未公开：历史学期讲义 PDF 位于 /afs/cs/academic/class/15418-*/public/（需 CMU 登录）；Ed 讨论区、Autolab、Canvas 均需登录。Fall 2026 授课教师为 Brian Railing 与 Dimitrios Skarlatos，课程由 Kayvon Fatahalian 创建（本讲未在讲义中指定具体作业/评分点，笔记中也不涉及）。
  • 注：f22_25_dnn.txt、s23_25_dnn.txt 两个抽取文件为空（仅 2 行），无可引用内容，本笔记未使用。

2. 核心概念与硬件/软件架构图解

2.1 从数据并行到模型并行：显存墙（Memory Wall）

定义与目的：数据并行（Data Parallelism）把训练集切成 batch 分给 N 张卡，每张卡用完整模型副本算自己那批数据的梯度，再用 AllReduce 聚合梯度，然后各自更新参数（w_i := w_i - γ∇L(w_i)）。它的目的是把”数据的并行度”变成算力，同时前向/反向过程中零通信（只在梯度同步时通信一次）——这是它最大的优点。

它是什么（直观解释）：数据并行就像一家连锁餐厅的 N 个分店，每家都买了全套厨具和完整菜谱：客人（数据）分流到各店，各店独立出餐，晚上汇总大家的经验（梯度）。问题是”菜谱”（模型参数）一旦厚到装不进单店的厨房（显存），这套模式就崩了。

为什么必须换路子（容量数字，来自课程前一讲的数据并行讲义）：

2800 GB / 80 GB ≈ 至少 35 张 A100 才”放得下”模型状态（还没算激活值）；再加数据并行要求每卡一份完整副本，这条路彻底走不通。于是必须切模型。

2.2 模型并行（Model Parallelism）：子图切分与设备放置

定义与目的：把模型（计算图）拆成多个子图，分给不同设备；设备之间传递中间结果（激活/梯度）；没有梯度同步（讲义的原文表述：No synchronization of gradients）。它解决的是”模型放不下”的问题，而不是”算得更快”的问题。

它是什么（直观解释）：模型并行像把一条装配线拆成几个工位，每个工位只装一部分零件（一层或几层），半成品用传送带送到下一个工位。注意：如果只切、”不分时”，同一时刻只有一个工位在干活，其他工位在等——这正是下一节流水线并行要修的毛病。

架构/机制图解（图 1：三种并行策略的”谁存什么、谁算什么”）：

                       (a) 数据并行 Data Parallelism
      ┌───────────────┐   ┌───────────────┐   ┌───────────────┐
      │    GPU 1      │   │    GPU 2      │   │    GPU N      │
      │  完整模型 W   │   │  完整模型 W   │   │  完整模型 W   │  ← 每卡一份完整副本
      │  minibatch 1  │   │  minibatch 2  │   │  minibatch N  │
      └───────┬───────┘   └───────┬───────┘   └───────┬───────┘
              │ ∇L₁               │ ∇L₂               │ ∇L_N
              └──────── AllReduce(求和, 每 iter 一次) ───────┘
        前向/反向: 零通信                 通信: 参数尺寸 M 的梯度, 1 次/iter

                       (b) 模型并行 (子图切分, 无 micro-batch)
   时间 →
   GPU1 [op1 op2]▓▓▓▓▓▓░░░░░░░░░░░░░░░░░░░░░░░░
   激活跨设备 ─────────┐
   GPU2               [op3 op4]▓▓▓▓▓▓░░░░░░░░░░   ← 只有 1 台在算, 其余在等
                      └────────────┐
   GPU3                            [op5 loss]▓▓▓▓▓
        设备利用率 ≈ 1/p, 增加 p 只会让等待更长

                       (c) 张量并行 (层内切分, 每层都要通信)
   GPU1:  W 的第 1 列块 → y₁ ┐
                             ├─ 每层内部一次 AllReduce / 点对点
   GPU2:  W 的第 2 列块 → y₂ ┘
        两个维度同时被切开: 参数切开(能装下) + 每层都通信(必须快)

关键性能特征：

• 延迟：模型并行的前向是串行链（op1→op2→…→loss），单样本延迟几乎不会因为加卡而下降，反而因为每层要跨设备传激活而上升（一次设备间激活传输 = α（启动延迟）+ bytes/BW）。
• 带宽：跨设备激活尺寸 ∝ B·S·H（batch × 序列长 × 隐层），与计算量（∝ B·S·H²）相比只差一个 H。H 越大、batch 越大，通信相对越便宜。
• 吞吐量：无流水线时吞吐 ∝ 1/p 恶化（见上图 b）。所以模型并行的价值在容量，吞吐要靠流水线并行补回来。

放置（placements）为什么难：讲义第 5 页给出两张对照图——”在 4 张 GPU 上训练 RNN”与”在 4 张 GPU 上训练常规网络”的最优切分完全不同（RNN 有时间步依赖，只能按层切；CNN/MLP 可以按通道切）。第 6 页给出的自动化思路是用 ML 来调 ML 的放置：RL Agent 输出 device placement → 在真实并行机上跑 → 得到 runtime performance 作为 reward → 迭代（引文：Device placement optimization with reinforcement learning, A. Mirhoseini et al.；图中还标了 ColocRL’s neural architecture）。第 7 页进一步给出数据并行 + 模型并行混用的方向——这是后面 3D 并行的雏形。

2.3 张量模型并行（Tensor Model Parallelism）：两种切分方式与通信量记账

定义与目的：不在”层与层之间”切，而是在一层/一个算子内部切参数与梯度（讲义原文：Partition parameters/gradients within a layer/operator）。目的：把权重的显存占用和 GEMM 的计算量同时按 E 份摊到 E 张卡上，同时让切分后的通信量比数据并行小几个数量级。

它是什么（直观解释）：模型并行是”把装配线切成工位”，张量并行是”把同一道工序的原料横向切开，几个人同时处理同一道菜“。两种切法：

• partition output（按输出切）：两个厨师拿到完全相同的食材（完整输入 x），各自按菜谱的一半行做菜，产出不同的菜（y₁、y₂）——无需合并。
• reduce output（按输入切）：两个厨师各拿到半份食材（x₁、x₂），各自做同一道菜，最后必须合锅（y = y₁ + y₂）。

架构/机制图解（图 2：Transformer 一层的张量并行数据流，Megatron-LM 写法）：

        ┌───────────────── Transformer 层 (张量并行度 E=2) ─────────────────┐
        │                                                                  │
 X (完整, 复制)                                                            │
   │  ── identity 层: 前向不通信, 反向做 AllReduce ──                        │
   ├──────────────► GPU1: A₁ [K × N/2] ──► Y₁ ─┐                            │
   │                                        ├─ GeLU（逐元素, 无需通信）      │
   └──────────────► GPU2: A₂ [K × N/2] ──► Y₂ ─┘                            │
        Y = [Y₁ | Y₂]  ← 拼接(不是求和), 通信被推迟到第二层之后              │
                     │                                                    │
   GPU1: B₁ [(N/2) × O] ──► Z₁ ┐                                          │
                               ├─ AllReduce(求和) ──► Z = Z₁ + Z₂  (g 层)   │
   GPU2: B₂ [(N/2) × O] ──► Z₂ ┘   ← 本层唯一通信: B·C_out 个元素           │
                     │   ── 反向: g 是 identity, f 做 AllReduce ──           │
        ┌────────────▼────────────┐                                       │
        │ 多头注意力: 按 head 切   │                                       │
        │ head i: Qᵢ,Kᵢ,Vᵢ 独立    │  Yᵢ = softmax(QᵢKᵢᵀ/√d) Vᵢ   (reduce 输出)│
        │ Z = Concat(Y₀..Y_h) W_o  │  ← W_o 行切, 需 AllReduce             │
        └─────────────────────────┘                                       │
        └──────────────────────────────────────────────────────────────────┘
   规则: 列并行(f, partition output) 后面接行并行(g, reduce output), 交替出现,
        把 AllReduce 放在「非线性之后、尺寸更小」的张量上。

通信量记账（讲义第 8–13 页的核心结论，务必按原口径理解）：设一层做 y = W x，W 是 C_out × C_in，batch 为 B：

读表要点：数据并行的通信量是参数尺寸 C_out·C_in（与 batch 无关，但每卡都要存全量参数）；张量并行的通信量是激活尺寸 B·C_in 或 B·C_out（与参数无关，所以千亿参数模型的张量并行通信量并不随参数量爆炸）。两者的量级差 C_out·C_in / (B·C_out) = C_in / B：当 C_in = 4096、B = 8 时相差 512 倍。这正是张量并行能扩展到数百卡的根本原因。

而代价是频率：数据并行每个 iteration 只同步一次梯度；张量并行每一层都要通信（Transformer 每层 2 次 AllReduce）。通信量小 × 频率高 × 对延迟敏感，三者共同决定了它只能跑在最快的链路上。

“最好的策略取决于模型和机器”（讲义第 13 页结论） ——CNN 的分层策略（讲义第 17–18 页）：

• 卷积层：占 90–95% 的计算量，只占 5% 的参数，但有非常大的中间激活。→ 用数据并行（参数小，副本无所谓；激活大，正好靠数据并行摊薄）。
• 全连接层：占 5–10% 的计算量，占 95% 的参数，中间激活很小。→ 用张量模型并行（参数大，必须切开；激活小，通信便宜）。
• 讲义的第 18 页直接给出结论行：Data parallelism for convolutional layers / Tensor model parallelism for fully-connected layers。

Transformer 的层内切分（讲义第 19–26 页）：一层 Transformer = 多头自注意力 + 两个全连接层。

• 注意力：Attention(Q,K,V) = softmax(QKᵀ/√d)V；多头版把输入分别做线性变换得到 Q_i,K_i,V_i，各头独立算 Y_i = softmax(Q_iK_iᵀ/√d)V_i，最后 Z = MultiHead(Q,K,V) = Concat(Y_0,…,Y_h)W_o。每个 head 分到一张卡（head 并行），W_o 按行切（reduce output，需 AllReduce）。
• 为什么用多头（讲义第 24 页的一行问答）：Why multi-head attention? → Reduce complexity of matrix multiply!（把 N×N 的大矩阵乘拆成 h 个 N×d/h 的小矩阵乘，且天然给出 h 路并行）。
• 两个全连接层：Y = GeLU(X A)（A 按列切，partition output）、Z = Dropout(Y) B（B 按行切，reduce output）。讲义第 21 页明确标注 Weights: Partition on Col dim、Y = [Y₁ Y₂]、Z = Z₁ + Z₂，并问了一句 “Why switch?”（为什么这样交替切）——答案见第 7 节思考题 Q2。
• 讲义第 26 页给出规模结论：Megatron-LM 用”数据并行 + 模型并行”组合扩展到 512 张 GPU。

2.4 流水线模型并行（Pipeline Model Parallelism）：micro-batch 与气泡

定义与目的：让被切成 p 级的模型真正同时在干活。做法是把一个 mini-batch 再切成 m 个 micro-batch，让不同 micro-batch 的不同级在同一时刻重叠执行（一级算 micro-batch i 的前向时，下一级在算 i-1 的前向）。目的：把 2.2 节图 1(b) 中”只有一台机器在算”的空转填满。

它是什么（直观解释）：像洗车流水线：一辆车（micro-batch）依次经过”冲水→打沫→擦干→上蜡”四个工位；如果一次只让一辆车走完全程，工位 2/3/4 大部分时间在晒太阳；同时放 8 辆车进去（m = 8），每个工位就一直在忙。但流水线要”灌满”和”排空”：开头的 3 个节拍（fill）和结尾的 3 个节拍（drain）没有满负荷，这段空转就是气泡（bubble）。

讲义给出的定量模型（第 30–31 页）：设 m = 一个 mini-batch 里的 micro-batch 数，p = 流水线级数，每级处理一个 micro-batch 的前向/反向时间都是 t_f / t_b，则

        BubbleFraction = (p-1) * (t_f + t_b)  /  ( m*t_f + m*t_b )  =  (p-1) / m

口径说明（非常重要，容易算错）：讲义这个式子的分母 m*(t_f+t_b) 是理想计算时间（没有任何气泡时单级应该忙的时间），所以 (p-1)/m 是”气泡相对于理想计算量”的比例，它不是气泡占墙钟时间的比例。整条流水线的墙钟时间 = (m + p - 1) * (t_f + t_b)（fill + steady + drain），因此气泡占墙钟时间的比例 = (p-1)/(m+p-1)，设备利用率 = m/(m+p-1)。两个口径都成立，只是分母不同；CMU 讲义的 (p-1)/m 用的是前者，论文/工程里常用后者。本笔记第 4 节会同时给出两者并做数值对照。

架构/机制图解（图 3：GPipe 调度 vs 定点重叠，含气泡）：

  (a) GPipe 调度: 全前向 → 全反向  (m=4 个 micro-batch, p=4 级; F=前向, B=反向)
      时间 →
 stage1 │F0│F1│F2│F3│  │  │  │  │B3│B2│B1│B0│
 stage2 │  │F0│F1│F2│F3│  │  │B3│B2│B1│B0│  │
 stage3 │  │  │F0│F1│F2│F3│B3│B2│B1│B0│  │  │
 stage4 │  │  │  │F0│F1│F2│F3│B3│B2│B1│B0│  │
         └──── fill ────┴── 稳定 ──┴── drain ──┘
         气泡 = (p-1) 个 (F+B) 时间槽, 总时间 = (m+p-1)(t_f+t_b)
         利用率 = m/(m+p-1) = 4/7 = 57%（m=4,p=4）

  (b) 增大 m（m=16, p=4）: 稳定段变长, 气泡占比从 3/7=43% 降到 3/19=16%
      时间 →
 stage1 │F0..F15│B15..B0│
 stage2 │ F0..F15│B15..B0│     每个级几乎全程忙碌
 stage3 │  F0..F15│B15..B0│
 stage4 │   F0..F15│B15..B0│
         └──────── 稳定段 ∝ m ────────┘

  (c) 1F1B（PipeDream/DeepSpeed 思路, 讲义未展开, 作为延伸）: 前向/反向交错,
      每个级最多只需缓存 ≈ p 个 micro-batch 的激活, 而不是 GPipe 的 m 个
      stage1 │F0│F1│F2│F3│B0│F4│B1│F5│B2│B3│B4│B5│...
                                    ↑ 前向与反向交替, 显存 ↓, 气泡仍 ≈ (p-1)/(m+p-1)

关键性能特征（延迟/带宽/吞吐）：

• 吞吐：稳态时每个时间槽有一级产出一个 micro-batch 的反向结果，所以吞吐 ∝ 1/(t_f+t_b)（每级），与 p 无关——加级数不提升吞吐，只提升”能装下的模型规模”。
• 延迟：单个 micro-batch 的端到端延迟 ≈ p*(t_f+t_b)（要穿过全部 p 级），p 越大单样本延迟越高。
• 带宽：级与级之间每个 micro-batch 传一次激活 B_micro·S·H·2 bytes（FP16）。因为可以与计算重叠（下一级的 F 与上一级的 transfer 并行），所需的带宽远低于张量并行。
• 气泡：(p-1)/(m+p-1)。三条改善路径与各自的代价（讲义第 31 页明确列出）：
  1. 增大 m（把 mini-batch 切得更细，或增大 mini-batch）——代价：large mini-batch sizes can lead to accuracy loss（大 batch 影响收敛精度）；切得太细则 small micro-batch sizes reduce GPU utilization（micro-batch 太小，GEMM 的 M 维变小，GPU 打不满，见第 4 节的 Roofline 算例）。
  2. 减小 p（流水线浅一点）——代价：increase stage size（每一级要装更多层，显存压力回到单卡）。
  3. （讲义外的工程手段，作为延伸）减小每级”空转粒度”：interleaved / virtual pipeline stages（把每一级再切成 v 个虚拟级交错调度），把气泡降到 ≈ (p-1)/(m·v)。

2.5 三种并行的对比与 3D 并行（DeepSpeed）

讲义第 32–33 页的总结表（Pros / Cons），完全按原文整理：

讲义第 34 页的收口结论：Training large models requires combining data/model/pipeline and other parallelization techniques —— 即 3D 并行（DeepSpeed 官方博客图示）：数据并行 × 流水线并行 × 张量并行三层嵌套。

架构/机制图解（图 4：3D 并行的硬件/软件映射，512 卡实例）：

                 一层节点内 (NVLink/NVSwitch, 300 GB/s 单向, ~1 µs)
      ┌──────────────────────────── 8 × A100 (80 GB) ────────────────────────────┐
      │  TP group = 8: 同一层的权重被切成 8 份, 每层 2 次 AllReduce              │
      │  ┌────────┐┌────────┐┌────────┐┌────────┐  ┌────────┐     ┌────────┐     │
      │  │ GPU 0  ││ GPU 1  ││ GPU 2  ││ GPU 3  │… │ GPU 6  │     │ GPU 7  │     │
      │  │TP rank0││TP rank1││TP rank2││TP rank3│  │TP rank6│     │TP rank7│     │
      │  └────┬───┘└───┬────┘└───┬────┘└───┬────┘  └───┬────┘     └───┬────┘     │
      │       └────────┴─────────┴─── NVSwitch ───────┴──────────────┘           │
      └───────────────────────────────┬──────────────────────────────────────────┘
                                      │ 流水线阶段边界: 点对点传激活
                                      │ (InfiniBand HDR 25 GB/s 单向, ~2 µs)
      ┌───────────────────────────────▼──────────────────────────────────────────┐
      │  节点 1 = PP stage 1 (存模型的第 1..L/p 层)  ←→ 节点 2 = PP stage 2 ...   │
      └───────────────────────────────┬──────────────────────────────────────────┘
                                      │ 数据并行副本: 每 iter 同步一次梯度
                                      │ (可用最慢链路, 因为频率最低)
      ┌───────────────────────────────▼──────────────────────────────────────────┐
      │  DP group: 第 0..7 份数据副本 (ZeRO 进一步切分优化器状态/梯度/参数)      │
      └──────────────────────────────────────────────────────────────────────────┘

   配置算术: 512 卡 = TP 8 × PP 8 × DP 8
             全局 batch = micro-batch(1 样本) × m(16) × DP(8) = 128 个序列
             每卡参数 = 175 B / (TP 8 × PP 8) = 2.73 B  ← 只有这种三层切分能装下

关键设计准则（贯穿本讲的一条线）：通信频率越高的并行维度，必须放在越快的链路上。张量并行每层都通信 → 必须留在 NVLink 域内；流水线并行每 micro-batch 通信一次且可与计算重叠 → 可以跨节点走 IB；数据并行每 iteration 同步一次 → 放到最外层，甚至可以用较慢的链路。

2.6 让张量并行的每一格方块跑满：GEMM 的执行模型（CS149 补充材料）

张量并行把 GEMM 的维度切开后，每个 GPU 上跑的还是 GEMM。讲义本身不重复 GEMM 优化，但”切分后的 GEMM 还能不能跑满”完全由 CS149 讲过的机制决定：

(1) 分块（blocking/tiling）与算术强度：直接三重循环的 GEMM 对 A、B 没有时间局部性，arithmetic intensity 很低；按 cache 层级分块（L2 → L1 → 寄存器）后，每次从 DRAM 取一个 TILE×TILE 的块就能做 TILE 次 FMA，片内算术强度 ≈ TILE/4 FLOP/byte。张量并行把 N 切成 N/E 后，可用的 tile 数变少，尾部（tail）浪费变大——这是切分粒度不能无限小的硬件原因。

(2) 向量化与数据布局：CS149 给了三种向量化分块 GEMM 的写法：(i) 向量化 i 维（warp 内不同 lane 算不同输出列，splat A 的标量乘 B 的向量）；(ii) 先转置 B 块再对最内层 k 维做向量点积；(iii) SIMD_WIDTH × SIMD_WIDTH 的寄存器块（每个 lane 维护 SIMD_WIDTH 个累加器，同时对 B 的 SIMD_WIDTH 个元素做 FMA），这是现代 GPU tensor core 之前的标准做法，也是”一个 warp 一次算 32×32 个输出”的来源。

(3) 算子融合（loop fusion）：CS149 的例子很直观——分开写 add/mul/add 三个循环的整体算术强度是 1/3（2 load + 1 store 换 1 个数学运算）；融成一个循环后中间结果不出 DRAM，算术强度升到 3/5。在 Transformer 里，这就是 flash-attention 的来源：softmax 的朴素实现要读 5MN+2M、写 3MN+2M 个元素；按行分块、把 softmax(QKᵀ)V 完全融合在片上完成，只需读 MN、写 MN，并且永不显式构造 N×N 的注意力矩阵（N 是序列长度，朴素实现要 N² 空间，长序列直接爆）。张量并行的”每层 2 次 AllReduce”之所以能被容忍，很大程度上靠这类融合把计算侧的带宽压力压下去。

(4) 显式 vs 隐式 GEMM：用现成的高性能 GEMM 库需要把卷积”im2col”成矩阵，DRAM 流量增加 R×S 倍（R、S 是卷积核尺寸）并多占大量存储；隐式 GEMM（implicit GEMM）只在片上共享内存里物化一个子块，既不额外占显存也不增加 DRAM 流量（CUTLASS/cuDNN 的做法）。这就是为什么训练框架的”算子→kernel”映射不是简单调库，而是与并行策略耦合的调度问题。

(5) 低精度：16 bit / 8 bit 已是主流，4 bit 正在进入，极端情况 1 bit。低精度直接使通信量按位宽线性下降——张量并行每层传的激活如果从 FP32 变 BF16，通信量立刻减半，这对通信受限的张量并行是”免费”的加速。

3. 代码示例与性能分析

3.1 示例 1：MPI + OpenMP 实现 Megatron 风格的张量并行 MLP（列并行 → GeLU → 行并行 → AllReduce）

这段代码把讲义第 21 页的 Y = GeLU(X A)（A 按列切，f 算子）+ Z = Dropout(Y) B（B 按行切，g 算子）完整实现成可运行的 MPI 程序，并用单 rank 串行版本做正确性对照。

// tp_mlp.cpp -- Megatron-LM 风格张量并行 MLP：列并行(f) -> GeLU -> 行并行(g) -> AllReduce
// 编译: mpicxx -O3 -march=native -fopenmp -std=c++17 tp_mlp.cpp -o tp_mlp
// 运行: OMP_NUM_THREADS=4 mpirun -np 2 ./tp_mlp
#include <mpi.h>
#include <omp.h>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <random>
#include <vector>

static const int M = 64;    // micro-batch 中的 token 数 (B*S)
static const int K = 256;   // 输入隐层维度 H
static const int N = 512;   // MLP 中间维度 4H
static const int O = 256;   // 输出隐层维度 H

static inline float gelu(float x) { return 0.5f * x * (1.0f + std::erff(x * 0.70710678f)); }

// C[M x n] = A[M x k] * B[k x n]，行主序；OpenMP 按输出行并行
static void gemm(int m, int k, int n, const float* A, const float* B, float* C, bool accumulate) {
#pragma omp parallel for schedule(static)
    for (int i = 0; i < m; ++i) {
        for (int j = 0; j < n; ++j) {
            float acc = accumulate ? C[i * n + j] : 0.0f;
            for (int p = 0; p < k; ++p) acc += A[i * k + p] * B[p * n + j];
            C[i * n + j] = acc;
        }
    }
}

int main(int argc, char** argv) {
    MPI_Init(&argc, &argv);
    int rank = 0, size = 1;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    if (N % size != 0) { if (rank == 0) std::printf("rank 数必须整除 N=%d\n", N); MPI_Finalize(); return 1; }
    const int Nl = N / size;              // 本 rank 的分片宽度 N/E
    const int reps = 20;                  // 重复次数，取最好一次（benchmark 惯例）

    // 所有 rank 生成完全相同的 X / Afull / Bfull
    std::vector<float> X(M * K), Afull(K * N), Bfull(N * O);
    std::mt19937 rng(2024);
    std::uniform_real_distribution<float> dist(-0.5f, 0.5f);
    for (auto& v : X) v = dist(rng);
    for (auto& v : Afull) v = dist(rng);
    for (auto& v : Bfull) v = dist(rng);

    // ---- 切分权重：A 列并行（partition output），B 行并行（reduce output）----
    std::vector<float> A(K * Nl), B(Nl * O);
    for (int k = 0; k < K; ++k)
        for (int j = 0; j < Nl; ++j) A[k * Nl + j] = Afull[k * N + rank * Nl + j];
    for (int i = 0; i < Nl; ++i)
        for (int j = 0; j < O; ++j) B[i * O + j] = Bfull[(rank * Nl + i) * O + j];

    std::vector<float> Y(M * Nl), Z(M * O, 0.0f);
    double best = 1e30;
    for (int r = 0; r < reps; ++r) {
        double t0 = MPI_Wtime();
        gemm(M, K, Nl, X.data(), A.data(), Y.data(), false);   // f: 列并行 GEMM，前向无通信
        for (auto& v : Y) v = gelu(v);                         // 逐元素非线性，无需通信
        gemm(M, Nl, O, Y.data(), B.data(), Z.data(), false);   // g: 行并行 GEMM，得到部分和
        MPI_Allreduce(MPI_IN_PLACE, Z.data(), M * O, MPI_FLOAT, MPI_SUM, MPI_COMM_WORLD); // 前向唯一通信
        best = std::fmin(best, MPI_Wtime() - t0);
    }

    // ---- 单 rank 串行参考实现（完整权重）----
    std::vector<float> Yref(M * N), Zref(M * O);
    gemm(M, K, N, X.data(), Afull.data(), Yref.data(), false);
    for (auto& v : Yref) v = gelu(v);
    gemm(M, N, O, Yref.data(), Bfull.data(), Zref.data(), false);
    float maxerr = 0.0f;
    for (int i = 0; i < M * O; ++i) maxerr = std::fmax(maxerr, std::fabs(Z[i] - Zref[i]));

    const double flops = 2.0 * M * K * N + 2.0 * M * N * O;      // 整个 MLP 的总浮点运算量
    const double comm_bytes_per_rank = 2.0 * (size - 1.0) / size * M * O * sizeof(float);
    if (rank == 0) {
        std::printf("E=%d threads=%d | M=%d K=%d N=%d O=%d\n", size, omp_get_max_threads(), M, K, N, O);
        std::printf("forward 最好一次 = %.3f ms  (%.2f GFLOP/s 全机器, 单 rank %.2f GFLOP/s)\n",
                    best * 1e3, flops / best / 1e9, flops / size / best / 1e9);
        std::printf("每次前向 AllReduce 装载字节/rank = %.1f KB, 最大误差 = %.3e\n",
                    comm_bytes_per_rank / 1024.0, maxerr);
    }
    MPI_Finalize();
    return 0;
}

【代码做什么？】

1. 每个 rank 用相同的随机种子独立生成 X (M×K)、A (K×N)、B (N×O)——这不是”偷懒”，而是真实训练的对应物：X（激活）由数据并行保证每张卡上的副本相同，A/B 在初始化时本来就由同一个随机种子生成（Megatron 只切分、不重采样）。
2. 切分：A 按列切成 E 份（第 e 份 = 每行的 [e·N/E, (e+1)·N/E) 列），partition output；B 按行切成 E 份（第 e 份 = 第 [e·N/E, (e+1)·N/E) 行，所有列），reduce output。两份分片的”缝”必须对齐（A 的列分片宽度 = B 的行分片宽度 = N/E），这是矩阵乘维度匹配的硬性要求。
3. 前向三步：本地 GEMM 得到 Y_e = X·A_e (M × N/E) → 对 Y_e 逐元素做 GeLU（注意：此时不需要任何通信，因为 GeLU 是逐元素算子，作用在”属于自己那一份”的元素上完全合法）→ 本地 GEMM 得到部分和 Z_e = Y_e·B_e (M × O) → 一次 MPI_Allreduce 把 E 份部分和相加得到完整 Z。
4. 正确性验证：用完整权重在单 rank 上串行重算一遍，比较最大绝对误差。
5. 循环 20 次，取最好一次的时间（benchmark 惯例，排除调度抖动）。

实测输出（在共享的登录节点上运行，OMP_NUM_THREADS=4）：

E=4 threads=4 | M=64 K=256 N=512 O=256
forward 最好一次 = 1.631 ms  (20.57 GFLOP/s 全机器, 单 rank 5.14 GFLOP/s)
每次前向 AllReduce 装载字节/rank = 96.0 KB, 最大误差 = 1.907e-05
E=8 threads=4 | M=64 K=256 N=512 O=256
forward 最好一次 = 0.734 ms  (45.70 GFLOP/s 全机器, 单 rank 5.71 GFLOP/s)
每次前向 AllReduce 装载字节/rank = 112.0 KB, 最大误差 = 1.717e-05

误差 ~2e-5 是 FP32 累加顺序改变造成的正常舍入差（不是 bug）。注意这台机器是多人共用的登录节点，E=1/E=2 的测量被内存带宽争抢严重污染（E=1 反而比 E=2 慢），因此趋势结论以 E=4 → E=8 为准：单 rank 吞吐几乎不变（5.14 → 5.71 GFLOP/s），总吞吐翻倍，说明通信开销在 E≤8 时还远未成为瓶颈。

【并行机制与性能解说】

• 进程/线程如何创建、工作如何分配：MPI 起 E 个进程（在真实系统上就是 E 张 GPU）；每个进程内用 #pragma omp parallel for 把 GEMM 的 M 行分给 4 个线程（schedule(static)，每线程 16 行）。两层并行是嵌套的：外层是跨设备的 SPMD（E 路），内层是设备内的 fork-join（线程数 = SM 数 / warp 数）。共享数据（X、A_e、B_e）只读，因此没有数据竞争；Y、Z 是每线程写自己的行，也不冲突（无伪共享风险，因为按行切分，行长度 256~512 字节远超 cache line 64 字节）。
• Work（总工作量）：W = 2·M·K·N + 2·M·N·O + O(M·N)（最后一项是 GeLU）= 2·64·256·512 + 2·64·512·256 + 64·512·10 ≈ 16.78 M + 16.78 M + 0.33 M ≈ 33.9 MFLOP。
• Span（关键路径）：S = W/(E·P_rank) + T_AllReduce，其中 P_rank 是单 rank 的有效算力，T_AllReduce = 2(E-1)·α + 2(E-1)/E · (M·O·4 bytes)/BW（ring AllReduce 的经典形式：2(E-1) 步，每步传 (M·O/E) 个元素）。
  • 代实测值 P_rank ≈ 5.7 GFLOP/s：W/(E·P_rank) 在 E=8 时 = 33.9e6/(8·5.7e9) = 0.743 ms，与实测 0.734 ms 吻合到 1%，说明这一档纯计算主导。
  • 通信侧：M·O = 16384 个 float = 64 KB，E=8 时每 rank 收发 2·7/8·64 KB = 112 KB；若走 IB HDR 25 GB/s，传输 ≈ 4.5 µs，加 2(E-1)=14 步各 ~2 µs 启动延迟 ≈ 33 µs，即 约为计算时间的 4%（实测总时间几乎无变化，一致）。
• 并行度 = Work/Span：通信可忽略时 ≈ E（线性扩展）；但通信项 2(E-1)α 随 E 线性增长，而计算项 W/(E·P) 随 E 下降，因此存在一个最优点：
  • E=32 时：计算 = 33.9e6/(32·5.7e9) = 0.186 ms；通信 = 2·31·2µs + 2·31/32·64KB/25GB/s ≈ 124 µs + 5 µs = 129 µs，已经相当于计算时间的 69%——再增加 E 只会更慢。
  • 这条曲线的形状（先线性上升、后饱和并回落）就是”张量并行度不能无限大”的定量原因；同时也是为什么张量并行必须放在 NVLink 域内：把 25 GB/s 换成 300 GB/s，通信时间降到 ~11 µs，E=32 才重新变得划算。
• 瓶颈总结：小规模时瓶颈是朴素 GEMM 的指令级效率（只有 5.7 GFLOP/s/rank，未用向量化分块，远低于单核峰值）；规模变大后瓶颈切换到集合通信的启动延迟（α 项）；M 很小（此处 64）时每一层的 GEMM 都”又短又频繁”，α 的相对代价被放大——这正是小 micro-batch 与高张量并行度不能同时使用的另一面。

3.2 示例 2：CUDA 上单个张量并行分片的 tiled GEMM

张量并行切完之后，每张 GPU 上的活就是一个 M × K × (N/E) 的 GEMM。下面这个 CUDA 程序实现这张分片 GEMM，并直接暴露”切分粒度 ↔ GPU 占用率”的关系。

// tp_shard_gemm.cu -- 张量并行分片上的 tiled GEMM（本 rank 只算 X * W_e，W_e 为列分片）
// 编译: nvcc -O3 -arch=sm_80 -Xptxas -v tp_shard_gemm.cu -o tp_shard_gemm
// 运行: ./tp_shard_gemm [M] [K] [Nlocal]     例如 ./tp_shard_gemm 1024 1024 512
#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <cuda_runtime.h>

#define TILE 32
#define CK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \
    std::printf("CUDA error %s:%d: %s\n", __FILE__, __LINE__, cudaGetErrorString(e_)); std::exit(1); } } while (0)

// Y[M x N] = X[M x K] * W[K x N]，分块装进 shared memory
__global__ void gemm_tiled(const float* __restrict__ X,
                           const float* __restrict__ W,
                           float* __restrict__ Y,
                           int M, int K, int N)
{
    __shared__ float Xs[TILE][TILE + 1];   // +1 避免 bank conflict
    __shared__ float Ws[TILE][TILE + 1];
    const int row = blockIdx.y * TILE + threadIdx.y;   // 输出行
    const int col = blockIdx.x * TILE + threadIdx.x;   // 输出列（本 rank 的局部列）
    float acc = 0.0f;
    for (int k0 = 0; k0 < K; k0 += TILE) {
        Xs[threadIdx.y][threadIdx.x] = (row < M && k0 + threadIdx.x < K) ? X[row * K + k0 + threadIdx.x] : 0.0f;
        Ws[threadIdx.y][threadIdx.x] = (k0 + threadIdx.y < K && col < N) ? W[(k0 + threadIdx.y) * N + col] : 0.0f;
        __syncthreads();
#pragma unroll
        for (int k = 0; k < TILE; ++k) acc += Xs[threadIdx.y][k] * Ws[k][threadIdx.x];
        __syncthreads();
    }
    if (row < M && col < N) Y[row * N + col] = acc;
}

int main(int argc, char** argv) {
    const int M = argc > 1 ? std::atoi(argv[1]) : 1024;   // token 数 B*S
    const int K = argc > 2 ? std::atoi(argv[2]) : 1024;   // 输入维度
    const int N = argc > 3 ? std::atoi(argv[3]) : 512;    // 本 rank 的分片宽度（= 全局 N / TP 度）
    const int iters = 100;

    float *hX = (float*)std::malloc(sizeof(float) * M * K);
    float *hW = (float*)std::malloc(sizeof(float) * K * N);
    float *hY = (float*)std::malloc(sizeof(float) * M * N);
    float *hRef = (float*)std::malloc(sizeof(float) * M * N);
    for (int i = 0; i < M * K; ++i) hX[i] = std::sin(0.001f * i);
    for (int i = 0; i < K * N; ++i) hW[i] = std::cos(0.002f * i);

    float *dX, *dW, *dY;
    CK(cudaMalloc((void**)&dX, sizeof(float) * M * K));
    CK(cudaMalloc((void**)&dW, sizeof(float) * K * N));
    CK(cudaMalloc((void**)&dY, sizeof(float) * M * N));
    CK(cudaMemcpy(dX, hX, sizeof(float) * M * K, cudaMemcpyHostToDevice));
    CK(cudaMemcpy(dW, hW, sizeof(float) * K * N, cudaMemcpyHostToDevice));

    dim3 block(TILE, TILE);
    dim3 grid((N + TILE - 1) / TILE, (M + TILE - 1) / TILE);
    for (int it = 0; it < 5; ++it)      // warmup
        gemm_tiled<<<grid, block>>>(dX, dW, dY, M, K, N);
    CK(cudaDeviceSynchronize());

    cudaEvent_t e0, e1;
    CK(cudaEventCreate(&e0)); CK(cudaEventCreate(&e1));
    CK(cudaEventRecord(e0));
    for (int it = 0; it < iters; ++it)
        gemm_tiled<<<grid, block>>>(dX, dW, dY, M, K, N);
    CK(cudaEventRecord(e1));
    CK(cudaEventSynchronize(e1));
    float ms = 0.0f;
    CK(cudaEventElapsedTime(&ms, e0, e1));
    ms /= iters;

    CK(cudaMemcpy(hY, dY, sizeof(float) * M * N, cudaMemcpyDeviceToHost));

    // 单线程参考实现（只做一次，O(M*K*N) 在 CPU 上很慢，所以取一个子块验证）
    const int VM = 16, VN = 16;
    double err = 0.0;
    for (int i = 0; i < VM; ++i) {
        for (int j = 0; j < VN; ++j) {
            double acc = 0.0;
            for (int k = 0; k < K; ++k) acc += (double)hX[i * K + k] * (double)hW[k * N + j];
            hRef[i * N + j] = (float)acc;
            err = std::fmax(err, std::fabs((double)hY[i * N + j] - (double)hRef[i * N + j]));
        }
    }

    cudaDeviceProp prop;
    CK(cudaGetDeviceProperties(&prop, 0));
    const double flops = 2.0 * M * K * (double)N;
    std::printf("分片 GEMM: M=%d K=%d Nlocal=%d | grid=(%ux%u)=%u blocks, %d SMs\n",
                M, K, N, grid.x, grid.y, grid.x * grid.y, prop.multiProcessorCount);
    std::printf("时间 = %.3f ms  吞吐 = %.1f GFLOP/s  （%u 个 block / %d SM）\n",
                ms, flops / (ms * 1e-3) / 1e9, grid.x * grid.y, prop.multiProcessorCount);
    const double bytesPerBlock = 2.0 * TILE * K * sizeof(float);   // 每个 block 从 DRAM 读 X 片 + W 片
    const double flopsPerBlock = 2.0 * (double)TILE * TILE * K;
    std::printf("每 block: 全局访存 %.1f KB, 计算 %.1f KFLOP -> 算术强度 %.1f FLOP/B (= TILE/4)\n",
                bytesPerBlock / 1024.0, flopsPerBlock / 1e3, flopsPerBlock / bytesPerBlock);
    std::printf("子块验证 (%dx%d) 最大误差 = %.3e\n", VM, VN, err);
    std::free(hX); std::free(hW); std::free(hY); std::free(hRef);
    cudaFree(dX); cudaFree(dW); cudaFree(dY);
    return 0;
}

【代码做什么？】

1. 每个 32×32 的线程块负责输出 Y 的一个 32×32 分块，线程 (tx,ty) 负责其中一个输出元素 (row, col)。
2. 沿 k 维循环 K/TILE 次：每次把 X 的一块（32×32）与 W 的一块（32×32）协作装载进 shared memory（每个线程搬 1 个元素），__syncthreads() 保证装载完成后才开始算；然后每个线程沿 k 做 32 次 FMA；再 __syncthreads() 保证所有人算完才覆盖共享内存（双缓冲缺失，这里是两次同步的原因，也是本 kernel 未优化的点）。
3. 累加器 acc 保存在寄存器里，整个 k 循环只从 shared memory 读，不回 DRAM——这就是”提高算术强度的分块”。
4. Xs/Ws 声明为 [TILE][TILE+1] 是经典的 padding 技巧：避免 Ws[k][threadIdx.x]（按列访问）产生的 bank conflict。
5. 计时用 cudaEvent，先 warmup 5 次再测 100 次平均；正确性用 CPU 单线程算 16×16 子块对照。

【并行机制与性能解说】

• 并行层次：host 启动一个 grid → 每个 SM 分到若干 block → 每个 block 32×32 = 1024 个线程 = 32 个 warp → 每个 warp 32 个 lane 执行同一条 FMA 指令（SIMT/SIMD 执行）。默认参数 M=1024, K=1024, N=512 时 grid = (512/32) × (1024/32) = 16 × 32 = 512 个 block，即 512 × 1024 = 524288 个线程；A100 有 108 个 SM，若每 SM 驻留 2 个这样的 block（2048 线程上限），一次只能跑 216 个 block → 需要 2.4 个 wave（尾部 wave 只填了 40% 的 SM，tail effect）。
• Work（总工作量）：W = 2·M·K·N = 2·1024·1024·512 = 1.074 GFLOP（单张分片上的量）。
• Span（关键路径）：每个线程串行执行 K/TILE = 32 轮，每轮 1 次 shared memory 装载 + 32 次 FMA，因此 Span ≈ 2K = 2048 FLOP 加上 2·K/TILE = 64 次 __syncthreads() 的同步延迟（每次约几百 ns 量级，是真实瓶颈之一）。
• 并行度 = Work/Span = 2MKN / (2K) = M·N = 524288 个独立输出元素（理论上界）。硬件可同时驻留的线程 = SM 数 × 每 SM 线程数 = 108 × 2048 ≈ 221K，所以并行度够用但余量只有 2.4×——这解释了 TILE 必须调大（64/128）配合寄存器分块（每个线程算 8×8 个输出）来把”每个 block 的复用次数”提上去，而不是靠堆 block。
• 瓶颈（带宽 vs 计算，用 Roofline 判断）：每个 block 从 DRAM 读 2·TILE·K·4 B = 256 KB，做 2·TILE²·K = 2.1 MFLOP，片内算术强度 = TILE/4 = 8 FLOP/byte（与 K、M、N 无关，只由 tile 尺寸决定）。
  • A100 的 FP32 峰值 19.5 TFLOP/s、HBM 带宽 2039 GB/s → 平衡点 = 19.5e12/2.039e12 = 9.6 FLOP/byte。
  • 8 < 9.6 → 该 kernel 落在带宽受限区。预计时间 = 总 DRAM 流量 / 带宽 = (512 block × 256 KB)/2039 GB/s = 134 MB/2039 GB/s ≈ 66 µs，而纯计算时间 = 1.074 GFLOP/19.5 TFLOP/s = 55 µs；取两者较大者 ≈ 66 µs（约 16 TFLOP/s，即 FP32 峰值的 84%）。（此为本环境无 GPU 时的模型预测值，非实测；运行程序会打印实测 GFLOP/s 供对照。）
  • 结论：TILE=32 的 FP32 分块 GEMM 刚好卡在 A100 的 roofline 拐点附近；要突破就必须（1）加大 TILE + 寄存器分块把 AI 提到 20+，（2）换 TF32/FP16 tensor core 把平衡点推到 100+ FLOP/byte。
• 张量并行切分带来的额外损耗（本示例的核心教学点）：N 被切成 N/E 后，grid 的 x 维变成 N/(E·TILE)。把 micro-batch 变小（B·S 从 1024 降到 64）时：
  • M/32 从 32 降到 2 → grid 变成 16 × 2 = 32 个 block，只有 32/108 = 30% 的 SM 有活干，设备利用率直接掉到 30%；
  • 同时每层的 GEMM 耗时缩短到几微秒量级，而张量并行的 AllReduce 延迟（α 项）不随 batch 变小而变小 → 通信/计算比急剧恶化。
  • 这两条正是讲义第 31 页那句 small micro-batch sizes reduce GPU utilization 的硬件级解释，也是”张量并行 + 流水线并行”必须小心选 micro-batch 尺寸的原因。

3.3 示例 3：用真实线程模拟 GPipe 流水线，实测气泡比例并与公式对照

这段代码用 std::thread + 条件变量邮箱（mailbox）模拟 p 级流水线、m 个 micro-batch 的 GPipe 调度（全前向 → 全反向），统计每级的忙碌时间，从而实测气泡比例，并与讲义的 (p-1)/m 口径和墙钟时间口径 (p-1)/(m+p-1) 对照。

// gpipe_sim.cpp -- 用真实线程模拟 GPipe 的 micro-batch 流水线，实测气泡比例
// 编译: g++ -O3 -std=c++17 -pthread gpipe_sim.cpp -o gpipe_sim
// 运行: ./gpipe_sim 8 4 2      # m=8 个 micro-batch, p=4 级, 前向 tf=2 ms (反向 tb=tf)
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <memory>
#include <mutex>
#include <queue>
#include <thread>
#include <vector>

using clk = std::chrono::steady_clock;

struct Mailbox {                       // 单生产者/单消费者邮箱（流水线级间的“激活/梯度”通道）
    std::mutex m;
    std::condition_variable cv;
    std::queue<int> q;
    bool closed = false;
    void push(int v) { { std::lock_guard<std::mutex> lk(m); q.push(v); } cv.notify_one(); }
    bool pop(int& v) {
        std::unique_lock<std::mutex> lk(m);
        cv.wait(lk, [&] { return !q.empty() || closed; });
        if (q.empty()) return false;
        v = q.front(); q.pop();
        return true;
    }
    void close() { { std::lock_guard<std::mutex> lk(m); closed = true; } cv.notify_all(); }
};

int main(int argc, char** argv) {
    const int m  = argc > 1 ? std::atoi(argv[1]) : 8;   // 一个 mini-batch 切成的 micro-batch 数
    const int p  = argc > 2 ? std::atoi(argv[2]) : 4;   // 流水线级数（= 设备数）
    const int tf = argc > 3 ? std::atoi(argv[3]) : 2;   // 单级前向耗时 (ms)
    const int tb = tf;                                  // 单级反向耗时 (ms)

    std::vector<std::unique_ptr<Mailbox>> fwd(p), bwd(p);
    for (int s = 0; s < p; ++s) { fwd[s] = std::make_unique<Mailbox>(); bwd[s] = std::make_unique<Mailbox>(); }
    std::vector<double> busy(p, 0.0), wall(p, 0.0);
    std::atomic<int> done_fwd{0};

    auto stage = [&](int s) {
        auto t_begin = clk::now();
        double local_busy = 0.0;
        // ---- 前向：从上一级收 micro-batch，算完转发给下一级 ----
        for (int i = 0; i < m; ++i) {
            int id;
            if (!fwd[s]->pop(id)) break;
            auto b0 = clk::now();
            std::this_thread::sleep_for(std::chrono::milliseconds(tf));   // 模拟本级前向计算
            local_busy += std::chrono::duration<double, std::milli>(clk::now() - b0).count();
            if (s + 1 < p) fwd[s + 1]->push(id); else ++done_fwd;
        }
        if (s + 1 < p) fwd[s + 1]->close();
        // ---- 反向：GPipe 调度下，反向任务在全前向完成后由末端按逆序流入 ----
        for (int i = 0; i < m; ++i) {
            int id;
            if (!bwd[s]->pop(id)) break;
            auto b0 = clk::now();
            std::this_thread::sleep_for(std::chrono::milliseconds(tb));   // 模拟本级反向计算
            local_busy += std::chrono::duration<double, std::milli>(clk::now() - b0).count();
            if (s > 0) bwd[s - 1]->push(id);
        }
        if (s > 0) bwd[s - 1]->close();
        busy[s] = local_busy;
        wall[s] = std::chrono::duration<double, std::milli>(clk::now() - t_begin).count();
    };

    auto t_all = clk::now();
    std::vector<std::thread> th;
    for (int s = 0; s < p; ++s) th.emplace_back(stage, s);
    for (int i = 0; i < m; ++i) fwd[0]->push(i);      // 主线程投喂 micro-batch 0..m-1
    fwd[0]->close();
    while (done_fwd.load() < p) std::this_thread::yield();  // 等所有级的前向都做完
    for (int i = m - 1; i >= 0; --i) bwd[p - 1]->push(i);   // 反向按逆序投喂（GPipe：全前向→全反向）
    bwd[p - 1]->close();
    for (auto& t : th) t.join();
    const double total = std::chrono::duration<double, std::milli>(clk::now() - t_all).count();

    double busy_max = 0, busy_sum = 0;
    for (int s = 0; s < p; ++s) { busy_max = std::fmax(busy_max, busy[s]); busy_sum += busy[s]; }
    std::printf("m=%d p=%d tf=tb=%dms\n", m, p, tf);
    std::printf("实测总时间      = %.1f ms   (理论 (m+p-1)*(tf+tb) = %.1f ms)\n", total, (double)(m + p - 1) * (tf + tb));
    std::printf("单级忙碌时间    = %.1f ms   (理论 m*(tf+tb) = %.1f ms)\n", busy_max, (double)m * (tf + tb));
    std::printf("实测空闲(气泡)  = %.1f%% 占墙钟时间\n", 100.0 * (total - busy_max) / total);
    std::printf("讲义口径 气泡/理想计算 = (p-1)/m = %.1f%%\n", 100.0 * (p - 1) / m);
    std::printf("时间占比口径 (p-1)/(m+p-1) = %.1f%% ; 平均设备利用率 = %.1f%%\n",
                100.0 * (p - 1) / (m + p - 1), 100.0 * busy_sum / (p * total));
    return 0;
}

【代码做什么？】

1. 为每一级 s 建两个 Mailbox：fwd[s] 收”上游传来的激活”，bwd[s] 收”下游传来的梯度”。
2. 每级一个线程，先后做两段循环：先做 m 次前向（sleep t_f 模拟计算，然后把 micro-batch id 传给 fwd[s+1]），再做 m 次反向（sleep t_b，把 id 传给 bwd[s-1]）。
3. 主线程投喂 micro-batch 0..m-1 给第 0 级；等所有级的前向计数到齐（done_fwd == p）后，按逆序 m-1..0 投喂反向给最后一级——这正是 GPipe 的”全前向 → 全反向”调度。
4. 结束时统计每级忙碌时间与墙钟时间，打印实测气泡比例并与两个解析口径对照。

实测输出（本机实跑，tf=tb=2 ms）：

m=8  p=4  tf=tb=2ms   实测总时间 46.0 ms (理论 44.0)  单级忙碌 33.2 ms (理论 32.0)
                      实测气泡 27.9% | 讲义口径 (p-1)/m = 37.5% | 时间占比 (p-1)/(m+p-1) = 27.3%
m=16 p=8  tf=tb=2ms   实测总时间 96.1 ms (理论 92.0)  单级忙碌 66.3 ms
                      实测气泡 31.0% | 讲义口径 43.8% | 时间占比 30.4%
m=64 p=16 tf=tb=2ms   实测总时间 328.9 ms (理论 316.0) 单级忙碌 264.6 ms
                      实测气泡 19.5% | 讲义口径 23.4% | 时间占比 19.0%

实测气泡比例与 (p-1)/(m+p-1) 几乎完全一致（27.9% vs 27.3%、31.0% vs 30.4%、19.5% vs 19.0%），而相对”理想计算时间”的口径 (p-1)/m 明显偏大——这就是第 2.4 节强调的”分母不同”问题。总时间比理论值高 ~4%，来自线程调度与 mutex/condvar 的开销（模拟器的开销，不是流水线本身的开销）。

【并行机制与性能解说】

• 线程如何创建、工作如何分配：p 个 std::thread（对应 p 张 GPU / p 个 MPI rank），每级只处理”属于自己的 micro-batch 序列”，级间用 Mailbox（mutex + condvar）传递数据，等价于真实系统中的点对点 send/recv 或 NCCL Send/Recv。共享数据：busy[]、wall[] 每个线程写自己的下标——但它们相邻存储，存在伪共享（false sharing）风险（8 个 double 共 64 字节，正好一条 cache line）！这里因为只在最后写一次、且不在热点路径上，影响可以忽略；但如果是”每处理一个 micro-batch 就累加一次计数器”，就必须给每个计数器加 padding 或用 std::atomic 分散到不同 cache line——这正是 15-418 反复强调的伪共享陷阱。
• Work（总工作量）：W = m·p·(t_f+t_b)（所有级所有 micro-batch 的计算时间之和，单位”级·毫秒”）= 8·4·4 = 128 级·ms。
• Span（关键路径）：S = (m+p-1)(t_f+t_b) = 11·4 = 44 ms（fill + 稳定 + drain，见第 2.4 节的甘特图）。
• 并行度 = Work/Span = m·p/(m+p-1) = 32/11 = 2.91（对 p=4 的理想上界 4 而言 = 72.7% 的利用率上限），与 1-(p-1)/(m+p-1) = 8/11 = 72.7% 完全一致，也与实测平均利用率 71.9% 吻合。
• 可扩展性上限：lim(m→∞) m·p/(m+p-1) = p——只有把 m 做得足够大，利用率才趋近 100%；而 lim(m 固定, p→∞) = m——流水线越深，需要越多的 micro-batch 才能填满。两个方向的代价分别是”micro-batch 太小 → GPU 打不满（见 3.2 节）”与”级数太多 → 每级只有极少层、通信占比上升”。
• 瓶颈：这个模拟器的瓶颈是主线程的串行投喂与 yield 自旋，以及 condvar 唤醒的调度延迟（这就是为什么实测比理论慢 4%）。在真实系统里，对应的瓶颈是点对点通信的启动延迟与带宽（第 4.5 节给了数值），以及”所有级必须等最慢的级”——负载不均（stage 划分不等长）会让真实气泡远大于公式值，因为公式假设所有级的 t_f/t_b 相同。

4. 性能模型与复杂度分析

4.1 三个代码示例的 Work / Span / 并行度汇总

4.2 数值算例 A：张量并行的通信-计算比（为什么 TP 必须留在 NVLink 域内）

假设（A100-80GB 单节点 8 卡，NVLink 单向 300 GB/s，BF16 训练）： B·S = 8192 个 token（micro-batch 8 × 序列 1024）、隐层 H = 4096、TP 度 E = 8、每 GPU 有效算力 150 TFLOP/s（A100 BF16 tensor core 峰值 312 TFLOP/s 的 48%，是千亿模型上比较现实的 MFU）。

计算量（每个 Transformer 层，每卡）：

• 注意力投影（QKV + 输出投影）：4 × 2·B·S·H² = 8·8192·4096² = 1.10e12 FLOP
• MLP（H→4H→H）：2 × 2·B·S·H·4H = 16·8192·4096² = 2.20e12 FLOP
• 层合计 3.30e12 FLOP → 每卡 3.30e12/8 = 412 GFLOP → 计算时间 = 412e9/150e12 = 2.75 ms

通信量（每层 2 次 AllReduce：注意力输出后、MLP 输出后）：

• 单次 AllReduce 消息 = B·S·H 个元素 × 2 B = 8192×4096×2 = 67.1 MB
• ring AllReduce 每卡收发量 = 2(E-1)/E × 67.1 MB = 117 MB；两次共 235 MB
• 传输时间 = 235 MB / 300 GB/s = 0.78 ms；再加 2 × 2(E-1) = 28 步的启动延迟（NVLink 域内 ~3 µs）= 0.084 ms → 通信 ≈ 0.87 ms

结论：通信/计算 = 0.87/2.75 = 32%，即约 24% 的层时间花在 AllReduce 上（0.87/(2.75+0.87)）。这还是在 300 GB/s 的 NVLink 上。如果把同一份张量并行放到跨节点（IB HDR 25 GB/s）：传输时间变成 235 MB/25 GB/s = 9.4 ms，是计算时间的 3.4 倍——层时间从 3.6 ms 劣化到 12.2 ms，整个训练慢 3.4 倍。这就是”张量并行绝不出节点”的定量依据（12 倍带宽差 → 3.4 倍端到端劣化）。

更重要的结构性结论：把两个比例写成公式， 通信时间/计算时间 = [4(E-1)/E·B·S·H·2/BW] / [24·B·S·H²/(E·P)] = 8(E-1)·P/(24·H·BW) ——它与 batch（B·S）完全无关，只与并行度 E、隐层 H、以及”算力/带宽比”有关。所以”把 micro-batch 调大”不能降低张量并行的相对通信开销（它只能改善 GEMM 的硬件效率）；要降低它只能：增大 H（用更大的模型换效率）、降低 E、或换更快的链路。这是一个非常容易被误解的性能事实。

4.3 数值算例 B：流水线气泡与配置选择

设 t_f = t_b = 1（归一化时间单位），两级口径的数值对照：

由表得到的实用规则：

• 想让气泡（墙钟口径）低于 10%，需要 m ≥ 9(p-1)。p=8 → m ≥ 63（即 64 个 micro-batch）；p=16 → m ≥ 135（128 还不够，要 144）。
• 结合 3.2 节的结论：micro-batch 太小会让 GPU 打不满，所以不能靠”无限切细”来增大 m，只能靠增大 mini-batch 总量——而大 batch 可能损害收敛精度（讲义第 31 页的 caveat）。这就是流水线并行的核心两难：气泡要大 m，大 m 要大 batch 或小 micro-batch，两者都有代价。
• 用 GPipe 的调度时，激活显存 ∝ m（必须保存所有 micro-batch 的中间激活直到反向）。m=64、每个 micro-batch 的激活约 100 MB（B_micro=8, S=1024, H=4096 量级）时，光激活就要 6.4 GB/卡；1F1B（前向反向交错）把这部分降到 ≈ p 个 micro-batch 的激活（~1 GB），这是工业界普遍不直接用 GPipe 而用 1F1B 的原因。

4.4 数值算例 C：Roofline —— micro-batch 变小如何把 GEMM 推进带宽受限区

考虑 MLP 的第一个线性层（列并行分片，K = H = 4096、N = 4H/E = 2048，BF16），比较两个 micro-batch 规模：

(a) 小 micro-batch：B·S = M = 64

• FLOP = 2·M·K·N = 2·64·4096·2048 = 1.074e9
• 强制 DRAM 流量（每个权重只读一次）：权重 2·4096·2048 = 16.8 MB + 输入 2·64·4096 = 0.52 MB + 输出 2·64·2048 = 0.26 MB ≈ 17.6 MB
• 算术强度 = 1.074e9/17.6e6 = 61 FLOP/byte < A100 平衡点 153 → 带宽受限
• 时间 = 17.6 MB/2039 GB/s = 8.6 µs → 有效算力 = 1.074e9/8.6e-6 = 125 TFLOP/s（峰值的 40%）

(b) 大 micro-batch：B·S = M = 8192

• FLOP = 2·8192·4096·2048 = 1.374e11
• DRAM 流量 = 权重 16.8 MB + 输入 67.1 MB + 输出 33.6 MB = 117 MB
• 算术强度 = 1.374e11/1.17e8 = 1174 FLOP/byte ≫ 153 → 计算受限
• 时间 = 1.374e11/312e12 = 0.44 ms（峰值）或 0.92 ms（150 TFLOP/s 有效）

结论：micro-batch 缩小 128 倍后，有效算力从计算受限时的高位掉到带宽受限的 40% 峰值（算术强度掉到平衡点以下）。这就是”流水线并行需要用大 m 来消气泡，但每个 micro-batch 又不能太小”的定量冲突；也是”张量并行要配合大 batch”的根本原因。工程上的缓解手段：算子融合（CS149：算术强度 1/3 → 3/5，flash-attention 把 5MN+2M 次读降到 MN）、低精度（流量按位宽线性下降）、以及把权重常驻片上（CUDA graph / persistent kernel）。

更完整的口径：情形 (a) 落在带宽受限区，天花板 = 2039 GB/s × 61 FLOP/B = 125 TFLOP/s（峰值的 40%）；情形 (b) 落在计算受限区，能跑到 GEMM 实现质量所允许的任何水平（本笔记用 150 TFLOP/s = 48% 作为保守值；成熟实现的大 GEMM 可到 275 TFLOP/s ≈ 88% 峰值）。关键不是绝对数字，而是”算术强度 61 vs 1174 跨越了 153 这个平衡点”——一旦跨过去，硬件再快也救不了。

补充（CS149 的同类数字，供对照）：V100 有 6 MB L2、16 GB HBM、900 GB/s 带宽、80 个 SM；一个 N=32, P=Q=256 的卷积层输出就有 256×256×128×32 = 256M 个输出 = 1 GB 的 float32 输出数据——单层激活就逼近整个显存容量，这直接说明为什么”激活”必须被切分（数据并行切 batch、张量并行切通道、流水线并行切层）以及为什么融合能省下成百 MB 级的 DRAM 流量。

4.5 数值算例 D：带宽/延迟（α-β）模型与”通信频率决定链路”

设一次跨设备传输 n 字节的时间为 T = α + n/BW。统一假设：模型 96 层、H = 4096（每层参数 ≈ 12H²：注意力 4 个 H×H 投影 + MLP 的 8H²，全模型 ≈ 19B 参数）、BF16、总规模 64 卡（TP8 × PP8）外加 DP = 8 副本、单个 micro-batch = 8192 个 token、每卡有效算力 150 TFLOP/s。一个 iteration 内三类通信的特征如下：

其中该 iteration 的计算时间 = 24·B·S·H²·L / (E·P) = 24·8192·4096²·96/(8·150e12) = 264 ms（每卡），与第 4.2 节的每层记账一致（24·B·S·H²/E = 412 GFLOP/层/卡 ÷ 150 TFLOP/s = 2.75 ms/层 × 96 层 = 264 ms）。

读表要点：

1. 张量并行的通信总量最大（22.5 GB/iter/卡）且消息最”碎”（192 次）、对 α（延迟）最敏感 → 必须 NVLink。注意它的总量之所以吓人，是因为它每层都发生；单看一条消息（67 MB）其实很小。若把同一批 token 切成 m 个 micro-batch，TP 的单条消息会缩小到 1/m、次数变成 m 倍，总量不变但 α 项比重上升——这也说明 micro-batch 尺寸与 α 敏感度是耦合的。
2. 数据并行的总量看似也大（1.05 GB/卡），但一个 iteration 只发生一次，且 ring AllReduce 可以拆成 ReduceScatter + AllGather 与反向计算重叠（前一讲结论：ring AllReduce 每卡只发 2(E-1)/E·M）→ 放在最慢的链路上也往往够用。
3. 流水线并行的通信量与 micro-batch 数 m 成正比（每级边界每 micro-batch 一次）：这里 m=1 时只有 0.47 GB；若把同一个 8192-token 的 batch 拆成 m=8 个 1024-token 的 micro-batch，量仍是 0.47 GB（因为每块变小、块数变多），但消息数变成 56 条、每条 8.4 MB，α 的相对代价上升——这是”micro-batch 不宜太小”的又一个理由。
4. Amdahl 视角：假设 PP 与 DP 的通信能被计算完全掩盖，而张量并行的 75 ms AllReduce 串在每层的关键路径上无法掩盖，则一个 iteration 里 75/(264+75) = 22% 的时间是”无法被计算掩盖”的通信开销。这 22% 就是 Amdahl 意义上的串行段：即使把计算部分加速到无穷，也只是让这 339 ms 变成 75 ms（最多再快 4.5 倍）；等价地说，该 iteration 的有效算力利用率被锁在 264/339 ≈ 78%。因此工业实现必须做通信/计算重叠（把 AllReduce 拆成 ReduceScatter/AllGather 与相邻 GEMM 并行、或使用异步流水线与 1F1B），否则任何并行策略都会撞上这堵 Amdahl 墙。

4.6 复杂度总览

5. 关键要点

1. 并行策略的选择本质上是”什么资源先耗尽”的问题：数据并行解决不了显存墙（每卡全量副本），张量并行解决显存与单层计算量，流水线并行把模型并行的空转填回吞吐。讲义第 33 页的收口结论是：训练超大模型必须叠加数据/模型/流水线等多种并行技术（DeepSpeed 3D 并行 = DP × PP × TP）。
2. 通信量由”切什么”决定，通信频率由”切在哪一层”决定：数据并行通信量 = 参数尺寸（C_out·C_in，每 iter 一次）；张量并行通信量 = 激活尺寸（B·C_in 或 B·C_out，每层 2 次），量级小 ≈ C_out/B 倍（C_in ≈ C_out 时即 C_in/B，取 C_in = 4096, B = 8 就是 512 倍），但频率高 层数 × 2 倍。频率决定可用链路：最高频的 TP 必须留在 NVLink 域内（300 GB/s vs IB 25 GB/s，代价是 3.4 倍端到端劣化）。
3. 张量并行的通信-计算比与 batch 无关，只与 (E-1)·P/(H·BW) 有关：调大 micro-batch 无法降低 TP 的相对开销，只能改善 GEMM 的硬件效率（Roofline）。要降低 TP 开销只有三条路：增大 H、降低 E、或换更快的链路。这条结论是”大模型训练用大隐层 + 小 TP 度 + 大 batch”的定量依据。
4. 流水线的气泡 (p-1)/(m+p-1) 是典型的 Amdahl 串行段：fill 与 drain 无法并行。要把它压到 10% 以下需要 m > 9(p-1)（p=4 → m ≥ 28；p=8 → m ≥ 64），而 m 的增大要么靠更大的 mini-batch（可能损害收敛），要么靠更小的 micro-batch（把 GPU 推进带宽受限区，算术强度从 1174 掉到 61 FLOP/byte，有效算力掉到峰值的 40%）。必须同时看这两个约束才能选对 m。
5. 切分粒度受硬件并行度的硬约束：TP 把 N 切成 N/E 后，CUDA grid 的 block 数按 1/E² 量级下降——M=64, N=512 的分片只有 32 个 block，而 A100 有 108 个 SM，70% 的 SM 闲置。所以每一层并行的最小切分尺寸要保证”分片后每张卡仍有足够多的 tile 填满 SM”，这也是 TILE 要调大、要用寄存器分块与 tensor core 的原因。

6. 常见陷阱与注意事项

• 把模型并行当成”加速手段”：模型并行（无 micro-batch）在时间轴上只有一台设备在算（讲义第 28 页：Under-utilization of compute resources / Low overall throughput），加设备只会加长等待。它是容量手段；吞吐要靠流水线并行或数据并行来补。评估任何”模型并行方案”时，先问”稳态时每张卡在干什么”。
• GeLU/softmax 等非线性之前做归约（数学错误 + 性能错误）：GeLU(X·A₁) + GeLU(X·A₂) ≠ GeLU(X·A)。所以列并行（partition output）之后不能立刻 AllReduce 成完整输出，必须让 Y = [Y₁\|Y₂] 保持拼接形态、把通信推迟到第二层 GEMM 之后。把 AllReduce 放在 4H 宽的中间激活上而不是 H 宽的层输出上，通信量会大 4 倍——这就是 Megatron 选择”列并行在前、行并行在后”的两个理由（见 Q2）。
• 只用”平均 batch”估算显存，忽略激活与优化器状态：显存 = 参数 + 梯度 + 优化器状态（FP32 的 master weights + Adam 的 m、v 共 12 B/参数）+ 激活（∝ m、∝ 序列长）。175 B 模型即使用 TP 8 × PP 8（64 卡）切分，每卡仍有 2.73 B 参数 → FP16 权重 5.5 GB + FP16 梯度 5.5 GB + 优化器状态 32.8 GB ≈ 43.8 GB，距 80 GB 只剩一半给激活与临时缓冲；如果这一层外面再套数据并行，优化器状态还会在 DP 组内被复制（这正是 ZeRO 要消掉的冗余）。
• 忽略通信的”频率 × 延迟”而只算总字节数：数据并行每 iter 传 1 GB 量级的梯度却几乎不伤性能（一次大消息、只发生一次、可重叠）；张量并行每层传 67 MB 却可能吃掉 24% 的层时间（96 层 × 2 次小消息、串在关键路径上）。只看总字节会把两者的代价完全看反；必须同时算 α 项（每次 AllReduce 有 2(E-1) 步启动延迟）与重叠可能性。
• 伪共享（false sharing）与同步粒度错误：流水线/张量并行的 worker 常常维护”每个 micro-batch 一个计数器/累加器”，若这些变量相邻存放（例如 std::vector<double> busy(p)，8 个 double 正好占满一条 64 B cache line），每个线程的写都会把其他核的 cache line 置为 Invalid，性能崩溃。解决：padding 到 64 B、或按 cache line 对齐的 per-thread 副本最后再归约；同时避免在最内层循环里用粗粒度锁（应像 3.3 节的 mailbox 一样只在级边界同步一次）。
• 把理论气泡当成”能消除的浪费”：(p-1)/(m+p-1) 是调度的固有下界；更常见的问题是级间负载不均（各层计算量不同，划分不等长），此时真实气泡会显著大于公式值——公式假设所有级的 t_f/t_b 完全相同。工程上要么按 FLOP 均衡划分（让每级层数不同），要么用 interleaved/virtual stages 把粒度打碎；此外 p 越大，单样本端到端延迟 ≈ p·(t_f+t_b) 越大，训练吞吐与推理延迟的取舍需要分开考虑。

7. 思考题（带答案）

Q1（定量）：某模型用 p = 4 级流水线、每个 mini-batch 切 m = 8 个 micro-batch，每级 t_f = t_b。（a）按讲义口径算气泡比例；（b）算气泡占墙钟时间的比例与设备利用率；（c）如果保持 p = 4、希望墙钟气泡低于 10%，m 至少要多大？（d）如果把流水线加深到 p = 8 并保持墙钟气泡 10% 以下，m 又要多大？请指出这里的实际矛盾。

【答案】 (a) 讲义第 30 页口径：BubbleFraction = (p-1)(t_f+t_b)/(m(t_f+t_b)) = (p-1)/m = 3/8 = 37.5%（分母是”理想计算时间” m(t_f+t_b)）。 (b) 墙钟总时间 = (m+p-1)(t_f+t_b) = 11(t_f+t_b)；气泡时间 = (p-1)(t_f+t_b) = 3(t_f+t_b) → 气泡占墙钟 = 3/11 = 27.3%，设备利用率 = m/(m+p-1) = 8/11 = 72.7%（与 3.3 节的实测 27.9% / 71.9% 吻合）。 (c) 要求 (p-1)/(m+p-1) < 0.1 → 3 < 0.1m + 0.3 → m > 27 → m ≥ 28。 (d) p = 8 时 7/(m+7) < 0.1 → m > 63 → m ≥ 64（正好是 2 的幂，工程上取 64）。 矛盾在于：m 翻倍后每张卡上每个 micro-batch 的 GEMM 的 M 维不变（m 是”批次个数”，不是 micro-batch 尺寸），要增大 m 就必须增大 mini-batch 总量或缩小 micro-batch 尺寸。前者可能损害收敛精度（讲义第 31 页 caveat：large mini-batch sizes can lead to accuracy loss），后者会让每层 GEMM 变小、把 GPU 推入带宽受限区（第 4.4 节：算术强度从 1174 掉到 61 FLOP/byte，有效算力掉到峰值的 40%；第 3.2 节：micro-batch 变小后 grid 只剩 32 个 block，只占 108 个 SM 的 30%）。另外 GPipe 的激活显存 ∝ m，p=8、m=64 时激活显存约为 p=4、m=32 时的 2 倍——所以工业界改用 1F1B 交错调度把激活降到 ≈ p 个 micro-batch。

Q2（理解 + 设计）：Megatron-LM 在 Transformer 的 MLP 里把第一个全连接层的权重按列切（partition output，f 算子），第二个全连接层按行切（reduce output，g 算子），中间夹一个 GeLU。讲义第 21 页特意问了一句 “Why switch?”。请回答：为什么必须是这个顺序？能否反过来（第一层按行切、第二层按列切）？如果用 H 表示隐层、4H 表示中间维度、B·S 表示 token 数，两种顺序的通信量各是多少？多头注意力的头切分属于哪一种？

【答案】 必须是”列并行 → 行并行”，两个理由：

1. 数学上：GeLU 是非线性逐元素函数，GeLU(Y₁) + GeLU(Y₂) ≠ GeLU(Y₁+Y₂)。反过来的话，第一层（行并行）算出的只是部分和，在送入 GeLU 之前就必须先 AllReduce 成完整的 Y——即在中间激活上通信，通信量为 B·S·4H 个元素。
2. 通信量上：按 Megatron 的顺序，中间激活 Y = [Y₁\|Y₂] 保持拼接形态不需要通信（列并行是 partition output），通信被推迟到第二个 GEMM 之后，作用在层输出 Z 上，通信量 = B·S·H 个元素（讲义第 13 页小结表：reduce output 的通信量是 B·C_out）。两者之比 = 4H/H = 4 倍。反过来做要付 4 倍的通信。 附带的好处还有：Z = Z₁+Z₂ 的 AllReduce 结果正好是下一层的输入，可以与下一层的计算/通信重叠；而前向的 AllReduce 在反向时对应 identity（讲义第 10–12 页的两张通信量表：partition output 在反向通信、reduce output 在前向通信），f 与 g 在前后向的角色是对偶的，这让实现可以复用同一套通信调度。 多头注意力：按 head 切分输入投影，每个 head 独立算 Yᵢ = softmax(QᵢKᵢᵀ/√d)Vᵢ，最后 Z = Concat(Y₀,…,Y_h)W_o——W_o 按行切、输出需要 AllReduce，所以属于 reduce output（行并行）（讲义第 25 页标注：Parallelizing across attention heads + Tensor model parallelism (reduce output)）。多层堆叠时，W_o 的行并行输出经 AllReduce 变为完整 Z，与下一层的列并行（partition output，前向无需通信）相接——整个 Transformer 就是”列并行 ↔ 行并行”交替的链条，把 AllReduce 放在每层的输出端（尺寸 H）而不是中间端（尺寸 4H）。

Q3（系统设计）：你要在 512 张 80 GB A100 上训练一个 175 B 参数的 BF16 Transformer（96 层、H = 12288），要求 (1) 模型状态装得下、(2) 单个 iteration 的流水线气泡不超过 ~15%、(3) 张量并行不跨节点（每节点 8 卡 NVLink）。请给出 TP / PP / DP 的配置，说明理由，并算出每张卡上的模型状态显存与全局 batch size。

【答案】 配置：TP = 8，PP = 8，DP = 8（512 = 8 × 8 × 8）。

• TP = 8 且必须放在同一节点内：TP 每层要 2 次 AllReduce，消息大小 B·S·H，只有 NVLink（300 GB/s）能承受（第 4.2 节：换到 IB 25 GB/s 会让层时间劣化 3.4 倍）。8 卡正好吃掉一个节点的 NVLink 域，所以 TP = 8 是上限而不是巧合。
• PP = 8：H = 12288 很大，TP 的通信-计算比 8(E-1)P/(24·H·BW) 随 H 增大而下降，所以 TP 不需要更宽；把层数（深）交给流水线：96 层 / 8 级 = 每级 12 层，级数少意味着级间通信（每 micro-batch 传 B_micro·S·H·2 字节）频率低、可重叠。
• DP = 8：剩下的并行维度，用数据并行提高吞吐；它的梯度 AllReduce 每 iter 只一次、量最大但可重叠，可以跨节点走 IB。
• 每卡模型状态：参数 = 175e9/(TP×PP) = 175e9/64 = 2.73e9。BF16 权重 = 2.73e9 × 2 B = 5.5 GB；BF16 梯度 = 5.5 GB；FP32 master weights + Adam 的 m、v（共 12 B/参数）= 32.8 GB；合计 ≈ 43.8 GB < 80 GB，剩下的 ~36 GB 给激活、通信缓冲与临时张量。（这正是 ZeRO 在数据并行维度上继续切分优化器状态的动机：这 32.8 GB 在 DP 组的 8 个副本里是重复的。）
• 气泡：取 m = 32 个 micro-batch → 墙钟气泡 (p-1)/(m+p-1) = 7/39 = 17.9%；取 m = 48 → 7/55 = 12.7%（满足 ≤15%）。工程上 m 通常取成流水线级数 p 的倍数以便均衡调度（如 4 的倍数）；实际训练多用 1F1B 调度而非 GPipe 以省激活显存。
• 全局 batch：global_batch = micro_batch_size × m × DP = 1 × 48 × 8 = 384 个序列（若 S = 2048 token，则约 78.6 万 token/step）。这里 micro-batch = 1 是为了在 PP = 8、TP = 8 下把激活显存压到最低——代价是每卡 GEMM 的 M 维只有 B·S = 2048，必须靠 tensor core + 大 TILE 才能维持效率（第 4.4 节的 Roofline 分析）。
• 一句话检验：TP×PP×DP = 512 决定了容量上界（175B/64 = 2.73 B 参数/卡）与吞吐上界（(p-1)/(m+p-1) 的气泡 + TP 每层通信）；三层并行各自服务于一个约束，缺一不可——这正是讲义最后一页 “3D parallelism in DeepSpeed” 的含义。