CS336 2026 · Lecture 8已发布

Parallelism

日期: Wed April 22
讲师: Tatsu
课程资料: 9 项
内容产物: 4 项

课程简介

本讲系统比较大模型训练的组合并行策略：从 ZeRO/FSDP 分片参数、梯度和优化器状态，到张量并行、流水线并行、序列并行与上下文并行。重点是根据模型规模、批量大小、节点拓扑和显存预算计算通信代价，并理解为何生产训练通常需要多种并行方式组成 3D 或 4D 并行。

讲义 / Slideslecture_08.pdf打开链接↗论文ZeRO - Memory Optimizations Toward Training Trillion Parameter Models参数分片↗论文Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM3D 并行↗论文Reducing Activation Recomputation in Large Transformer Models序列并行↗论文GPipe - Efficient Training of Giant Neural Networks using Pipeline Parallelism流水线并行↗论文PyTorch FSDP - Experiences on Scaling Fully Sharded Data Parallel参数分片实践↗论文DeepSpeed-Megatron - System Optimizations Enable Training Deep Learning Models with Over 100 Billion Parameters系统案例↗文档Megatron Core Developer Guide工程文档↗视频课程视频课程视频↗

CS336 Lecture 8：并行化（Parallelism）——中文讲义

学习目标

完成本讲后，你应该能够：

理解数据并行（Data Parallelism）与模型并行（Model Parallelism）的根本区别，以及各自的适用场景。
掌握 ZeRO 优化器三个阶段（尤其是 FSDP/ZeRO-3）的工作原理，能解释其“免费”内存节省的通信代价。
理解流水线并行（Pipeline Parallelism）的 Bubble 问题及其缓解策略。
理解张量并行（Tensor Parallelism）的前向/反向传播对偶性，以及它与序列并行（Sequence Parallelism）的组合方式。
了解专家并行（Expert Parallelism）的特殊约束，以及它与张量并行的取舍。
能够通过简单的内存核算（参数 + 优化器状态 + 激活值）判断一种并行策略是否合理。
掌握组合多种并行策略（“3D/4D 并行”）的基本优先级原则。

课程主线

课程首先回顾集合通信原语，然后引入并行化的两大动机——算力扩展与内存扩展。随后依次介绍数据并行（含 ZeRO/FSDP）、流水线并行、张量并行、专家并行与上下文并行，最后给出组合策略的实践指南，并以近期大规模训练案例收尾。

一、并行化的动机与硬件背景

1.1 两个核心瓶颈

并行化要解决两个问题：

算力不足：单个 GPU 无法提供训练大模型所需算力。
内存不足：模型参数、优化器状态、梯度与激活值无法放入单卡显存。

1.2 节点内 vs. 节点间

节点内（intranode）：GPU 间通过高速互联（如 NVLink）连接，通信开销低，可承受频繁通信的并行策略。
节点间（internode）：通过网络（如 InfiniBand）连接，速度较慢，应选择通信量小、模式规则的并行策略。

1.3 TPU 与 GPU 的硬件哲学差异

TPU：采用环形网格（toroidal mesh）拓扑，每个芯片只与邻居通信，邻居数量恒定，可无限扩展且成本低。适合通信模式固定、规则的工作负载（如稠密模型的张量并行）。
GPU：采用胖树（fat-tree）拓扑，支持全对全通信，灵活性高。适合通信模式不确定的工作负载（如 MoE 的 token 路由）。
趋同趋势：现代工作负载（如 MoE）越来越需要全对全能力，新的 TPU 设计也在向更接近 GPU 的交换网络靠拢。

结论：硬件的通信拓扑直接影响并行策略的选择，不同的工作负载应匹配不同的硬件。

二、集合通信原语回顾

2.1 关键原语

all-reduce：所有 GPU 的梯度求和后广播给所有人。
reduce-scatter：将各 GPU 的梯度按索引分块求和，每个 GPU 只持有自己负责的那一块。
all-gather：收集所有 GPU 持有的分块数据，拼接成完整向量。

2.2 关键恒等式（本讲反复使用）

all-reduce = reduce-scatter + all-gather

该恒等式意味着：如果某种并行策略的通信量可以表达为 reduce-scatter 加 all-gather，那么它的通信成本与一次 all-reduce 完全相同——即“免费”实现内存节省。

三、数据并行与 ZeRO

3.1 朴素数据并行（DDP）

将批次 B 切分为 M 份，分别送入 M 个 GPU，各自计算梯度后做一次 all-reduce 同步。

计算扩展：理想，只要每个 GPU 上有足够样本。
通信成本：每个批次需通信 2 倍参数量。
内存扩展：完全无节省——每个 GPU 持有完整模型副本。

3.2 内存核算：为什么 DDP 糟糕

每个训练步骤需要存储：

模型参数（权重本身）
梯度（更新前的累加器）
优化器状态（SGD 的高精度累加器；Adam 的一阶矩 + 二阶矩，通常需高精度存储）

粗略经验法则：每参数约需 16 字节存储（约为 5 份权重复制）。其中优化器状态占大头，是内存消耗的主因。

3.3 ZeRO 三个阶段

阶段	分片内容	通信成本	内存节省
1	优化器状态	与 DDP 相同（零额外成本）	优化器状态除以 GPU 数
2	优化器状态 + 梯度	与 DDP 相同（零额外成本）	上述两者均分片
3（FSDP）	优化器状态 + 梯度 + 参数	略多于 DDP（多一次 all-gather）	三者均分片

3.4 ZeRO-1 与 ZeRO-2 为何“免费”

ZeRO-1：每个 worker 只负责更新参数的一个切片。先对梯度做 reduce-scatter（分发到各 GPU 的更新对应部分），各 GPU 完成本地更新后，再对更新后的参数做 all-gather。两个操作之和等价于一次 all-reduce，通信成本与 DDP 相同。
ZeRO-2：在反向传播时增量计算梯度——每算完一层的梯度立即做 reduce 并发送给对应 worker，然后释放该梯度。通信总量不变，但无需物化完整梯度向量。

3.5 ZeRO-3 / FSDP：按需抓取参数

FSDP 的关键机制：

前向：逐层 all-gather 该层权重 → 计算 → 释放权重。
反向：逐层 all-gather 权重 → 计算梯度 → reduce-scatter 梯度 → 释放权重。
更新：每张卡更新自己持有的参数切片。

两个关键设计：

沿计算图扫描：通信紧贴计算发生，用完立即释放内存。
通信与计算重叠（overlap）：在前向计算第 k 层的同时，预取第 k+1 层参数。若计算量足够大且网络足够快，FSDP 可实现近乎零额外开销。

注意：FSDP 中每个 GPU 都遍历整个模型（与流水线并行不同，不是按层切分）。

3.6 FSDP 的局限

不能减少激活值内存——这在长序列任务中往往是主要瓶颈。
并行度受限于全局批次大小：如果批次大小为 8，最多只能使用 8 个 GPU。

四、模型并行：切分层与矩阵

4.1 核心区别：从通信参数到通信激活值

FSDP / 数据并行：通信的是参数（权重），计算本身是单 GPU 的完整前向/反向。
模型并行：通信的是激活值——各层分布在不同的 GPU 上，GPU 间传递的是层间的中间结果。

4.2 流水线并行（Pipeline Parallelism）

概念：按层切分模型，不同层放在不同 GPU 上，前向传播时激活值逐层传递，反向传播时梯度逐层回传。

问题——Bubble：朴素实现中，任意时刻只有一个 GPU 在计算，其他 GPU 全部空闲。利用率极低。

解决——微批次（micro-batch）流水线：将批次切分为多个微批次，第一个微批次算完第 0 层后立即交给第 1 层，同时第 0 层开始处理第二个微批次。利用率 ≈ 1 - (阶段数 / 微批次数量)。

通信特性：仅做点对点通信，数据量为 b × s × h（批次 × 序列长度 × 隐藏维度），远小于通信完整参数矩阵。这是所有模型并行策略中通信最温和的。

实践建议：流水线并行最常用于慢速链路（如跨数据中心、跨 Pod），因为它的通信模式最温和。

“零气泡”（ZeroBubble）思想：在反向传播中，有两件事必须做——(a) 将偏导数向计算图深处传播（必须尽快完成），(b) 计算当前权重的梯度（可以随时完成）。将两者解耦：优先做 (a)，把 (b) 推迟到计算间隙执行，可几乎完全填满流水线。

4.3 张量并行（Tensor Parallelism）

概念：沿宽度维度切分矩阵乘法。MLP 的输入投影按列切分，输出投影按行切分；注意力同理。层归一化、dropout、路由器等小操作完全复制到每张卡上。

前向/反向对偶性：

前向：输入复制（f = identity），输出做 all-reduce（g = all-reduce）。
反向：对偶——g 变 identity，f 变加法（偏导数需要累加）。

通信代价：每次矩阵乘法后都需要一次 all-reduce，通信量约为 2 × b × s × h，且是全局通信。因此张量并行只能在节点内（约 8 张卡）使用，跨节点性能急剧下降。

与 TPU 的关系：TPU 的环形网格天然适合高带宽、规则通信的张量并行，因此 TPU 上可以做到更大的张量并行度；GPU 上张量并行超过 8 后性能下降明显。

4.4 激活值内存的精确核算

基础经验公式（存储所有中间结果时）：

激活值 ≈ 34 × s × b × h + 5 × a_s / h

34 项中，MLP 占 24，注意力占 10。
第二项来自注意力 softmax 的 dropout，可通过 Flash Attention + 重计算消除。

张量并行的效果：

MLP 与注意力部分除以张量并行度 t：✓ 可扩展
残差输入、层归一化输入等不随 t 缩减：✗ 固定项（约 10 × s × b × h）

序列并行（Sequence Parallelism）：将剩余的不随张量并行缩减的项沿序列轴分片存储。需要时 all-gather，用完 release。这与 FSDP 的思路一致（分片存储、按需物化），前向的 all-gather 在反向中变为 reduce-scatter。

激活内存下限公式：

激活值下限 ≈ 34 × s × b × h / t  （t = 张量并行度，配合激活重计算）

此公式可用于快速估算模型能否装入单卡。

五、专家并行（Expert Parallelism）

5.1 基本思想

将 MoE 的 FFN（专家）分布到不同设备。与前几种策略类似，专家并行也属于模型并行家族——切分 MLP 而不切分注意力。

5.2 与张量并行的比较

优先选择专家并行而非张量并行（Megatron 官方指南）：

张量并行把矩阵切小，降低 GPU 利用率。
MoE 的 token 路由天然稀疏，专家并行可以精确路由，避免无谓计算。
专家并行同样能减少激活内存。

代价：专家并行需要大量全对全通信，延迟极其敏感。工程实现非常复杂（如 DeepSeek 的 DPP 库甚至使用了未公开的 PTX 指令来优化网络通信）。

5.3 关键约束：DP 与 EP 的耦合

朴素实现中，EP 的副本与 DP 的副本相同——即 EP 是 DP 切分的子集。这限制了 EP 可扩展的最大并行度。

5.4 注意力层与 MLP 层的张量并行冲突

注意力层需要较高的张量并行度才能有效切分。
MLP 层如果张量并行度过高，矩阵被切得太碎，利用率下降。
现代系统将注意力层的张量并行与 MoE 层的张量并行解耦，分别独立设定。

六、上下文并行（Context Parallelism）

将长序列的激活值沿序列维度拆分到多个加速器，采用环形传递方式。主要用于长上下文扩展阶段和推理服务。与序列并行不同，它是独立于张量并行的并行原语，多用于 TPU 等具备环形拓扑的硬件。

七、并行策略总览与组合原则

7.1 各策略的局限总结

策略	优点	局限性
FSDP / ZeRO-3	内存节省大、通信可重叠	不省激活值；受全局批次大小限制
流水线并行	通信最小、适合慢链路	气泡开销大、依赖批次大小
张量并行	无气泡、省激活值	通信量大、只能节点内用
专家并行	省激活值、适合 MoE	全对全通信延迟敏感、工程复杂
上下文并行	处理超长序列	需要专门硬件拓扑支持

7.2 实用组合策略（简单的经验法则）

两层策略：

装进内存：先用张量/专家并行处理快速互联（每机 8 卡），再用流水线或 FSDP 处理慢速互联。
充分利用算力：模型装下后，剩余 GPU 全部用于数据并行。批次太小时用梯度累积弥补。

优先级顺序（从高到低）：

最小化模型并行，最大化数据并行。
GPU 上保持张量/专家并行在节点内（NVLink 范围）。
跨节点用流水线并行。
MoE 模型优先使用专家并行。
长序列使用上下文并行（环形注意力）。

7.3 性能分析框架

对给定模型和并行配置，计算每层的计算量与通信量，绘制“计算时间 / 通信时间”图表。只要计算时间 > 通信时间，通信就可被完全隐藏，GPU 利用率达到 100%。策略选择的目标是：在不同批次大小下都保持计算受限。

八、实践案例

Olmo（7B 稠密）：纯 FSDP 训练。小模型规模下 FSDP 扩展性良好。
DeepSeek V1：数据并行 + ZeRO-1 + 张量并行 + 序列并行 + 流水线并行。
DeepSeek V3：MoE 模型，用大规模专家并行（64 路）替代部分张量并行，结合流水线技巧避免气泡。
Yi：数据并行 + 张量并行 + 流水线并行（经典组合）。
Llama 3 405B：预训练阶段为 TP=8 + CP=1 + PP=16 + DP=128；长上下文阶段将 CP 提升、DP 降低。
Gamma 2：仅用 FSDP + 张量并行 + 序列并行，无流水线（利用 TPU 大网格拓扑）。
Mixtral 8x22B：EP=8 + PP=4 + TP=4（用于注意力层）。
Qwen 3：EP=32 + PP=8 + TP=2。

共性规律：

数据并行最大程度使用。
张量并行不超过 8（GPU 环境）。
专家并行可做大（受 MoE 工作负载影响，DeepSeek 推动了大专家并行规模）。

九、课程作业关联

作业要求：给定网络拓扑（节点内带宽、节点间带宽）与目标模型，寻找最优并行化策略组合。你需要：

核算各策略的每层通信量。
核算各策略的内存占用（参数 + 优化器状态 + 激活值下限公式）。
判断通信能否被计算隐藏。
可能要求自行实现 FSDP 包装器（对任意模块做 all-gather → 计算 → release → reduce-scatter 的循环）。

常见误区

认为 FSDP 省一切内存：FSDP 不省激活值，长序列场景下激活值才是主要瓶颈。
混淆序列并行与上下文并行：序列并行是张量并行的附属（切序列分片存储，目的是省激活值）；上下文并行是处理超长序列的独立并行策略。
认为流水线并行的气泡可以忽略：气泡随微批次数量下降，需要足够大的批次才有效率。
在节点间使用张量并行：超出 8 张卡范围，通信开销急剧上升。
混淆 all-reduce 与 reduce-scatter + all-gather 的成本：二者等价，这是 ZeRO-1/2 免费的关键。
认为专家并行可无限扩展：EP 的扩展受限于 DP 切分结构，且全对全通信延迟极难优化。
忽视上下文并行在长上下文阶段的作用：模型服务与长上下文扩展中，上下文并行是必需的。

复习清单

[ ] 能说出数据并行与模型并行的根本区别（通信参数 vs. 通信激活值）。
[ ] 能解释 all-reduce = reduce-scatter + all-gather 和内存节省的对应关系。
[ ] 能列出 ZeRO 三个阶段的分片对象与各自的通信成本。
[ ] 能讲清 FSDP 如何通过“按需抓取”和“通信计算重叠”实现近零开销。
[ ] 能解释流水线并行的气泡来源、微批次的作用、以及 ZeroBubble 的核心思路。
[ ] 能写出张量并行的前向/反向对偶关系（f ↔ g 互换）。
[ ] 能计算激活值下限：34 × s × b × h / t。
[ ] 能说明序列并行要解决什么问题、为何与张量并行组合。
[ ] 能说明为什么 MoE 优先选专家并行而非张量并行。
[ ] 能按优先级列出组合策略的实践原则（DP 最大化、TP/EP 在节点内、PP 跨节点、CP 长上下文）。

从单卡到数据中心：大语言模型训练中的并行化策略全景

导语

当我们在单张 GPU 上训练一个 70 亿参数的模型时，内存会立刻告急；当我们要训练一个数千亿参数的模型时，算力更是遥不可及。大语言模型规模的指数级增长，迫使我们把目光从单个芯片移向整个数据中心。但如何把成千上万张 GPU 组织起来，让它们像一个巨大的虚拟芯片一样高效运转？这正是并行化要回答的核心问题。

在 CS336 课程的第八讲中，我们深入探讨了现代大语言模型训练中并行化的方方面面。这一讲建立在第七讲对通信底层机制的基础上，将抽象的网络原语与实际的工程实践连接起来，最终呈现出一幅完整的并行化策略图谱。本文将从数据并行讲起，逐步过渡到各种模型并行策略，最后探讨它们如何组合成我们在前沿模型训练中看到的大规模系统。

两个瓶颈与两种网络哲学

并行化的必要性源于两个根本瓶颈：算力与内存。单个芯片永远无法提供训练千亿参数模型所需的算力，模型本身也远远超出任何单卡的内存容量。于是，我们把多台机器连接起来，把模型切碎，把工作分发下去。

但这里立刻浮现出一个关键区分：节点内（intranode）并行与节点间（internode）并行。节点内的 GPU 通过超高速互联连接，通信开销极低；节点间的连接则慢得多。这个区别贯穿了整讲的内容——我们选择哪种并行策略，很大程度上取决于通信发生在哪一层。

有趣的是，不同的硬件厂商对网络拓扑有着截然不同的设计哲学。TPU 采用环形网格（toroidal mesh）拓扑，每个芯片只与固定数量的邻居相连，邻居之间环绕相接。这种设计的优势在于：无论网络扩展得多大，每个芯片的邻居数不变，通信模式极其规则且可预测。相比之下，GPU 网络更像一棵胖树（fat tree），从底层的快速连接逐层向上汇聚，理论上支持任意两个节点之间的全对全通信，但代价是随着节点数增长，网络复杂度急剧上升。

这两种拓扑直接影响了并行策略的选择。TPU 擅长那些通信模式规则的并行方式，而 GPU 网络则更适合通信模式不那么可预测的场景——比如混合专家（Mixture of Experts）模型，token 被动态路由到不同的专家，通信模式本身就带有随机性。值得注意的是，讲者提到谷歌最新发布的 TPU 已经转向更接近全对全的拓扑，这被视为硬件与工作负载共同演化的一个信号。

数据并行：最朴素的想法及其局限

数据并行（data parallelism）是最容易理解的并行策略。假设我们有一个大小为 B 的批次，把它均匀分到 M 台机器上，每台机器拿到 B/M 个样本，计算各自的部分梯度，然后同步所有梯度。

从算力角度看这很完美——只要每个 GPU 上有足够多的样本。但问题在于通信和内存。每处理一个批次，我们需要通信约两倍参数量大小的数据；而在内存方面，每个 GPU 都持有一份完整的模型副本，完全不节省内存。

更糟的是，现代优化器让内存问题雪上加霜。讲者给出了一个粗略的经验法则：我们需要大约五份权重复制，每参数 16 字节。这包括模型参数本身、梯度累加器，以及优化器状态（optimizer state）——包括一阶动量和二阶动量。在使用 Adam 优化器时，优化器状态往往占据了内存的大头，而这也是数据并行最浪费资源的地方：每个 GPU 都在冗余存储这些状态。

ZeRO 与 FSDP：数据并行的内存革命

如果优化器状态占内存最多，为什么不把它们切碎分布到不同 GPU 上？这引出了 ZeRO 的三个阶段。

ZeRO 阶段 1：只切分优化器状态。每个工作节点负责更新参数的一个切片。流程是：先计算完整梯度，然后做一次 reduce-scatter 把梯度分发到各节点，各节点完成自己的参数更新，最后 all-gather 收集更新后的参数。关键洞察是：reduce-scatter 加 all-gather 的组合在通信量上等价于一次 all-reduce。因此，ZeRO 阶段 1 与朴素数据并行相比，通信成本完全相同——内存节省是白拿的。

ZeRO 阶段 2：进一步切分梯度。这带来一个微妙的问题：既然梯度也被分片了，我们就不再能完整物化整个梯度向量。解决方案是沿着计算图反向传播时，计算完某一层的梯度后立即做归约并发给正确的工作节点，随后释放。这是增量式地完成原本一次性做的事情，通信成本不变，但梯度也不再重复存储。

ZeRO 阶段 3，也就是广为人知的 FSDP，最激进也最优雅：一切都被分片——参数、梯度、优化器状态。每个 GPU 在任何时刻只持有模型的一个切片。前向传播时按需 all-gather 当前层的参数，计算完成后立即释放；反向传播时再次 all-gather 参数，拿到梯度后立刻 reduce-scatter 出去。通信量比阶段 2 多了一次 all-gather，似乎更昂贵了。

但 FSDP 真正精妙的地方在于它把通信和计算重叠（overlap）起来。在等待第 0 层参数到位并开始前向计算的同时，可以发起对第 1 层参数的 all-gather。只要网络足够快、计算量足够大，FSDP 可以做到几乎零开销。讲者强调，FSDP 的实际 GPU 利用率可以非常接近单 GPU 性能，而作业中要亲自实现一个 FSDP 包装器——只需把任意模块包装起来，做一堆 all-gather、计算、释放，然后反向过程中重复。

FSDP 的收益是巨大的：在 8 块 A100 上，从连 70 亿参数的模型都放不下，到可以装下 500 亿参数的模型。然而，FSDP 有两个根本局限：其一，它消耗全局批大小，如果你的批次大小是 8，就永远不可能用超过 8 个加速器有效训练（受临界批大小约束，批次增大带来的收益会递减）；其二，它无法削减激活值内存——那些在前向过程中产生、反向传播时必需的中间量。

模型并行：切分层、矩阵与专家

当数据并行的内存节省到达极限，我们必须求助于模型并行（model parallelism）。与 FSDP 切分参数但仍在每个 GPU 上遍历完整模型不同，模型并行把不同部分真正放置到不同设备上，通信的对象也从参数变成了激活值。

流水线并行

流水线并行（pipeline parallelism）是最直观的模型并行：把层切分开，放在不同 GPU 上，激活值在层间传递。朴素实现的问题是灾难性的气泡（bubble）——任何时候只有一个 GPU 活跃，利用率低得令人绝望。

解决方案是微批次（micro-batches）：像流水线一样，微批次依次进入，上层完成一个微批次的计算后立刻传递给下一层，而不是等整个批次算完。利用率大约等于阶段数除以微批次数量，因此它消耗批大小来换取效率。实践中，流水线并行的通信特征非常好——只依赖 b·s·h 大小的激活值，而且是点对点通信——因此它最适合放在最慢的链路上，比如跨数据中心。

进一步优化空间在于更聪明的调度。讲者提到了"零气泡流水线"的思想：反向传播实际上包含两件事——把偏导数传回计算图更深处，和计算当前权重的梯度。前者有严格的时间依赖，后者是叶节点，随时可算。把这两者解耦，优先做前者，在空闲间隙填补后者，几乎可以填满流水线。

张量并行

张量并行（tensor parallelism）沿宽度方向切分矩阵乘法。简单的想法：把 A 矩阵切成两块，B 矩阵也切成两块，各自并行计算，最后汇总。前向传播时 f 是恒等操作（把输入复制到两份计算中），g 是 all-reduce（把结果合并）；反向传播时角色反转。

实际应用中，MLP 的输入投影和注意力机制的投影做按列切分，MLP 的降维投影和注意力输出做按行切分，而层归一化等小层完全复制。张量并行对通信极为敏感——每次矩阵乘法都伴随一次所有规模的通信，而且非常频繁——因此通常只在节点内部使用。对于 GPU 用户，这通常意味着最多 8 卡。但 TPU 的环形网格拓扑恰好适合这种规则的通信模式，因此 TPU 上可以把张量并行扩展到更大的规模。

序列并行

讨论激活值内存时，我们得到一个经验公式：激活值约需 34·s·b·h（序列长度 × 批大小 × 隐藏维度）加上注意力项。张量并行可以削减其中大部分，但层归一化、dropout 等逐点运算不参与切分，仍然留下约 10·s·b·h 的不可削减项。

序列并行（sequence parallelism）的名称极具误导性——它并不是把序列分给不同设备各自独立处理（那是上下文并行），而是沿序列轴切分这些逐点操作需要的激活值。这非常像 FSDP 的思路：分片存储，按需 all-gather 具体化，反向传播中角色反转。通过序列并行，激活值内存可以实现对张量并行规模的完全线性依赖，配合激活重计算消除注意力项，最终得到激活内存下限：34·s·b·h 除以 t。

专家并行

专家并行（expert parallelism）把 MoE 层中的专家（FFN 组件）分布到不同设备上。它在多个方面与张量并行相似：高带宽消耗、减少激活内存，但它有几个独特优势：矩阵保持更大因此 GPU 利用率更高；token 可以精确路由到需要去的专家；在 MoE 模型中，优先选择专家并行而非张量并行是 NVIDIA Megatron 官方指南的明确建议。

然而专家并行极其复杂，因为涉及大量的全对全分发，且要求低延迟——计算在等待 token 到达。讲者提到 DeepSeek 的 DPP 库甚至使用了未公开的 PTX 指令来加速网络通信——这已经是接近机器码级别的优化了。

专家并行还有一个系统层面的纠结：MoE 只改变 MLP，不改变注意力。注意力层需要高张量并行来切分，而 MLP 层如果同时使用高张量并行和高专家并行，矩阵会被切成碎片，利用率急剧下降。近年来的解决方案是把注意力和 MLP 中的张量并行解耦，各自使用独立的并行策略。

上下文并行

上下文并行（context parallelism）或环形注意力，用于处理极长序列：把激活值沿序列轴拆分到不同加速器上，以环形方式传递。这主要用于长上下文扩展阶段和服务场景，概念上与前述内容有大量重叠。

组合策略：从乐高积木到 4D 并行

所有并行策略都各有优劣。FSDP 不能处理激活值、消耗全局批大小；张量并行能削减激活值但不影响批大小，却需要高速网络；流水线并行通信效率高但依赖大批次。没有任何一种策略是全局最优的——这个观点通过并行策略对比表格中的红色标记（缺点）被反复强调。

实践中的策略其实非常简单：首先，只要模型还不能完全放进内存，就想尽一切办法切分。在快速互联（节点内部）使用张量/专家并行，通常为 8；用流水线并行或 FSDP 处理剩余部分。一旦模型能放进去了，就全部使用数据并行来利用剩余 GPU。如果批大小太小，用梯度累积弥补。

这也正是 Megatron 的 MoE 并行化指南的核心内容：最小化模型并行，最大化数据并行；在 NVLink 范围内（即单机箱内）使用专家/张量并行；跨节点用流水线并行；MoE 优先专家并行；长序列用上下文并行。

现实案例：前沿模型的并行化配置

课程最后展示了一系列真实模型的并行化方案，让我们看到理论如何落地。

Olmo 7B 完全使用 FSDP 训练，展示了 FSDP 在小模型上可以扩展到大量 GPU。许多 7B 级别的模型都是用纯 FSDP 训练的。

DeepSeek V1 使用了数据并行加 ZeRO 阶段 1，加上张量、序列和流水线并行。DeepSeek V3 作为 MoE 模型，用专家并行取代了张量并行，并且采用了独特的 64 路专家并行——8 台机器组成一个专家并行域，借用了流水线的技巧来实现大规模专家并行。

Llama 3 405B 是少数公开详细并行化配置的模型之一：主预训练阶段使用张量并行 8、上下文并行若干、流水线并行 16、数据并行 128。长上下文扩展阶段提高上下文并行、降低数据并行。讲者顺带提到，405B 训练期间 GPU 故障了 148 次——大规模训练不仅是算法问题，也是分布式系统的挑战。

Gemma 2 仅用 FSDP 加张量并行和序列并行——讲者认为这基本实现了谷歌的主张：对于 TPU，你不需要流水线并行，只需一个大网格加张量并行。

近期模型如 Mixtral 8x22B（专家并行 8 + 流水线并行 4 + 张量并行 4）、Nemotron 3 Super（大量专家并行加上下文并行）、Qwen 3（专家并行 32 + 流水线并行 8 + 张量并行 2），都遵循了同一个模式：最大化数据并行、张量并行保持在 8 以下、专家并行可以做得很大。

结语

从单张 GPU 到整个数据中心，并行化的本质是一场资源的精算：算力、内存、通信带宽、批大小——每种资源都有限，每种并行策略都在某种资源的约束下换取另一种资源的释放。数据并行给我们算力但不省内存；FSDP 免费削减了参数、梯度和优化器状态的内存，却对激活值无能为力；张量并行削减激活值但吞噬通信带宽；流水线并行通信高效却依赖大批次；专家并行在很多方面更好但它带来的系统复杂度也最高。

理解这些权衡之后，组合各种策略其实遵循着简单的经验法则。而现实中那些训练出最强模型的公司，正是在这些经验法则的基础上，结合对自己硬件拓扑的深刻理解，构建出了高效的大规模并行系统。在课程的整体路径中，这讲标志着我们从"如何实现一个模型组件"走向"如何在规模上训练一个系统"。下周讲到的缩放定律，将告诉我们规模的极限在哪里，以及如何智能地分配资源。

官方资源：

视频：CS336 Lecture 8 官方视频
课程讲义及代码：CS336 Language Modeling from Scratch

CS336 2026 Lecture 8：完整中文逐字稿

视频：官方视频
来源：英文人工字幕提取后翻译为中文；保留段落起始时间戳，只修正明显的断句、重复和语序。
说明：这是逐字稿，不是讲义摘要；相邻且语义连续的片段已合并为阅读段落；术语保留必要英文。

正文

[00:05.240] 上一讲中，Percy 介绍了并行化的一些底层机制，我们姑且这么称呼它。今天，我要做我最喜欢的事情，就是聊聊各种你可以了解的知识、细节和花絮，关于现代大语言模型训练的并行化是如何运作的。我们今天的核心目标，是真正深入理解当你在庞大的集群上训练巨大的模型时所涉及的众多复杂性。事实证明，你可以同时使用许多不同的并行化策略。在最后，我们会讲到一张幻灯片，上面写着“4D 并行”，因为你可以同时做四件不同的事情。实际上，我还会谈到更多。而且事实证明，要在最大规模上进行并行化，你需要使用所有这些策略，或者至少是其中大部分。

[00:56.720] 我们会讨论为什么。你们作业的一部分将是，针对给定的特定网络拓扑和特定模型，找出最优的并行化策略是什么。最后，我会带着大家看一些近期大规模训练运行的例子，这些例子有公开的细节，并向你们展示这些方法在实际中的应用。好的，第一部分我会讲得稍微快一些。这部分某种程度上是对基础集合通信、基础网络原语的复习。但是，我们需要所有这些并行化的东西，原因在于我们将要解决两个瓶颈，而解决它们的方式就是使用多台机器和多个 GPU。其中一个瓶颈是算力。我们需要比单个芯片能提供的更多的算力。我们的解决方案就是拥有非常非常多的机器，并把它们连接起来。是的，世界上最快的超级计算机拥有百亿亿次（exaflops）的算力，而单个 GPU 远远达不到这个数字。

[01:55.080] 我们需要持续并行扩展的另一个原因是内存。模型太大了。我们没法把它们全部塞进一个 GPU 里，所以需要以某种方式将它们切分成更小的部分。这就会带来复杂性，也就是要在多台机器上使用多个 GPU。今天整个讲座中我们会反复遇到的一个重要概念性对象，是节点内（intranode）并行和节点间（internode）并行的区别。节点内连接速度非常快，因此你可以做通信开销很大的事情。而节点间通信则会慢一些。正因为这些通信较慢，我们将不得不采用更尊重信道通信约束的流程。所以我们两者都得做。请记住我们今天要做的所有事情。实际上，我们并不会真的去发送数据包或做类似的事情。

[02:44.860] 我们将在集合通信原语的层面来处理这些。所以，每当我们计算完成了多少通信量时，我可能会说，好的，我们需要做一次全规约（all-reduce），或者我们需要做一次全收集（allgather）来实现这些操作中的某一个。为了高效地实现其中一些操作，你可能需要深入到硬件层面。但我们今天的讨论将是算法层面的。所以我将谈到的所有网络原语都停留在集合通信的层面。记住，Percy 讲过这个，但这对于——对于讲座的某一部分非常重要，那就是全规约（all-reduce）与归约散播（reduce-scatter）加全收集（all-gather）之间的等价关系。也就是说，有一种算法，一方面你做的是全规约，另一方面你做的是两个不同的步骤。

[03:31.180] 而这两个操作在某种意义上成本相同的事实意味着，我们可以“免费”执行右边的算法。我们稍后会再次看到这一点。这只是个快速回顾，确保你们还记得 Percy 讲过的内容。好了，基础回顾就到这里。在进入那些与具体硬件和底层平台无关的算法之前，我还想简单提一下另一个区别——哦，抱歉，怎么了？你特别提到……你是在问，我为什么特别提到这个归约散播（reduce scatter）和……好的，问题是，我是不是因为什么特殊原因才特别讲这个？我们将——[听不清] 是的，是的，还有其他分解方式。但这个特定的分解方式对一个算法特别有用，这就是我强调这个特定分解的原因。没错。当然，你也可以做其他的分解。好的。那么，在我开始之前，我要讲的另一件事是——我们今天的讨论大部分将是与硬件无关的。

[04:27.810] 在很大程度上，一旦你理解了基本概念，你就能把它应用到任何硬件平台上。但我确实想在开始之前，简要地谈谈底层的硬件栈。这既因为有趣，也因为我认为它导致了不同模型训练公司之间采用了不同的并行化策略。还记得 TPU 吧。TPU 是谷歌的加速器。我在 GPU 那讲提到过，TPU 就像是个轻量级的 GPU。但不同之处在于网络。那么，网络有什么不同呢？简单来说，在一张幻灯片里，TPU 的网络被称为环形网格（toroidal mesh）。你可以把它想象成一个网格，然后允许网格的两端互相通信。实际情况稍微复杂一点，它是一个三维的图形。Marcel，你有那个可视化的东西，对吧？你能发到讲座频道里吗？好的。

[05:21.870] 有一个很有趣的网页可视化，你可以在里面看到三维的环形网格，非常酷。但你的心智模型大致可以参照这张图，可以忘掉那个更复杂的。所有芯片都与邻居相连，并且邻居是环绕相接的。那这能让你做什么呢？这让你拥有一个非常简单的网络拓扑，可以非常简单地无限扩展。无论你的网络变得多大，你的邻居数量都是相同的。相比之下，GPU 的网络理念则非常不同。它更像是全对全（all to all）的理念。它的组网方式像一棵树，通常被称为胖树（fat tree）。在最低层，你的 GPU 以非常非常快的速度连接。然后，你有一组一组的 GPU 连接成一个 Pod（机柜组）。

[06:05.850] 然后这些 Pod 可能有脊交换机（spine switches）让它们之间可以互相通信。这意味着随着节点数量的增长，这棵树会变得越来越庞大。你的通信成本会变得更高，或者你的通信拓扑会变得更复杂，因为你在连接所有这些点，实现全对全通信。这种理念上的差异意味着，如果你只和邻居通信，TPU 很棒。它们成本效益高，并且你可以在相同功耗下把每条连接做得更“粗壮”。但如果你是以非常随机的方式、非常随机应变且不可预测的方式进行连接，那么这种 GPU 网络会更加灵活。所以你可以看到 TPU 和 GPU 之间存在非常不同的设计理念。在我们继续并行化讲座的过程中，我会试着指出这些差异体现在哪里。上一讲我已经把这个发到讲座聊天区了，因为有人问过这两者有什么区别等等。

[07:02.310] 是的，为了支持我的观点，Bill Daley 和 Jeff Dean 曾讨论过不同的硬件以及它们各自擅长什么。GPU 擅长混合专家（Mixture of Experts）这类任务，因为 token 会路由到不同的专家，通信模式不太明确。如果你有一个密集模型，并且要进行非常可预测的切分，TPU 是很棒的。我稍后会在本讲中介绍一个叫做“张量并行”的概念时讨论这一点。所以这些硬件各自适用于不同的工作负载。我本来打算就讲到这里，直到今天早上，谷歌决定他们要发布 TPU8i。今天早上我看着这个准备讲义，我说，天哪，TPU8i？那是树状拓扑。他们已经转向了更接近全对全的连接方式。如果你仔细想想，这是有道理的。

[07:54.630] 现代语言模型都是 MoE。如果你要服务一个 MoE，当你把 token 路由到不同的专家时，会有各种各样的带宽需求飞来飞去。你可能确实需要更多的全对全连接，因为在推理时，这些通信是真正的瓶颈。你需要处理它们。即使在训练时，TPU8t 这个训练芯片的跨机架连接看起来也更像 GPU 了。他们正在构建一个更像全对全交换的网络，叫做 Virgo 网络。那是一个更高层的网络协议栈。所以我认为现在正发生着非常有趣的事情。TPU 和 GPU 之间有一点趋同演化的味道，实际上工作负载正在定义我们需要的网络。好。关于硬件的最后一点——我非常喜欢谈论中国 AI 领域正在发生的事情——是这样一个想法——我记得在 GPU 那讲中，有人

[08:55.670] 问过，如果 SRAM 这么好，为什么不全部用 SRAM？只用 SRAM 芯片。那就是 Grok，对吧？它效果很好。现在，你们中的一些人可能会想，如果全对全连接很好，为什么不把一切都连起来，然后用非常快的光纤芯片连接它们呢，对吧？然后你就得到了华为昇腾 910，对吧？如果你看一下华为昇腾的规格，它其实非常有意思，因为每颗芯片都比 H200 差很多。矩阵乘法慢很多。很多方面都更慢。但是，如果你愿意用一个巨大的光纤交换机机架来连接数量多得多的芯片，就像他们那样，在一个大机架里放 384 颗，那么也许你可以通过连接数量大得多的这些芯片来应对单颗芯片较慢的问题，对吧。所以我认为这向我们展示了一个非常有趣的权衡。因为一旦你看到这种架构，你就会意识到这个东西的功耗高得离谱。

[09:58.440] 它是同等 NVIDIA 系统的四倍。所以这就变成了，如果你愿意承担功耗成本，你可以用蛮力解决很多通信问题，并且可以更激进地扩展到 300 颗芯片。但你是在付出功耗的代价，对吧。这和 SRAM 的故事类似：如果你想设计一个在功耗和制造方面都高效的硬件，你会得到一个特定的设计结果。如果你想用蛮力解决问题，你会得到一个截然不同的结果。对，所以我认为这是了解这些硬件设计权衡的一个非常有趣的点。好的。所以现在，把所有这些放在一起，你首先应该想到的是，计算的新单元不再是 GPU，而是整个数据中心。我们需要实现几件事。我们希望能够控制我们使用的内存量。我们希望能够控制我们可以调用的计算量。

[10:51.680] 而且我们希望所有这些都无损。对，我们希望使用我们所拥有的所有资源。然后我们将讨论通过这些通信原语组件来实现这一点的算法。好的。现在，本讲的第二部分将是内容的主体。这些是我们将用来并行化语言模型的算法。这里会有几个不同的重要思想。实际上，我们可能应该对它们进行分组。数据并行是一组思想，本质上就是把你将要训练的数据分发到各个 GPU 上。所以是数据在移动。然后是模型并行，你将模型的不同部分切分开并分散放置。事实证明，这两者之间的界限有点模糊，因为这里有一种算法实际上也会切分参数。但我们稍后会明白为什么我把所有这些算法放到这些不同的类别里，对吧。从概念上讲，我们既要扩展计算，也要扩展内存，这就是为什么我们需要这一大堆不同的技巧。酷。好的。那么数据并行，我认为是最容易理解的。

[12:04.360] 它在概念上非常简单，我认为这是你尝试并行化像 LLM 训练这样的算法的标准方式。假设我们不用 Adam，我们暂时用朴素的 SGD。那么我们有一个大小为 B 的批次，对吧。所以每一步更新，我要取 B 个元素，然后计算梯度之和。然后进行更新，对吧。所以最朴素的并行方式就是把我的批次 B 切开，分到 M 台机器上。如果它们能整除，每台机器会得到一个大小为 B/M 的更小的批次。然后我要做的是每台机器计算梯度，然后我同步所有梯度来得到总和。好的，太棒了。从计算扩展性来看，这很完美。只要每个 GPU 上有足够多的样本，就没问题。通信开销呢？

[12:53.800] 我们需要彼此通信多少？嗯，每处理一个批次，我们基本上需要通信两倍于参数数量的数据，对吧。因为我必须做的是基本上来回通信我的梯度。好的，然后在内存扩展方面，完全没有。每个 GPU 都有相同的模型副本。你需要同样大小的激活值，所以在这种方案中你完全不会节省任何内存。那么问题是什么？显然我们在一定程度上解决了计算问题。但我们完全没有解决内存问题。所以让我们来看看这个内存问题，做一点核算。最初，可能看起来内存情况很糟糕。但事实证明，我们的内存情况只能用“糟糕透顶”来形容。原因在于，如果你只是对参数做一下核算，对吧，你会发现你需要存储很多东西。

[13:49.308] 对，我们不仅需要存储很多东西，实际上比需要的参数量还要多。粗略的经验法则——这取决于你的精度。但经验法则，你可能会说，我们需要大约五份权重复制，每参数 16 字节来存储我们的模型。所以我们当然需要存储我们的模型参数。我们需要一个地方放我们的梯度，对吧。也就是在把它们更新到模型参数之前，那个累加器。所以我们需要存储那些东西。但实际上，你可能需要存储更多的东西。在 SGD 中，你可能需要一个更高精度的累加器来进行累加。这可能是临时的，但你可能需要它。如果你用的是 Adam——这才是真正糟糕的地方——你需要追踪梯度的一阶矩和二阶矩。而且你需要随着时间推移持续追踪它们，对吧。所以你需要两个这样的东西。

[14:40.700] 而遗憾的是，根据你的稳定性特征，你可能需要以高精度存储它们。而这些可能相当昂贵，对吧。我们通常把这些东西称为优化器状态（optimizer state）。对，就是你要累加进去的那个东西。一阶动量，二阶动量。对吧，这些实际上相当昂贵。如果你看一下账目，从内存角度看，这是执行一次 SGD 更新的大部分成本。所以我认为，在这个领域你能做的一件非常优雅和美妙的事情，就是开始削减你所需的内存量，并将它们分布到各个加速器上，但要以这种结构化的方式来做，这样你就能基本上免费获得这些收益。所以正如我上一张幻灯片所说，优化器状态占了你内存的大部分。所以这里，优化器状态是绿色的。那是你内存的大部分。橙色是你的梯度，蓝色是你的参数。注意参数和梯度的大小是相同的，而你的状态更大，对吧。

[15:39.650] 如果你做朴素的數據并行（data parallel），你基本上会把这些状态在所有 GPU 上复制一遍。这并不太好，对吧。因为你在某种意义上消耗了远超所需的内存，而且你的内存消耗是随着加速器数量线性增长的。不太好，对吧。我们可以做的是，也许可以将优化器状态分片（shard）到不同的 GPU 上。这将是一个巨大的节省。如果我们把我们需要跟踪的其他部分也分片，那也会是巨大的节省。也许我们还可以把梯度分片到不同的 GPU 上。或者极端一点，也许我们可以把所有东西都分片到不同的加速器上。然后你可以看到，在最右边所隐含的内存收益将非常惊人，从 120 一路下降到 1.9。现在，我们接下来要研究的主要问题是，这样做的代价是什么？

[16:31.170] 对吧，你可能想得到，天下没有免费的午餐。基线（baseline）将会是最快的。当我们往下走到这里时，我们将不得不付出巨大的通信成本来实现这些内存节省。我认为值得注意的是，其中很多将是免费的。但我要一步一步地讲解。让我们从人们所说的 ZeRO 阶段 1 开始。这里有三个阶段，分别对应这些步骤中的每一步。阶段 1 是最简单的。我唯一要分片的就是我的优化器状态。我将在 GPU 之间对它进行分片，好吗。每个人都将拥有参数和梯度。所以你可以这样想：每个工作节点（worker）将负责更新参数的一个切片。如果我是 GPU 0 号，我的责任是更新这边的这第一个切片。

[17:18.490] 我不负责任何其他的更新。那么这如何运作呢？对吧，我们要做的是每个人都在计算。每个人都有不同的数据项，对吧。所以他们将在他们的数据项上计算一个完整的梯度。这是第一步。现在我要拿我计算出的梯度，并对其进行 reduce-scatter 操作，分发到所有其他机器上。因为其他所有人都需要我的梯度，对吧？但他们不需要全部的梯度。他们只需要与他们的更新相关的那些梯度。对，所以每个工作节点将从我这里得到我的梯度，但那是参数空间的一部分，对吧。所以这相当于一定数量的参数通信成本。因为我发送我的梯度，但我把它切分开，发送到每个 GPU。

[18:02.050] 现在每台机器都能够进行它们的更新了。它们有它们的参数，有它们从其他所有人那里得到的梯度，还有它们自己追踪的状态，对吧。现在，参数已经被更新了。现在我要把它们 all-gather 回来。对，我负责，比方说，四分之一的参数。那么这些更新后的参数，我要发送回给其他所有人，对吧。所以这是一个 all-gather，对吧。这就是等价性有用的地方。我做了多少次操作？嗯，在朴素的數據并行情况下，我要做的是执行一次 all-reduce。对，我必须与所有人交换所有的梯度，然后进行更新，对吧。所以通信成本是参数大小的 2 倍。现在在 ZeRO 阶段 1，我做了两个操作。一个是 reduce-scatter 用于发送梯度，另一个是 all-gather 用于收集更新后的参数。

[18:52.230] 但我们已经说过，reduce-scatter 加上 all-gather 等价于 all-reduce。所以 ZeRO 阶段 1 将拥有与朴素 DDP 完全相同的通信特征。对吧，所以这是免费的。免费的内存节省。太棒了，对吧？我们来看一下内存。我们基本上拿走了参数。以及任何属于优化器状态的东西，我除以了 n 个 GPU，对吧。太棒了，太棒了，太棒了。现在，我们还能做更多吗？嗯，ZeRO 阶段 2 现在要说的是，我要把梯度拆分。然后我也要把它分片。现在，这会更加棘手，因为之前我依赖于能够计算出完整梯度这一事实，而现在我甚至无法完整地物化（materialize）整个梯度。那我该怎么办呢。嗯，这其实只是一个系统技巧。当我通过网络反向传播时，我不需要计算完整的梯度。

[19:42.670] 我可以在计算过程中将梯度发送给对应的工作节点。所以这里的关键步骤是，现在不再是一次性物化整个梯度向量，而是沿着计算图向后走。在计算完某一层的梯度后，我立即进行归约（reduce）并发送给正确的工作节点，对吧。一旦在反向图中不再需要梯度，就立即释放它。对吧，所以你现在是边计算边增量地发送梯度。正如你所见，增量地做还是一次性做完，是等价的，对吧？所以这并没有什么区别。所以现在你可以在最后，再次 all-gather 参数。我在某种意义上分片了梯度，而没有任何额外的成本，对吧。好的，最后一个也是最棘手的事情——这在我第一次学到的时候，总是感觉像完全的魔法——是我们到目前为止已经免费得到了一些东西，所以让我们继续推进。

[20:36.390] 我也要把参数分片，贯穿所有状态。所以我们将使用在 ZeRO 阶段 2 中使用过的类似想法。记住，这个想法是，我要增量地做事，对吧。所以我要做的是，在逐步遍历计算图时，按需发送和接收参数。这真的很棘手，因为每个 GPU 在任何时候只能看到参数、梯度和优化器状态的一个切片，对吧。这就是 ZeRO 阶段 3，也就是众所周知的 FSDP。如果你曾经用 Torch 库并行化过任何东西，你可能以前处理过这个。它在底层的工作方式是，你要在过程中执行两个 all-gather 和一个 reduce-scatter。所以我们要从这里从左到右逐步进行。这将是整个更新过程。你要做的是加载你的模型。

[21:32.030] 你在顶部和底部有两个 GPU。你首先要为这一层收集你的权重。我要对这一层进行前向传播。然后我可以释放我的权重，因为我已经完成了前向传播，对吧。我不再需要这些参数了。现在，假设我只有一层。那么我就可以进入反向过程。要做反向步骤，我需要激活值，这些我已经保留下来了。我还需要这一层的参数，我可以按需进行全收集（all-gather）。我全收集之后，就做反向计算。然后我一拿到梯度，就立刻把梯度散射（scatter）回去。没错，所以我得到了梯度。我把梯度做归约散射（reduce-scatter）发出去。然后我释放权重。我不再需要它们了，然后我重复这个过程，对吧。所以我是按需抓取参数。

[22:10.750] 我计算反向。我把梯度发出去，然后重复这个过程。接着我更新权重。全部完成，对吧。那么通信成本是多少呢？对，我们做了多少次全收集？我们做了两次全收集，一次归约散射。也就是说，我们多了一次全收集，这看起来不太好。这是一方面。第二点是，从开销的角度来看，这似乎有点疯狂，因为我们一直在做通信。每一层，我们都在做这种通信。你可能会认为这极其昂贵且糟糕。但我觉得其中一个反直觉且重要的点是，ZeRO 阶段 3 用了两个非常重要的思想来让这基本上没有额外开销。第一个思想我们已经见过了，就是沿着计算图扫描，然后做必要的通信，之后立刻释放内存。对，这个思想我们已经体验过了。

[23:07.800] 另一个思想我还没讲到，就是我们要把通信和计算的开销进行重叠（overlap），对吧。这一点一开始不太好解释。所以我觉得我拿了一张来自某篇论文的图，那篇论文我记得是在解释 PyTorch FSDP 的架构。那么它的工作方式是这样的——为了……我们这里有不同的流。这是 CPU 上发生的事情。这是 GPU 上做计算发生的事情，然后最底下是 GPU 上做通信发生的事情。FSDP 的工作方式是，我必须先拿到参数才能做计算。所以我要先对第 0 层做全收集，对吧。一旦我全收集完成，我就可以做前向了。但注意，在我做前向的同时，我正在请求下一层的参数。我在做这个计算的同时，正在对第 1 层做全收集。现在我对第 1 层全收集完了，我就可以做第 1 层的前向了，对吧。

[24:02.680] 而在开始这个过程的同时，我可以开始对第 2 层做全收集。所以我在做前向的过程中，把计算和通信重叠起来了。我可能会花一些时间来释放参数。我也可能会花一些时间计算另一个前向。比如，如果你复用了那个 W₀，我就在两次前向中复用了第 0 层的前向。所以这让你可以做到的是，如果你有足够的计算量并且你的网络足够快，通过把通信和计算重叠起来，对吧。FSDP 基本上可以做到零开销。对，你确实会看到一些气泡（bubble）。这些是开销。但如果你的通信非常快，而你的计算量非常大，对吧，你很容易看到计算时间会比通信时间长——抱歉，计算可以比通信花更长时间，于是你最终在内存使用上得到了非常棒的改善，而几乎不用付出多少内存代价。好，所以当我们看不同的 ZeRO 阶段时，对吧，DDP，即朴素数据并行，成本是 2 倍参数量。

[25:06.800] 对，它是一次全归约（all-reduce）。两个 ZeRO 阶段，即阶段 1 和 2，对吧，也就是切分梯度和优化器状态的那些，这在字面意义上是零成本的，因为它的通信量和原来一样，对吧。这只是全归约等价于两个操作这个恒等式而已，对吧。现在，ZeRO 阶段 3 在技术上通信量稍微多一点。你需要额外的一次全收集，对吧。但也没那么糟糕，因为这些库可以用非常巧妙的方式把通信成本隐藏到计算下面。所以实际上，FSDP，如果你运行它，你会看到 GPU 利用率非常接近单 GPU 的性能。FSDP 确实好得令人惊叹。而且它在概念上也非常简单。你们在作业中实际上需要自己写一个 FSDP 实现。你只需要写一个包装器（wrapper），它可以把任何模块包装起来，变成它的 FSDP 版本。

[26:02.620] 对，从概念上讲，你要做的就是做一堆全收集、计算、释放，然后在反向传播中重复这个过程。好，很好。有问题吗？[听不清] ——需要下一个 GPU 的梯度来计算上一个 GPU 的梯度，对吗？当你说——我们不是……我觉得你刚才描述的，也就是把一个 GPU 的梯度拿到下一个 GPU，那更接近流水线并行（pipelining）。流水线并行是指一个层在一个 GPU 上，另一个层在另一个 GPU 上。但这里不是这种情况。所以在所有这些情况里，每个 GPU 都会遍历整个模型，对吧，从头到尾。但区别在于，没有任何一个 GPU 会同时持有全部参数。

[26:49.000] 所以我们要做的是，如果我只有一部分网络，没有第 0 层，那我就先请求第 0 层的参数，然后做完整的计算。然后我再继续到下一层，对吧。所以这是同样的计算。这和朴素数据并行是相同的单 GPU 计算。但你可以把它理解为在每两个计算操作之间请求参数、释放参数。好的？所以你的意思是每个 GPU 拥有每一层的一部分，而全收集就是为了这个？是的，没错。图片就在这里，对吧。参数、梯度和优化器状态都切分在不同的 GPU 上。没错。对，很好。哦，有问题是吗。为什么通信成本不会也乘以层数呢？操作次数确实乘以了层数。这是对的。但每一块通信的数据量更小了，对吧。因为这只是一小个 MLP 之类的。而我是在拿它和全网络做一次全归约的成本做比较。对，所以在某种意义上，加起来总量是差不多的。对，很好。好的。好，很好。

[27:55.080] 所以实际上，如果你看能放得下的最大模型尺寸，比如说在一块 A100 GPU 上，如果从基线开始，你连 70 亿参数的模型都放不下。而到 ZeRO 阶段 3，你现在可以放下 500 亿参数的模型等等，对吧。所以用 FSDP 来装这些大模型就变得容易很多了。对，你在内存方面得到的改善是非常显著的。好，是的。那么你能回到上一页幻灯片吗？关于这个在 8 块 A100 上的数字，那在 H100 或 P200 上会有什么限制？好，我没法心算乘以 141 再除以 80，但那就是 H200 的数字。以此类推。对，会是线性的，对吧。不过从一个加速器往上推，并不会真正改变计算的本质。好，所以 FSDP，我觉得非常干净。非常优雅。我很喜欢它。我真希望这节课到这里就能结束。但不幸的是，我们还得继续进入更难看、更棘手的内容。

[29:05.910] 我们必须这样做有两个原因。其中一个，也是并行化中一个重要的思想，就是在某种意义上，数据并行消耗了一种重要资源。而这个资源就是批大小（batch size），可以这么说，对吧。如果你的批大小是 8，你就永远不可能用超过 8 个加速器。对吧。你可能会想，好吧，那我们就随着加速器数量的增加，把批次大小做得越来越大、越来越大。但你不能这么做，因为在某个临界点上，存在一个叫做临界批次大小（critical batch size）的东西，此时增加一个批次元素所带来的收益，反而不如你针对那单个元素再执行一次 SGD 步骤来得大，对吧。

[29:43.950] 所以在某个点之前，在批次大小较小时，添加额外的元素是完全有帮助的，这等同于执行更多 SGD 步骤。但在某个临界点，收益递减效应就出现了。无限大的批次大小，并不比无限多的步骤更好，对吧。无限多的单步更新。正因为如此，我们现在面临一个艰难的权衡。我们是选择较小的批次大小，让我们的 GPU 更加空闲呢？还是选择较大的批次大小，承受优化方面的损失呢？很难说，对吧。我们还有另一个问题，那就是内存问题依然存在，ZeRO 阶段 1 和阶段 2 实际上并不能让你扩展内存。而 ZeRO 阶段 2 可以让你切分参数内存。但不幸的是，对于激活内存以及其他类型的内存使用，这种方式根本无法切分，或者说根本无法减少，对吧。所以我们需要更好的方法，更细粒度的方法，来切分模型，以便真正开始把内存使用量推到这些点以下，对吧。

[30:41.290] 这就把我们带到了各种模型并行（model parallelism）的思想上。模型并行基本上就是这样一种思想：就像 FSDP 一样，我们把参数拆分到不同的 GPU 上。但我认为，我们之前讨论的内容和接下来要讨论的内容之间的一个重要概念区别是：现在我们要通信的是激活值（activations）。对，在 FSDP 中，我们切分参数。但在某种意义上，它只是一个包装器。我们仍然在进行正常的计算。我们只是在进行计算之前把参数来回发送。对吧，所以参数是飞来飞去的东西。现在我们要做的是，改为来回通信激活值，对吧。

[31:21.510] 所以，比如说，如果一层位于 GPU 0 上，下一层位于 GPU 1 上，对吧，那我需要通信的就是这两层之间的激活值，对吧。所以这将是最大的区别。这将涉及到流水线并行（pipeline parallel），也就是切分层；张量并行（tensor parallel），也就是切分矩阵；以及专家并行（expert parallel），也就是把专家（experts）分片到不同的位置。对，Percy 已经讲过张量并行了，所以你们大概有了一些概念模型。我将从流水线并行或层间并行开始讲起。我认为这在概念上非常简单。早在我最初学习并行计算的时候，我就觉得，啊，是的，这很简单。

[32:01.090] 只需切分层，然后放到不同的 GPU 上，对吧，非常简单的事情。然后你在前向传播中向前传递激活值，接着在反向传播中向后传递部分梯度。那么，好吧，如果我这样做，会发生什么呢？你会得到一张像这样的、看起来很糟糕的图。非常非常令人沮丧。我有四个加速器，每个处理模型的四分之一。那么，计算是如何随时间变化的呢？时间从左向右流动。不同的加速器对应不同的行。首先，我的第零个加速器，也就是底部这个，它处理第一层，会做一些计算。然后它会停止，对吧。然后它把计算结果交给第二个加速器。然后它在剩下的时间里什么都不做，直到最后。

[32:44.672] 然后我基本上就是在来回传递东西。一次只有一个 GPU 处于活跃状态。然后回到反向传播时，一次也只有一个 GPU 在处理反向传播。所以在大部分时间里，我的 GPU 都是空闲的。我的利用率真的是非常非常糟糕。这被称为气泡（bubble）。这是流水线中你什么都不做的那部分。解决方案本质上就是做批处理或流水线化。而不是简单地一层接一层地执行，你会有一个流水线，对吧。在这个例子里，这张图中，你有四个需要处理的数据元素。一旦你处理完一个，你就开始处理下一个元素，然后把已经完成的那个交给你的下一层，对吧。所以你可以这样做，基本上就是把元素向上传递，然后再向下传递。

[33:33.390] 所以如果你想让这个变得合理，你需要的是这个气泡时间与有效计算时间的比率。在某种意义上，利用率将等于阶段数除以微批次数量（micro-batches）。所以我们需要很大的批次大小，才能把这个气泡减少到接近 0。对吧，所以如果我们希望气泡为 0，它会随着 1 除以微批次大小而下降。这就是为什么我说批次大小是一种有用的资源。因为现在我们可以在另一个方面使用它，那就是通过更多的流水线化来减少流水线并行中的空闲时间。好了，流水线很糟糕。而且根据经验法则，人们会说类似这样的话：并行化代码看起来还算合理，人们也能理解，直到你实现流水线并行。那我们为什么还要用流水线并行呢？嗯，与 DDP 相比，流水线可以节省内存。

[34:28.420] 因为我们已经切分了层，而且你还可以把它与数据并行（data parallel）结合起来。所以你可以两者都用。但另一个原因是，这是一个非常重要的需要学习的思想，而且在某些类型的拓扑结构中你总会用到流水线，那就是与 FSDP 以及你能做的其他任何类型的事情相比，它具有非常好的通信特性。它只依赖于激活值，也就是 b 乘以 s 乘以 h。它也是点对点（point to point）的通信。一层与另一层通信，而不是全对全（all to all）。因此，根据你如何安排这些通信的顺序，它们可以成为利用网络的一种非常高效的方式。

[35:01.980] 所以在实践中，我们要做的是，流水线将位于最慢的网络链路上。所以如果你有多个数据中心或多个计算舱（pods），它们之间的通信速度较慢，你会使用流水线并行在这两者之间进行并行化，因为这是你能在此处使用的最通信高效的并行方式。好的。正如我所说，流水线性能确实依赖于批次大小。所以我会反复提到这篇论文，但我会说明是哪一篇。这就是 NVIDIA 的 Megatron 论文。他们做了一堆非常漂亮的参数扫描，展示了在不同配置下，利用率会如何变化。所以，如果你有较大的批次大小，并且选择了较大的流水线规模，你可以获得相当不错的利用率，接近完全不使用流水线并行的水平。但是，你需要大的批次大小才能保持高利用率。否则，流水线并行会迅速降低你的利用率和性能。

[36:03.472] 实际上，我想在这里停顿一下，回答一些问题，因为下一张幻灯片比较复杂。有请？你之前在上一张幻灯片中提到，流水线有很好的通信特性，比如 FSPD。这背后的具体机制是什么，为什么会这样？因为需要通信的数据量更小。就是批次大小乘以序列长度乘以隐藏维度。而且这几乎总是比尝试通信整个参数矩阵的数据量要小，对吧。人们已经尝试找出更聪明的方法来减少流水线规模。而且你可以用一些巧妙的方法，把不同层或不同微批次的工作阶段切分开来。这样你就能在不同的反向传播阶段之间，穿插安排不同的前向传播元素。通过这种巧妙的调度，你可以进一步缩小气泡的大小。

[37:01.920] 这张图来自 DeepSeek 的论文。通过极其精细地控制前向和反向计算的时间安排，你可以在气泡大小方面获得一些改善。但我觉得更巧妙的方法，叫做“零气泡流水线”，它本质上并不完全关乎调度本身。它其实关乎的是仔细思考反向传播的结构。在反向传播中，实际上有两件事必须发生。你可以想一想，你在计算图中往回走。在这个过程中，你要做两件事。第一件事，你要把偏导数继续向计算图的更深处传播。记住，反向传播是从末端开始的，你把这些偏导数往回传播，对吧。所以你必须做这件事。但你还想计算当前权重的梯度。对，所以你在计算图的每个节点上要做两件事。你要往回传播，你还要计算权重的导数。在思考这个计算图时，重要的一点是，这两个任务中有一个对流水线来说非常关键。

[38:08.740] 而另一个你什么时候做都行。把偏导数往回传播，这件事非常重要。因为如果你的下一阶段还没有收到那个传播回来的信号，它就无法开始任何工作。另一方面，关于权重的导数，它在这个图上可以说是叶节点。所以我可以随时计算它，对吧。现在要做的巧妙之处在于，你把这两件事分开。对，你在这里有这些 B。这些 B 是反向传播，即把信号往回传。然后这些 W，这些 W 是计算权重的导数。而它们可以随时计算。所以你要做的就是尽快计算 B。对，所以这些优先。然后你可以把 W 的计算推迟到后面，在计算间隙的时候再做。当你这样做的时候，基本上可以几乎完全填满流水线。

[38:56.180] 而“零气泡流水线”这个方法非常复杂，或者比通常你愿意处理的复杂度要高得多。但它能让你几乎完全解决这些流水线问题，具体效果取决于 W 和 B 各自涉及的计算量等等。但我觉得看到这个方法之后，我的反应是，哇，这太聪明了。用系统能做出这么巧妙的事情。好了，这就是流水线并行。有什么问题吗？好。我觉得你可以想到一个自然的类比或对照。流水线并行就像是深度。也就是你的深度，你沿着深度维度把网络切开。但我们有两种扩展方式，对吧。我们有深度轴，也有宽度轴。那如果我们开始沿着宽度来切分模型呢？对，这就是张量并行，之前 Percy 解释过。

[39:50.580] 而且这确实是一个简单的想法：矩阵乘法可以被切分成更小的矩阵乘法，然后再把部分和加回来。这其实和分块（tiling）也是同一个思路。所以这个核心原语会反复出现很多次。那么在实际操作中，这样做到底是什么样的呢？假设这里是一个 GeLU 的计算。通常你会怎么做？你的输入 x 从这张幻灯片最右边进来。然后我乘以一个矩阵 A。我应用非线性函数 GeLU。然后我再乘以 B 来做降维投影。最后输出我的 z。对，这就是我通常做的串行计算。在张量并行中，我要做的是把 A 切分成两个子矩阵，把 B 也切分成两个子矩阵，然后并行运行这些计算，最后再把它们汇聚起来。

[40:45.180] 现在有一点很重要，就是前向和反向传播之间存在对偶性。在前向传播中，我对 x 做的只是把它复制两份，然后向前运行计算。所以 f 是恒等函数，就是这个 f 函数。然后当我到达并行部分的末端时，g 函数是一个全归约（all-reduce）操作，把结果重新汇聚起来，对吧。但在反向传播中，当我在计算梯度时，这个关系会反转。我的 g 会变成——抱歉，我的 g 会变成恒等操作。对吧，因为偏导数是往回传的。然后我的 f 是那些需要被累加起来的偏导数，以便得到这个模块起始处的偏导数。对，所以它们会反转。如果你要编写张量并行代码，理解这种对偶性非常重要。好了。你可能已经注意到了，对吧。当我们切分东西的时候，有些方式是按列切分。具体来说，按列的张量并行会发生在输入侧。

[41:44.990] 对，MLP 的输入、注意力机制中的投影在每个 transformer 块中都会按列切分。所以这是按列切分。然后对应的第二阶段会有按行切分。也就是在 MLP 的降维投影中，以及注意力的输出部分，你会做按行切分。然后所有那些小层，比如层归一化、非线性函数或者 MoE 中的路由器，这些都完全复制到每台机器上。你不会想去为了切分这些而承受额外的通信开销。好了。那么，什么时候用张量并行呢？张量并行对通信的需求极大。你做的事情是，每次做矩阵乘法的时候，你都需要进行一些通信。对吧，你有一个全归约操作。或者在反向传播中，这里也有全归约。前向的时候在 G 那里有一次。

[42:44.330] 而这些都是激活值大小级别的通信，而且你非常频繁地做这些操作。所以这是非常消耗通信资源的。因此你只希望在单个节点内部使用它。对于 GPU 来说，GPU 会在一个机箱内紧密地通过网络连接在一起。所以在最多八张卡的范围内，你可能愿意做一些张量并行。一旦你超过这个范围，你就要跨节点了。那些互联、那些连接速度要慢得多。所以一旦你越过单台机器的边界，性能就会大幅下降。所以张量并行，你要记住它是对通信极其敏感的。你通常会在手头最快的互联环境中使用它。说到这里，我要回到之前提到的 TPU。对，记住对于 TPU 来说，你实际上没有那种“八张快速 GPU”和“其他机器”的区别。

[43:37.470] 你只有一个大的网状结构。所以 TPU 的人会告诉你的一大优势是，与 GPU 世界相比，你可以用张量并行扩展到非常大的规模，因为他们能够在这个网格上以非常规则的通信模式实现高带宽。对吧，所以如果你在 TPU 和 GPU 之间做并行化，这是一个很大的区别。你使用的张量并行程度与流水线并行的比例，实际上会有很大的不同。那我们怎么比较这两种方法呢？它们都是模型切分策略。这两者在参数和可能的激活值上都能节省。张量并行的优点在于，不一定会有流水线气泡。如果你的网络足够快，你或许能实现完全利用。它复杂度也低，不过是拆分一些矩阵乘法而已。

[44:25.410] 但缺点在于通信量要大得多。以前你只需要做 b 乘以 s 乘以 h 的点对点通信。而现在，我要做的则是类似 a 乘以 b 乘以 s 乘以 h，再大约乘上一次全归约（all-reduce）通信量，对吧。所以这甚至不只是点对点，而是需要发生的全局通信。因此，每当我们有高速互连时，张量并行是极好的；在其他任何情况下，我们可能更希望使用流水线并行。好的，关于内存使用，我想讨论的最后一点，是更仔细地去思考内存本身。我认为在深度学习和语言建模中，对内存最朴素的观点大概是这样的：内存就是参数而已。就是这里的绿色部分。对吧，你需要一定数量的比特来存储你的参数，仅此而已。我已经讨论过优化器状态，这很重要，就是这边的黄色方框。

[45:23.713] 但如果你看这张图，这是在一个 PyTorch 分析器中记录的实际内存使用情况。你会发现，这其实并不是全部的故事。在计算过程中还需要存在大量动态内存。你需要存储所有的激活值。就是这些大的红色凸起部分。然后，当你计算梯度时，你需要在反向传播过程中存储它们，对吧。所以真正的内存使用峰值发生在最大激活点之后一点。当你开始反向传播计算梯度时，你仍然需要保留大量的激活值，那些通常才是内存的最大占用点，对吧。所以为了减少这种情况，也许我们需要处理这个巨大的红色驼峰，也就是激活值。如果你开始关注更多现代的语言模型工作负载，你会发现，随着参数规模越来越大——先别管那一列写着“当前工作（present work）”的。

[46:17.950] 我稍后会讲到这个。但如果你只看这张图的基础列，随着模型变得更大，对于中等长度的序列，你的激活值内存将远远超过你所需的参数内存。因此，任何内存节省策略都必须考虑激活值，才能完全有效。我们不能只处理参数内存，那对我们来说确实不够。所以，如果我们要存储所有东西——这是一个很好的，我也不确定，可以作为起点的统计数据或经验法则。如果我们需要存储所有内容，我们需要的激活值量大约是 34 乘以 s 乘以 b 乘以 h，即序列长度乘以批大小再乘以隐藏层维度大小，再加上 5 乘以 as 除以 h。as 是注意力头数乘以序列长度，再除以隐藏层维度大小。而这个 sbh 的依赖关系是根本性的，因为我们预期要为序列中的每个元素存储一些东西。

[47:21.833] 我们预期要为每个批次元素存储一些东西，我们也预期要为每个隐藏维度大小存储一些东西。对吧，所以当我们尝试不同的减少激活内存、计算总激活内存量的方法时，我们会反复看到这个 sbh 项。如果你好奇这些项是从哪里来的，这个看起来奇怪的项，5 as 除以 h，来自二次的注意力项，包括 dropout 项。如果我们使用 Flash Attention，我们可以通过重计算来去掉这一项。好的，那如果我们做张量并行呢？我在谈论模型并行时向你们承诺的其中一点是，我们可以相当多地减少激活值。那么现在让我们想想会发生什么。嗯，张量并行会拆分注意力和 MLP 中的矩阵乘法。那么这些是哪些部分呢？

[48:11.760] 嗯，MLP 占了之前我提到的 34 项中的 24 项。而注意力头是 5 as 除以 ht。对吧，就是这里的第二项。现在，我可以将其除以我的张量并行规模 t 来减少它。这很好。这并不坏。因为我已经占据了大部分内存，并使其随着张量并行规模 t 线性扩展。所以这很好。但事实证明，你还需要做其他事情。你可能有一个层归一化（layer norm）。你可能有一个 dropout。你需要保留注意力和 MLP 的输入。对吧，因为每一层的输入都需要作为残差存储起来，以供反向传播使用。不幸的是，这些并不会随着你的张量并行规模而减少。对吧，张量并行并不拆分层归一化。因此，它也不会拆分这些部分的激活值。所以这很不幸，因为你原本希望，如果你是谷歌，有 1000 个不同的 GPU 做张量并行，你可能能够大幅减少这些激活值。

[49:12.780] 不幸的现实是，做所有这些，你仍然要承受 10 倍 sbh 的代价。因此，通常与张量并行组合使用的最后一部分，是这个称为序列并行（sequence parallel）的想法。这是一个极具误导性的名称，因为通常当我说序列并行时，实际上有一个叫做上下文并行（context parallel）的东西，它对这个名字来说更自然。但序列并行是以下这个想法。所以你要做的是，我们还有这些剩余的项，10 sbh。而这些项是非常轻量级的。它们就像层归一化，是计算量不大的东西。而我要做的是，我要以某种方式把它们拆分。我要沿着序列轴而不是隐藏轴来拆分它们。这将会在每次需要这些项的操作之前，涉及到全收集（all-gather）和归约分散（reduce-scatter）。如果你想一想，这非常让人联想到 FSDP。

[50:13.475] 对吧，我们有这个需要的东西。我们需要注意力——抱歉，我们在某个时刻需要激活值。我们需要在前向过程中计算激活值，并且在反向传播求导数时也需要这些激活值，对吧。但我们不是现在就需要它们。所以我们要通过沿着序列轴拆分它们来存储，并按需将它们具体化。所以我会说，这在概念上与 FSDP 风格的想法非常相似，我们以这种分片格式存储它们，然后在需要时收集它们。然后正如我之前所说，这里存在一个对偶性。所以在正向传播中，g 是全收集（all-gather）。g bar 是归约分散（reduce-scatter）。在反向传播中，这将会反转。对吧，全收集和归约分散在 g 和 g bar 之间将会互换。

[50:57.600] 所以，如果我现在这样做，我会得到什么？嗯，记住，在无语并行的情况下，最开始的部分，我告诉过你，你会得到这个激活内存。使用 34 乘以 sbh，并带有注意力头分量。使用张量并行，我可以将 MLP 部分和注意力部分除以 t，但所有这些其他逐点运算则不行。如果我做序列并行，我也已经通过加速器拆分这些，只是沿着序列轴，这让我现在有了完全的线性依赖。如果我做激活重计算，我甚至可以去掉这第二项。因为记住，这只是我为 softmax 需要的存储。所以我可以通重计算来降低这部分开销，得到 sbh 乘以 34 除以 t。记住这个值很有用，因为这是在合理情况下，常规训练中激活内存所能达到的下限。所以如果你想手动估算模型能否装进一块 GPU，记住这个值会很有帮助。

[51:59.220] 激活内存至少需要这么多。优化器状态还需要一些内存，参数和梯度也需要一些内存。一旦你有了这些数字，你就大概能估算出你的模型需要多少内存了。好的，我在这里暂停一下。到目前为止，我已经讲了流水线并行和张量并行，大家可能有点困惑。在我继续讲专家并行之前，我们可以先讨论几个问题。好的，嗯？[听不清] 24 除以 t 是 MLP 部分。MLP 部分不能通过重表示来优化吗？MLP 也可以重计算。你可以做的计算量比这里列出的多得多。但这里列出的重计算在计算上非常昂贵，比如对 MLP 做重计算意味着在反向传播中要再次运行 MLP，这通常是你不想做的。

[53:04.880] 注意力的重计算通常更便宜，是的。因为它在时间上也是可行的。而且你也不想付出二次方的代价。那相当昂贵。二次方的代价，确实很痛苦。好的，行。那么我现在要讲标准并行工具包中的最后一个组件，也就是专家并行，对吧。如今 MoE 已经成为工具包中非常标准的一部分，现在大多数大模型都是 MoE。这让我们能够利用它们的一大优势，也就是专家并行。而专家并行，我认为它和张量并行有些类似。对，从某种意义上说，你是把 MLP，也就是网络中的 FFN 组件拿过来。你要用各种方式把它们拆分开来，你会产生通信开销，就像你在拆分矩阵时一样，只不过这里你是在把 MLP 分布到不同的设备上。

[54:07.870] 而专家并行的系统行为——哦，抱歉，专家并行——在某种意义上和张量并行也差不多，因为它是一种高带宽的并行原语，同样也能减少激活内存。但有一点很重要，如果你在使用 MoE，你几乎总是会优先选择专家并行切分策略，而不是张量并行。我截了 Megatron 的并行指南中的一张图，Megatron 是 NVIDIA 的并行库。这是一个非常出色的库和文档。其中有一条指南说，如果你要做某种并行，无论是 EP 还是 TP，你都应该选择 EP 而不是 TP。为什么呢？嗯，这其实很好地列出了张量并行的缺点。如果你有一个矩阵，然后把矩阵切得太细，矩阵就会变小，GPU 利用率就会下降。所以你希望矩阵乘法尽可能大，而张量并行会降低这一点。

[55:10.070] 对于 MoE 层来说，路由稀疏的 token 激活要比可能路由这些稠密的大张量并行矩阵乘法激活容易得多。而且你可以做到将 token 精确地路由到它们需要去的地方。如果你本来就要用 MoE，你还可以跳过一些计算——或者说，抱歉，跳过一些开销，不如直接把专家数量分布到所有设备上。这——抱歉，是的。这可能会让人觉得，好吧，EP 很棒。每个人都应该一直用专家并行。不幸的是，专家并行仍然非常复杂，仍然非常困难。我觉得我在 FSDP 之后讲的所有内容——张量并行、流水线并行——从系统角度来看，要真正做好都是相当不容易的。我只想提两个不同的库，不是为了告诉你们它们内部发生了什么细节，而是想告诉你们，这真的是非常复杂的工程。

[56:16.110] DeepSeek 有一个叫做 DPP 的库，我认为这是 DeepSeek V3 时代我最喜欢的东西之一，这基本上是他们自己做专家并行路由和分发的一个库，他们在研究非常底层的 GPU 网络原语，试图基本上把不同的操作合并起来，或者让硬件来处理某些操作，从而让 MoE 路由尽可能高效。类似地，NVIDIA 有一个叫 Hybrid EP 的库，也是在做一个类似的工作，提供非常高效的底层硬件实现来进行专家并行的分发。简单来说，对吧，这些东西之所以难，是因为你在做大量的全对全分发。如果你看 MoE 的通信模式，和张量并行很像，每次你有一个 MLP，你都得把 token 路由到各种不同的地方，而且这需要以非常低延迟的方式完成。因为你的计算在等待 token 到达。

[57:20.070] 所以降低这种分发的延迟极其、极其重要。最后一件关于 DPP 库的趣事是，深度求索的工程师们太拼了，为了优化每一丝性能，他们基本上是去找了未公开的 PTX 指令——就是类似 GPU 机器码或机器码级别的东西。他们找到了如何使用这些未公开的指令来进一步加速网络通信。所以，真的，要到达并行效率的前沿，你就需要做到这个程度。最后，到目前为止，我讲的都是孤立的各种并行原语，比如数据并行等等。隐含的假设是你像搭乐高积木一样把它们组合在一起就行。这在大多数情况下对基本上所有其他并行策略都是成立的。比如数据并行和张量并行就很容易组合。但专家并行有一些额外的约束，你需要在一定程度上加以注意。

[58:28.110] 通常，数据和专家并行的朴素实现方式——我认为很多老库都是这么做的——是 DP 和 EP 的副本是相同的。换句话说，假设我有一个八路数据并行。我要把专家切分到那八个副本上。然后我的数据会被切分到这八个副本上，以各种方式路由。对吧，换句话说，你根据 DP 来切分 GPU，然后你的 EP 是这个切分的一个子集。对吧，这是很自然的事情，因为你在路由 token。但如果这样做，EP 的并行度有一个上限。它会约束 DP 和 EP 的方式。所以我觉得这是最后的一个复杂性。我很抱歉专家并行在很多方面确实有点混乱。但正因为这些复杂性，在我结束讲 EP 之前，其实有一个有趣的系统层面的东西要告诉你们，

[59:30.900] 就是下面这个事实。记住，MoE 只改变 MLP，对吧。它们完全不改变注意力。对吧，因为我们在切分 MLP，但没有切分注意力，这意味着专家并行在模型上是不均匀地应用的。所以当我在专家层面做并行时，我在并行化 MLP，但我并没有并行化注意力，对吧。所以为了并行化注意力机制，我可能也想使用张量并行，对吧。所以我需要较高的张量并行级别来切分我的注意力层。但这就会影响我的MLP。因为如果我的张量并行和专家并行级别都很高，我的矩阵就被切成了非常小的碎片，利用率会非常糟糕。所以这里存在一个奇怪的权衡：注意力层需要高张量并行，但MLP层却需要低张量并行。所以这就产生了冲突。那该怎么办呢？嗯，在过去几年里，人们提出了更复杂的系统解决方案，将MLP和注意力中的张量并行解耦。

[01:00:39.440] 于是注意力层使用一种张量并行，而MoE层使用另一种张量并行。这样你现在就可以完全解耦并行策略了。这也带来了EP、张量并行以及数据并行之间更复杂但更有效的组合。好的。那么在这个领域里，最后我想讲的是上下文并行。上下文并行，或称环形注意力，其基本思想是将一个非常长的序列中的激活值拆分到不同的加速器上。你可以按照TPU的网格拓扑结构，以环形的方式将激活值传递给需要的设备。环形注意力——我认为这是最早提出这一方法的论文——证明了这种方法在TPU上效果非常好。上下文并行是执行这种操作的标准化并行策略。这两种方法都用于长上下文扩展阶段和模型服务中。

[01:01:44.740] 我不打算详细展开讲，主要是因为我觉得它在概念上与我们已经讨论过的很多内容有重叠。好的，这就是第二部分内容的结尾。为了把所有内容整合起来，我做了一个并行策略的大表格，包括最上面两行的数据并行，以及下面各种模型并行策略。我用红色标出的是我主观上认为的各种方法的缺点。我之所以用红色标记，是想让大家明白，没有哪一种并行策略是绝对占优的。全都是各种各样的权衡，你需要以某种方式妥善处理这些权衡，才能得到好的结果。比如，FSDP很好，它是一个很棒的策略。但它不能帮助你处理激活值的问题。它需要消耗你的全局批量大小来进行并行化。所以它是有局限性的。为了解决这些问题，你也许需要将它与你手头的其他技巧组合起来使用。

[01:02:46.892] 也许你想使用一些张量并行来减少激活值的内存占用。它也不影响全局批量大小。但如果你这样做，就需要快速的网络连接和高带宽。而且，为了利用你较慢的连接，你可能需要使用其他一些并行化策略。对，你们看，这些不同的优势意味着，这些并行化策略并非适用于所有架构，但在许多大规模架构中都有用武之地，你会大量组合使用这些策略。有问题吗？[听不清] 它们仍然适用于MoE吗，还是说那严格来说是——是的。所以问题是，流水线并行和，呃，FSDP适用于MoE吗？是的，是的，它们绝对适用。如果非要说的话，标准的建议是专家并行就像张量并行一样。所以你倾向于用到8或者快速连接。人们现在其实不再严格遵守这个了。

[01:03:48.620] 但即便如此，FSDP或一般的数据并行，还有流水线并行和张量并行，在所有前沿模型中都被广泛使用。是的。好的。好的，你们将在作业中做这件事或类似的事情。但关于并行化，有一点很酷也很好，那就是你可以做一些数学计算，对吧。你实际上可以说，好的，对于我的不同切分策略，我每层要做多少计算？对于我的不同通信策略，我每层要做多少通信？如果我将不同的设备分配给不同的通信策略，这些指标会如何扩展？然后我可以弄清楚，对于这些策略的组合，它们是如何扩展的？你可以开始绘制图表，比如，如果我有某个批量大小和某个并行化策略——我这里展示了三种不同的并行化策略——我的GPU利用率是多少？即计算时间与通信时间的比率。只要我的计算时间比通信时间长，原则上你就可以把通信隐藏在你的计算之下。所以你可以做出这样的图表，最上面的那条虚线，就是你效率高的边界。

[01:05:04.520] 高于这条线意味着完全利用了你的计算资源。低于这条线则处于你Roofline模型中表现不佳的部分，你在等待通信完成以便进行计算。这里有一件非常有趣的事情。我觉得基于我告诉你们的这些，你们应该会直观地感受到这一点。如果你的批量大小足够大，那么只用FSDP就很好。对吧，如果你的批量大小是每芯片2000，那你最多能扩展到多远？嗯，FSDP表现极其出色。你完全受计算限制，对吧。但重要的是，随着批量大小的减小，纯FSDP会达到一个点，此时你受到通信限制。那现在该怎么办呢，对吧？嗯，现在你要做的就是为了持续提升效率，必须将MP——在这里指的是张量并行——加入到组合中。对吧，现在你可以把那条曲线再往外推一点，即使是在更小的批量范围内也能受计算限制，对吧。

[01:05:59.310] 你不断添加策略来推动这条曲线外扩，以获得更好的结果。这就是人们所说的3D或4D——我不认为有5D——并行。也就是说，你把所有这些策略结合起来，以便在不同的通信拓扑下保持你的计算单元得到充分利用。其实有一个非常简单的方案。我认为到目前为止，我一直在比较笼统地讲述。所以可能感觉有点复杂。但实际上，并行化的策略是非常简单的。那就是，只要你的模型还不能完全放进内存，你就需要想尽一切办法去切分你的模型，对吧。你会怎么切分你的模型呢？对于快速互联，你会使用张量并行或者专家并行。所以，如果你每台机器有8个GPU，你就会使用张量/专家并行度为8的设置。然后你会使用流水线并行或ZeRO-3（也就是FSDP）来处理剩下的部分，以确保模型能够放进去。

[01:06:53.070] 一旦你的模型能放进了，接下来就全部使用数据并行。对吧，就用数据并行来利用剩余的GPU。如果你的批量大小最终太小，你可以通过梯度累积来获得更好的GPU利用率。很简单的策略。这也是你可能会用到的大规模并行训练库所实际采用的方法。我在这里添加了Megatron的MoE并行化指南的链接。但这个想法实际上也适用于稠密模型。你会看到我所说的完全相同的内容，只是顺序相反。指导原则第一条是：最小化模型并行，最大化数据并行。对吧，如果可以的话，用你大部分GPU来切分数据。现在，如果你使用的是GPU，要保持在NVLink范围内。也就是一台机器，或者说一个机箱内，对吧。将专家并行和张量并行限制在一个机箱内。

[01:07:44.190] 然后，如果你要跨节点训练，就用流水线并行来处理多节点。如果你用的是 MoE，那就优先考虑专家并行。最后，如果你在处理长序列，就用上下文并行。对，就是环形注意力。这正是我整节课想传达给大家的内容，只不过这里呈现的是实践者的视角。好的，最后一个部分——其实，或许我可以在这里停一下，以防有人对四维并行化感到好奇。嗯。我知道你说过上下文并行对于减少并行来说是个更好的名字[听不清]。哦，对对对。嗯，我的意思是，序列并行通常是你和张量并行一起使用的，目的是减少激活值。它更多是一种附加手段，用来减少激活值。在很多情况下，它并不是一个独立的东西。是吗？有人猜测这就像循环 Transformer，那样的话，会让事情变得更容易还是更困难？好吧。我不太确定我有多相信循环 Transformer 的那些传闻。

[01:08:57.430] 也许是真的。那会很有意思。但关于系统结构会如何改变这个问题，问得好。我的意思是，我猜它可能会打乱像 FSDP 这样的东西。你确实不能把权重丢弃掉。不过它的优点在于，你不必反复获取权重。它在参数效率上要高得多，因此也许很多模型并行的工作就不会那么重要了。那会很有趣。对于很多 FSDP 风格的东西来说，肯定会非常不同。因为 FSDP 策略本质上依赖于获取权重、丢弃权重、获取权重、丢弃权重这样的循环。而在循环结构里，你什么都不用丢弃。好，行。好的。所以这是一篇比较老的论文，但它实际上可能是最好的资源之一。而且，网络基础知识并没有改变太多。所以我认为这里的经验教训在今天依然非常适用。

[01:09:53.050] 这是 NVIDIA 的一篇论文，还有几位很不幸曾经是斯坦福的人，Matei 和 Deepak，他们做了大量大规模并行训练的实验，跨越了许多不同的配置，基本上是为了告诉你，当你扩展规模的时候，你的并行化策略应该随着不同规模如何变化。你看到的就是我之前一直在讲的这套方案，也就是说，你先把数据并行推到极致。然后你的张量并行增加到 8，之后你就停下来，因为你不希望再高了。从那个点开始，你的流水线并行不断增加、不断增加、不断增加，对吧。而在足够大的规模下，数据并行会下降。在最后，DP 是 6。这是因为你确实需要这么多张量并行和流水线并行才能装下这些东西。对吧，然后你把剩下的预算用在数据并行上。

[01:10:47.130] 如果你把它们全部放在一起，即使你去到极其庞大的 GPU 数量，你的利用率也会保持得非常平稳、非常好。这就是为什么我认为你会看到这些巨大的数据中心建设。人们想出来的并行化策略实在太有效了。而且通信硬件也非常好，所以即使在巨大的数据中心，甚至是跨数据中心训练，你也能在这些不同的模式下获得非常好的利用率。他们有非常漂亮的量化证据，向你展示对于不同的批次大小，张量并行应该在什么地方停下来。很明显是张量并行大小为 8 的时候。一旦超过这个值，你就会有麻烦。而流水线并行，你也需要很大的批次大小才能有效利用机器。这就是为什么蓝色和橙色之间有一个差距。那是流水线并行。

[01:11:44.700] 流水线并行越大，橙色和蓝色之间的差距就越大。然后我觉得另一件你可能已经从 Flash Attention 中内化但也许还没有的事情是，也许你应该做大量的激活重计算。因为如果你非常巧妙地做激活重计算，那么你就能获得更大的批次大小，而更大的批次大小反过来又能让你获得更好的利用率。这最初是一个非常反直觉的观察，就是实际上你应该通过重计算来做更多的计算，以便最终获得更好的利用率，因为它让你节省内存，而内存可以转化为批次大小。好了。以上就是关于如何选择并行化策略的一个更量化的视角。最后我想讲的是，我们来看几个非常少的训练实例。通过观察这些训练实例，也许我们能了解现代并行化在实践中是如何被使用的，以及它如何演变。

[01:12:46.720] 所以 AI2 的 Dolma 是一个 7B 的开源模型。抱歉，是 Olmo，它是在 Dolma 数据集上训练的。这不是讲细节的 Dolma 论文。它是一个 7B 模型，完全使用 FSDP 在大量的加速器上训练。我记不清确切是多少个了。但这是一个 FSDP 出人意料地能扩展到非常多 GPU 的例子。只要你训练的是小模型，FSDP 实际上是一个非常好的并行化策略。我认为很多 7B 级别的模型纯粹是用 FSDP 训练的。DeepSeek 由于他们的系统工程能力，确实做了一些相当独特且有用的东西。

[01:13:31.560] 所以 DeepSeek V1 只是用数据并行加 ZeRO 阶段 1，再加上张量、序列和流水线并行训练的，我之前跟你们讲过。DeepSeek V3 是一个 MoE 模型，如果你们还记得我的 MoE 讲座的话。所以它们不仅有流水线并行，而且用专家并行代替了张量并行，对吧。而且他们的专家并行实际上有点独特。因为他们有 64 路并行。他们把 8 台不同的机器分到一组，这就是他们的专家并行域。为了实现这种大规模的专家并行，他们基本上使用了从流水线并行中得来的那些流水线技巧，以确保专家并行不会出现利用率低下的时期。所以这实际上是一个相当复杂的技巧，他们用这个技巧来实现非常大的专家并行切分。另一个例子，Yi，来自中国的开源权重训练环境。他们又一次使用了 ZeRO 阶段 1 和张量并行加流水线并行。

[01:14:32.480] 对，这非常合理。这是经典的数据并行、张量并行、流水线并行组合。而且又一次，一旦你转向 MoE，你就用专家并行取代张量并行。对，它们服务于相似的目标。但专家并行只是更高效一点点。在真正大型稠密模型的另一个极端，你有 Llama 3。Llama 3 的技术报告，我相当喜欢，因为它们是少数几份真正完整拆解了它们并行化策略的报告之一。不只是摘要里简单提一下，而是覆盖了所有不同的阶段。所以你会看到一个很酷的现象，它们有一个预预训练阶段，一个在刚开始的小热身阶段。第一行可以忽略。然后中间这一行是它们主要的预训练阶段。这是一个非常标准的并行化策略。

[01:15:21.140] 张量并行为 8，一个上下文并行，16 个流水线并行，以及 128 个数据并行。很大的数据并行切分。然后在最后，因为他们想要一个长上下文模型，他们做了长上下文扩展。为此，他们提高了上下文并行，降低了数据并行。而这就是他们如何为长上下文扩展并行化这个极其消耗内存的阶段。注意，我们这里用的是张量并行，而不是专家并行。因为Llama3 405B是一个巨大的稠密模型，不是MoE模型。顺带提一句，有些内容我不细讲，但如果你将来成为这些公司的基础设施工程师，你会发现GPU经常故障。据说在Llama3 405B训练期间，GPU故障了148次。因此你不仅需要快速的并行策略，还需要冗余机制来应对训练过程中可能发生的各种糟糕情况。

[01:16:21.135] 所以这里也涉及分布式系统的挑战。Gamma 2，谷歌的一个开源模型，使用了FSDP加张量并行和序列并行，而且只用这两种策略。我之所以提到Gamma 2，是因为据我理解，这基本上实现了谷歌的一个主张：对于TPU，你真的不需要做流水线并行。你只需要拿一个非常大的环面网格，然后在这个大网格网络上做张量并行。不清楚这种方法能否无限扩展下去。这一点对我来说一直不太明确。但至少在Gamma的规模上，确实如此。最后，我查看了今年训练的一些模型，比如Mistral 8x22B。在调查过程中，我发现的一件事是，如果你对实际中的模型训练和并行配置感兴趣，NVIDIA有一个名为Megatron Bridge的仓库，他们在那里发布了很多针对不同模型规模和模型设置的推荐

[01:17:25.590] 训练配置。你可以看看他们为8x22B MoE或DeepSeek V3建模推荐的正确模型配置是什么。对于Mixtral，他们采用专家并行为8、流水线并行为4、外加张量并行为4。这个配置是针对注意力层的。你会看到这遵循了建议：让专家并行大概保持在8左右。Nemotron 3 Super沿用了DeepSeek V3的模型，有大量的专家并行以及上下文并行，因为他们至少在这篇文章里是做扩展训练的。然后我想讲的最后一个是Qwen 3，它采用了DeepSeek的方案，即相当大规模的专家并行，数量为32。他们还有8路流水线并行和2路张量并行来拆分诸如注意力矩阵之类的东西。

[01:18:23.790] 实际上很有趣的是，NVIDIA的人还针对不同模型的并行化做了各种基准测试。你会看到，我不会详细讲这是什么。如果你想看的话，这在NVIDIA Bridge仓库里。即使只在张量并行内部，也有各种不同的配置可供选择，这些配置会显著影响性能。所以系统集成的重要性相当高。那么为了把所有这些整合起来，我决定也许可以做一个总览。真正共同的一点是，尽可能多地，所有模型都使用数据并行。它们尽可能最大化数据并行的维度。张量并行几乎总是保持在8以下。而专家并行现在有时可以做得很大。我认为这部分是因为DeepSeek V3以及他们为大规模专家并行训练构建的大量基础设施。好的。那么把所有内容整合起来，这堂课最高层次的观点是，我们需要考虑多GPU、多节点，甚至多数据中心的并行。

[01:19:31.750] 要在这种场景下思考效率，我们需要真正用上我们拥有的所有方法。因为你有快速链路、有慢速链路、有利用批次大小和各种不同资源的技术，而所有这些资源你都希望尽可能充分利用。但考虑到所有这些，事实证明，组合这些不同形式的并行有一些相当简单的经验法则，最终能实现对计算硬件的高效充分利用。太好了。下周，我想我们要讲的是缩放定律。

CS336 2026 · Lecture 8 · Parallelism

视频：官方视频
字幕：en-US · manual captions
处理：连续字幕合并为自然段；每段保留首条字幕的起始时间戳
状态：原语言字幕稿，待校对

正文

[00:05.240] So last lecture, Percy covered some of the underlying mechanics, let's call it, of parallelism. Today, I'm going to do my favorite thing of talk about all sorts of things you can know, like knowledge, details, and trivia about how modern parallelism for language model training works. And what we're going to try to do is really get to grips with the many, many complexities that are involved when you are having to train huge models on huge clusters. It's going to turn out that there are many different parallelization strategies that you can use at once. At the end, we will get to a slide where I say 4D parallelism, because there's four different things that you can do at once. Actually, even more that I'll talk about. And it turns out that to parallelize at the largest scale, you'll need to use all of them or most of them.

[00:56.720] And we'll discuss why. Part of your assignment will be figuring out, given a particular network topology and a particular model, what is the optimal parallelization strategy for that? And then finally, I'll walk through some examples of recent large scale training runs with details that are published and show you this stuff in action. OK, so the very first part, I'll go through somewhat quickly. It's partially a review of basic collectives, basic networking primitives, right. But the reason why we need all of this parallelism stuff is because there's going to be two bottlenecks that we'll solve by going to multiple machines and multiple GPUs. So one of it is compute. We need more compute than we can bring to bear on a single chip. And the way we solve that is just to have many, many machines and link them up together. Right, the fastest supercomputers in the world have exaflops of compute, whereas a single GPU is nowhere near that number.

[01:55.080] The other reason why we need to continue to scale in parallel is because of memory, right. Models are big. We can't fit them all into one GPU, so we need to shard them into smaller pieces somehow. So this is going to lead to complexities of having many GPUs across many different machines. And an important conceptual object that we will deal with throughout our lecture today is the difference between intranode parallelism, where your connections are very fast, and so you can do things that are very communication expensive. And you're going to have slower internode communication. And because these are slow, we will have to use processes that respect the communication constraints of our channel more, right. So we're going to have to do both of those. And remember all of the things that we will do today. We're not actually going to be sending packets or anything like that.

[02:44.860] We're going to be dealing with this at the level of collective communication primitives. Right, so whenever we do accounting for how much communication is done, I might say, OK, we're going to have to do all-reduce or we're going to have to do allgather to implement some of these operations. So in order to implement some of these efficiently, you're going to have to go all the way to the hardware level. But our discussion today is going to be algorithmic. And so all of the networking primitives I'll talk about is going to be at this level of a collective. And remember, Percy talked about this, but this is important for-- very important for one part of the talk, which is going to be that the equivalence between all-reduce and reduce-scatter plus all-gather. Right, so there's going to be one algorithm where on one hand, you do an all-reduce, and on the other hand, you do two of these different steps.

[03:31.180] And the fact that these two have the same cost in some sense means that we can do the algorithm on the right for free. We'll see that again later. This is just a quick refresher just to make sure you remember what Percy said. OK, so that's it for the really basic review. The other thing I want to touch on briefly before we get into the algorithms, which are going to be hardware and substrate agnostic, is the difference-- oh, sorry, yes? Is there any particular reason why you are saying that old [INAUDIBLE] reduce scatter and-- OK, the question was, is there any reason why I'm talking about this specifically? We will-- [INAUDIBLE] Yes, yes, there are other decompositions. This one in particular will be useful for a algorithm, which is why I'm emphasizing this particular decomposition. That's right. There's others that you could do, of course. Yes. OK. So the other thing I will talk about before I start-- most of our discussion today will be hardware agnostic.

[04:27.810] For the most part, once you understand the basic concepts, you'll be able to apply it to any hardware substrate. But I do want to briefly before we start, just talk about the low level hardware stack. Both because it's interesting. But also because I think it leads to different parallelization strategies across different model training companies. So remember the TPU. Right, TPU is Google's accelerator. And I mentioned during the GPU lecture, a TPU is like a lightweight GPU. But what's different is the networking. So how is the networking different? Well, in one slide, the TPU network is what's called a toroidal mesh. Right, so you can think about a grid and then allowing the ends of the grid to communicate with each other. It's a little bit more complicated. It's a 3D picture. Which, Marcel, you have the visualization, right? Can you post that in the lecture channel? Yeah.

[05:21.870] So there's a fun web visualization that you can look at that visualizes the toroidal mesh in 3D, which is very cool. But your mental model, I think, can roughly be this picture. You can forget the more complex one. All the chips are networked to neighbors and the neighbors wrap around, right. So what does this let you do? This lets you have a very simple networking topology that can be indefinitely scaled up very simply, right. The number of neighbors you have is the same, no matter how large your network gets. In contrast, the GPU networking philosophy is very different. It's really much more of an all to all philosophy. And the way that it's networked is it's like a tree. It's often called a fat tree. You have your GPUs connected very, very quickly at the lowest layers. And then you have groups of GPUs that are connected in a pod.

[06:05.850] And then the pods might have spine switches that allow them to communicate to each other. And this means that as the number of nodes grows, the tree gets bigger and bigger. Your communication costs is going to get higher, or your communication topology will get more complicated, because you're connecting all these all to all as you go. Now, this difference in philosophy means that if you're communicating only to neighbors, TPUs are great. They're cost effective. And you can make each connection more beefy, right, with the same power consumption. But if you're connecting in very random ways, very stochastic, unpredictable ways, then this GPU networking will be more flexible, right. So there's this very different design philosophy that you see between TPUs and GPUs. And I'll try to mention where this shows up as we go through the parallelism lecture. And I posted this into the lecture chat last lecture. Because someone asked, what's the difference between these two and so on?

[07:02.310] Right, but just to back up my point, Bill Daley and Jeff Dean were talking about the different hardware and what are they good for, right. GPUs are good for things like mixture of experts, where the communication is much less clear because tokens route to different experts. TPUs are great if you have this dense model that you're doing very predictable partitions of. I'll talk about that in an idea called tensor parallelism later in this lecture. And so these hardware things are good for different workloads. And this was where I was going to leave it for the lecture until this morning, when Google decided that they were going to announce TPU8i and t this morning. And then I looked at this as I was preparing my lecture, and I said, oh my god, TPu8i? That's a tree topology. They've switched to a much more all to all connection. And if you think about it, it makes sense.

[07:54.630] Modern day language models are MoEs. And if you're going to serve a MoE, you're going to have all sorts of bandwidth just flying around as you route tokens to different experts. You probably do want much more of an all to all connectivity where, at inference time, these communications are real bottlenecks. You want to deal with them. And even at training time, the TPU8t, which is a training chip, the cross rack connectivity looks a lot more like a GPU. There's much more like an all to all style switch that's being built out called the Virgo network. That's a higher level networking stack. So I think there's really interesting things happening. It's a little bit of a convergent evolution across both TPUs and GPUs, where really the workloads are defining the network that we need to have. OK. And then the final thing about hardware-- I just love talking about what's happening in the Chinese AI scene-- is this idea of-- I think in the GPU lecture, someone

[08:55.670] asked, if SRAM is so good, why don't you just all use SRAM? Only SRAM chip. And that's Grok, right? It works really well. Now, some of you might be thinking, if all to all connectivity is good, why don't you just connect everything and connect them with really fast fiber optic chips, right? And then you would get the Huawei Ascend 910, right? So if you look at the specs of the Huawei Ascend, it's actually really interesting, because each chip is a lot worse than a H200. It's a lot slower at matmuls. It's slower in many ways. But then, if you're willing to dedicate a giant rack of fiber optic switches to connect a much larger number of chips, like in their case, 384 in a big rack, then maybe you can deal with the fact that each chip is slower by just connecting a much, much larger number of these guys, right. And then so I think what this shows us is a really interesting tradeoff. Because once you look at this architecture, you realize the power consumption of this thing is insanely high.

[09:58.440] It's four times that of the equivalent NVIDIA system. And so it's something like, well, if you're willing to pay the power cost, you can solve a lot of communication problems brute force, and you can scale out much more aggressively to 300 chips. But you're paying a power cost, right. It's similar to that SRAM story of, if you want to design a hardware that is efficient to use both power and manufacturing wise, you end up in a certain place. If you want to brute force things, you end up in a very different place. Right, so this is, I think, just a very interesting place to learn from about these hardware design tradeoffs. OK, good. So now, just to put this all together, the first thing you should think about is the new unit of compute is not the GPU. It's the entire data center. And we want to have several things. We want to be able to control the amount of memory we use. We want to be able to control the amount of compute that we can bring to bear.

[10:51.680] And we want all of these to be lossless. Right, we want to use all of the resources that we have. And then we're going to talk about the algorithms that will enable this through all these communication primitives components. OK, good. So now the second part of this lecture is going to be the bulk of the content. And these are the algorithms that we will use to paralyze language models. And there's going to be a couple of different important ideas. Actually, we should probably maybe even be grouped. Data parallel is one set of ideas for essentially taking your data that you're going to train on and distributing them across GPUs. So the data is the ones that move around. And then model parallelism, you're going to chop up different parts of the model and split them around. Turns out that the boundary between these two is a little leaky, because one of the algorithms here will actually cut parameters up, too. But we will see why I've put all of these algorithms into these different bins, right. And conceptually, we're trying to both scale out compute and memory, and that's why we will need a whole bunch of these different tricks. Cool. All right. So data parallelism, I think, is the easiest to understand.

[12:04.360] It is conceptually very simple, and I think it's the standard way that you would try to parallelize an algorithm like LLM training. So forget that we're doing Adam. We're going to do just naive SGD for the moment. Then we have a capital B sized batch, right. So each update step, I'm going to take capital B elements, and I'm going to sum up the gradients. And I'll take my update, right. So the most naive form of parallelism that one can do for this is to take my batch, B. I can cut it up across M machines. If they're divisible, each machine will get a smaller B over M sized batch. And then what I'm going to do is each machine computes gradients, and then I will synchronize all the gradients to get my sum. OK, great. So compute scaling wise, this is perfect. As long as there's enough examples per GPU, we're good. Communications overhead.

[12:53.800] How much do we have to talk to each other? Well, every batch, we have to basically communicate two times the number of parameters, right. Because what I have to do is I have to basically communicate my gradients back and forth. OK, and then for memory scaling, there's absolutely none. Every GPU has the same copies of the models. You need to have the same sized activations so you don't save any memory at all in this scheme. So now what is the problem? Clearly we've addressed compute to some extent. But we have not addressed memory at all. So let's take a look at this memory problem and just do a little bit of accounting. Initially, it might seem that the memory situation is bad. But it turns out our memory situation is actually just terrible. And the reason why the memory situation is terrible is if you just do some parameter accounting, right, it just turns out you need to store a lot of stuff.

[13:49.308] Right, and not only do we need to store a lot of stuff, it's actually more than the parameters that we need to store. So roughly the rule of thumb-- it depends on your precision. But the rule of thumb, you might say, is that we need something like five copies of weights and 16 bytes per parameter to store our model. So we will need to, of course, store our model parameters. We need a place to put our gradients, right. The accumulator before we put them into the model parameters. So we need to store those guys. But actually, you need to maybe store more stuff. You might need to have a higher precision accumulator that you accumulate into in SGD. This might be temporary, but you might need this. And if you're doing Adam, and this is really where the bad stuff comes in, you're going to need to track both the first and second moments of your gradients. And you need to track them over time, right. So you need two of these.

[14:40.700] And unfortunately, you might need to store them in high precision depending on your stability characteristics. And these might be quite expensive, right. And we generally call these things optimizer state. Right, the thing that you're accumulating into. The first moment, the second moment. Right, and these are actually quite expensive. If you look at the accounting, this is most of the cost memory wise of doing an SGD update. And so I think one of the really elegant and nice things that you can do in this space is to start cutting up the amount of memory you need and distributing them across the accelerators, but doing so in this structured way that gets you these gains for basically free. So as I said in the last slide, the optimizer state is most of your memory. So here, the optimizer state is green. That's most of your memory. Orange is your gradients, and blue is your parameters. And notice how parameters and gradients are the same size, and your state is bigger, right.

[15:39.650] If you do naive data parallel, you're just going to replicate all of this state across the GPUs. And this is not so good, right. Because you've, in some sense, consumed way more memory than you need and your memory consumption is linear with your number of accelerators. Not so good, right. What we can do instead is perhaps we can shard the optimizer state onto different GPUs. And this would be a big savings. If we shard different other parts of what we need to track, that would be also big savings. Maybe we can also shard the gradients into different GPUs. Or in the extreme, maybe we can shard everything into different accelerators. And then you can see that the implied memory gains on the very right would be quite dramatic as we go all the way down from 120 to 1.9. Now, the main thing that we will study now is, what is the cost of doing this?

[16:31.170] Right, you might imagine, there's no free lunch. Baseline is going to be the fastest. And as we go downwards to here, we're going to have to pay a big communication cost in order to achieve these memory savings. I think the remarkable thing is going to be that a lot of this will be free. But I'm going to walk through this step by step. So let's start with what people call ZeRO stage 1. So there's going to be three stages corresponding to each of these steps. And so stage 1 is going to be the simplest. The only thing I will shard is my optimizer state. And I will shard it across the GPUs, OK. Everyone is going to have both the parameters and the gradients. And so what you can think of is each worker is going to be responsible for updating one slice of the parameter. So if I'm GPU number 0, my responsibility is updating this very first slice over here.

[17:18.490] And I'm not responsible for any of the other updates. So how would it work? Right, so what we're going to do is everyone computes. Everyone has a different data item, right. So they're going to compute a full gradient on their data item. That's the first step. Now I'm going to take the gradient I computed and I'm going to reduce scatter the gradients onto all the other machines. Because everyone else needs my gradients, right? But they don't need all the gradients. They only need the gradients associated with their update. Right, so every worker is going to get from me, my gradient, but they're part of the parameter space, right. So this is going to be equivalent to a number of parameters communication costs. Because I'm sending my gradient, but I'm cutting it up and sending it into each GPU.

[18:02.050] Now each machine is going to be able to now do their update. They have their parameters, they have their gradients that they got from everybody else, and they have their own state that they've tracked, right. And now, the parameters have been updated. And now I'm going to all gather them back. Right, I was responsible for, let's say, one fourth of the parameters. So those updated parameters, I'm going to send back to everybody else, right. So this is a all-gather, right. So this is where the equivalence is useful. How many operations did I do? Well, in the naive data parallel case, what I had to do was I had to do one all-reduce. Right, I had to exchange all the gradients with everybody and then do the updates, right. So the communication cost is 2 times the parameters. Now in ZeRO stage 1, I did two operations. One was a reduce-scatter to send the gradients out and then an all-gather to collect the updated parameters.

[18:52.230] But we already said a reduce-scatter and an all-gather is equivalent to an all-reduce. So ZeRO stage 1 is going to have the exact same communication characteristics as naive DDP. Right, so this was free. Free memory savings. Wonderful, right? And we look at the memory. We've basically taken the parameters. And anything that's an optimizer state, I've divided by n GPUs, right. Wonderful, wonderful, wonderful. Now, can we do more? Well ZeRO stage 2 is now going to say, I'm going to split the gradients up. And then I'm going to shard that as well. Now, this is going to be trickier because before I was relying on the fact that I can compute the entire gradient, now I can't even materialize the whole gradient. So what do I do. Well, this turns out to just be a systems trick. As I sweep backwards through my network, I don't need to compute the full gradient.

[19:42.670] I can send out the gradients to the associated workers as I go. So the key step here is now, instead of materializing the gradient vector all at once, I'm going to walk backwards through the compute graph. And after computing a layer's gradients, I'm going to immediately just reduce that to send that to the right worker, right. And once gradients aren't needed in the backwards graph, just immediately free it. Right, so you're now incrementally computing and sending gradients as you go. And as you can see, doing it incrementally or doing it all at once, same thing, right? So it doesn't really make a difference. So now you can, at the end, once again all gather the parameters. I've started the gradients in some sense without incurring any additional costs, right. OK, so the last and the most hairy thing-- and this always, when I first learned about this, seemed like total magic-- is we've gotten something for free so far, so let's keep pushing it.

[20:36.390] I'm going to shard the parameters as well all the way across all the states. And so we're going to use the same kind of idea that we used in ZeRO stage 2. And the idea, remember, was I'm going to do things incrementally, right. And so what I'm going to do is I'm going to send and receive parameters on demand while stepping through the compute graph. And this is really tricky, because each GPU only sees a slice of both the parameters, gradients, and optimizer states at any one time, right. And this is ZeRO stage 3, also known as FSDP. If you've ever parallelized anything in a Torch library, you've probably dealt with this before. So the way that it works under the hood is what you're going to do is you're going to do two all-gathers and one reduce-scatter as you go. So we're going to step from left to right here. And this is going to be the entire update process. So what you do is you're going to load your model.

[21:32.030] You have your two GPUs at the top and the bottom. You're first going to gather your weights for this one layer. I'm going to forward this one layer. And then I can free my weights, because I've already done the forward, right. I don't need these parameters anymore. Now, let's say I only have one layer. Now I can move to the backwards process. Now to do the backwards step, I need the activations, which I've kept. And I need the parameters for this layer, which I can all-gather on demand. I all-gather it, I do the backwards. And then I scatter the gradients back right when I get them. Right, so I got the gradients. I reduce-scatter the gradients out. And then I free the weights. I don't need them anymore, and I repeat this process, right. So I grab the parameters as I need them.

[22:10.750] I compute my backwards. I send the gradients out, and I repeat this process. And then I update my weights. All done, right. So what's the communication cost? Right, how many all-gathers did we do? We did two all-gathers, one reduce-scatter. Also, so we have one extra all-gather, which seems bad. That's one thing. The second thing is this seems a little crazy in terms of the overhead, because we're doing communication all the time. Every layer, we're doing this communication. You might think that this is extremely expensive and bad. But I think one of the things that's counterintuitive and important is ZeRO stage 3 uses two really important ideas to make this essentially overhead free. So the first one, which we've already seen, is this idea that we're going to sweep through the compute graph, and then we're going to do the necessary communication and then immediately free the memory. Right, that's already an idea we've experienced.

[23:07.800] The other idea, which I haven't talked about yet, is we're going to overlap the communication and computation cost of this process, right. And this is hard to explain initially. So I think I've taken this plot from this paper, which I think was explaining the PyTorch FSDP architecture. So the way that this works is in order-- so we have different streams here. This is what's happening in the CPU. This is what's happening in the GPU for computation, and then what's happening in the GPU for communication at the very bottom here. Now, the way FSDP works is I have to first get my parameters before I can do my computation. So I'm going to all-gather 0, layer 0, right. And once I've all-gathered, I can now do my forward. But as I'm doing my forward, notice that I'm requesting the next layer's parameters. I'm all gathering 1 while I'm doing this computation. And now I all-gathered 1, I can now do forward 1, right.

[24:02.680] And I can start all-gathering 2 while I'm starting this process. And so I'm overlapping computation and communication as I'm doing this. And I might spend some time freeing my parameters. I might also spend some time computing another forward. For example, if you reuse that W nought, I'm reusing this forward 0 twice. And so what this allows you to do is if you have enough computation and your network's fast enough, by overlapping communication and computation, right. FSDP can basically be free. Right, you do see some bubbles. These are overheads. But if your comms is very fast and your computation is very big, right, you can easily see how the computation can run faster than the communication-- I'm sorry, the computation can take longer than the communication, and you end up having this really great improvement in memory usage without paying for almost any amount of memory use. OK, so when we look at the different ZeRO stages, right, DDP, the naive data parallel, costs 2 times parameters.

[25:06.800] Right, it was a single all-reduce. The two ZeRO stages, 1 and 2, right, the ones that shard gradient and optimizer state, this is free in the literal sense, because it has the same amount of communication cost, right. It's just the identity of all-reduce being equivalent to two operations, right. Now, ZeRO stage 3 is technically a little bit more communications. You need an extra all-gather, right. But it's not so bad because of the really clever ways in which all these libraries can hide the communication cost underneath the computation. So in practice, FSDP, if you run it, you're going to see GPU utilization that's very close to just the single GPU performance. It is actually quite remarkable how good FSDP is. And it's also very conceptually simple. You will actually have to write an FSDP implementation as part of your assignment. And you can just write a wrapper that will wrap any module and turn it into its FSDP version.

[26:02.620] Right, conceptually, all you're going to do is do a bunch of all-gathers, compute, free, and then repeat on the backwards pass. OK, good. Yes? [INAUDIBLE] --gradients from the next GPU to calculate the gradients for the previous GPU, right? When you say-- so we're not-- I think what you just described, which is taking gradients from one GPU to the next, that's closer to pipelining. Pipelining would be if one layer was on one GPU and another layer was on another GPU. That's not the case. So in all these cases, every GPU goes through the entire model, right, from the start to finish. But the difference is, no GPU is going to hold the entire parameters at the same time.

[26:49.000] And so what we're going to do is if I only have part of the network that doesn't have layer 0, I'm going to demand the layer 0 parameters first and then do the full computation. And then I'll move on to the next layer, right. So it's the same computation. It's the same single GPU computation as naive data parallel. But you can just think of it as request parameters, free parameters in between every computational operation. Yes? So are you saying that every GPU has a part of every layer, and that's what the all-gather's for? Yes, that's right. So the picture is here, right. Parameters, gradients, and optimizer states are all sharded across the different GPUs. That's right. Yeah, good. Oh, yes. Why does the communication cost not multiply by the number of layers as well? So the number of operations multiplies by layer. So that's correct. But each of these communication pieces is smaller, right. Because this is just one tiny MLP or something. And I'm comparing that to the cost of all-reducing a whole network. Right, so it all adds up in some sense. Yeah, good. All right. OK, good.

[27:55.080] So in practice, if you look at the maximum size that can fit on, let's say, a A100 GPU, if you go from the baseline, you can not even fit a 7B model. And you go to ZeRO stage 3. You can now fit 50 billion parameter models and so on, right. So it becomes a lot easier to fit these big models using FSDP. Right, the improvements that you get in terms of the memory improvement is quite significant. OK, yes. So can you go back to the previous slide? So regarding this on 8x1 A100, what would be the limitation on H100 or P200? OK, I can't do the mental math of multiply by 141 and divide by 80, but that would be the number for an H200. And so on and so forth. Yeah, it would be linear, right. It wouldn't really change the math very much, though, to go from one accelerator to up. OK, so FSDP, I think, is very clean. It's very elegant. I like it a lot. I wish I could end the lecture here. But unfortunately, we will have to proceed into uglier and more hairy things.

[29:05.910] And there's two reasons why we have to do that. One of them, and this is an important idea for parallelization, is that in some sense, data parallel is consuming an important resource. And that resource is the batch size, so to speak, right. If you have a batch size of 8, you can never have more than eight accelerators, right. And you might think, OK, let's just make the batch size bigger and bigger and bigger as we have accelerators. You can't do that, because at a certain point, there's something called the critical batch size, where the basically gain that you get from the additional batch element is less than if you had taken another SGD step on that single element, right.

[29:43.950] So there comes a point where at small batch sizes, adding extra elements is perfectly helpful. It's the same as taking more SGD steps. But at a certain point, diminishing return hits in. An infinitely large batch size is not infinitely better than infinite steps, right. Infinite single steps. Because of this, now we have faced a hard tradeoff. Do we want small batch sizes and let our GPUs be more idle? Do we want big batch sizes and take the hit from optimization? Hard to say, rigt. There's another issue that we have, which is that we still have memory issues, which is ZeRO stages 1 and 2 don't really let you scale memory at all. And ZeRO stage 2 lets you cut up the parameter memory. But unfortunately, there's activation memory and other kinds of memory use that this does not split up or this does not reduce at all, right. And so we need better ways, more fine grained ways, of cutting up the model in order to really start pushing down the memory use beyond these points, right.

[30:41.290] And so this takes us to various model parallelism ideas. And so model parallelism is going to be this idea of essentially like FSDP, we're going to split up the parameters across GPUs. But I think the important conceptual difference between what we talked about and what we will talk about next is that now we are going to communicate activations. Right, in FSDP, we cut up the parameters. But in some sense, it was just a wrapper. We were still doing the normal computation. We were just sending parameters back and forth before doing it. Right, so parameters were the things flying around. Now what we're going to do instead is we're going to communicate activations back and forth, right.

[31:21.510] So if, for example, one layer lives on GPU 0, the next one lives on GPU 1. Right, what I have to communicate is the activation of the layer in between. Right, so that's going to be the big difference. And this will involve pipeline parallel, which is cutting up layers, tensor parallel, which is cutting up matrices, and expert parallel, which is sharding experts into different places. Right, and Percy's talked about tensor parallel already, so you have roughly some mental model of this. I'll start by talking about pipeline parallel or layer wise parallel. This is, I think, conceptually very simple. Back in the days when I first learned about parallelism, I was like, ah, yes, it's simple.

[32:01.090] Just cut up the layers and put them on different GPUs. Right, very simple thing. And so you're going to pass activations forward in the forward pass, and then you're going to pass like partial gradients backwards in the backward pass. Now, OK, if I do this, what will happen? You will get a terrible looking picture like this. Very, very depressing. I have four accelerators, each handling a quarter of the model. And so what computation happens as a function of time. So time is sweeping from left to right. And the different accelerators are the different rows. So first, my zeroth accelerator at the bottom here, that's handling the first layer will do some computation. And then it will stop, right. And then it will hand off the computation to the second one. And then it will do nothing for the rest of time until the very end.

[32:44.672] And then I will basically be passing things back and forth. And one GPU will be active at a time. And then coming backwards, one GPU will be active handling the backward pass at a time. And so most of the time, my GPUs are idle. My utilization is truly, truly terrible. And this is called a bubble. It's the part of the pipeline where you're not doing anything. And the solution to this is to essentially do a batching or pipelining. So instead of doing a simple one layer after another, you're going to have a pipeline, right. Where, in this case, in this picture here, you have four elements to be processed. And as soon as you process one, you start on the next element, and then you pass off the one that's already done to your next layer. Right, so you can do this to essentially pass elements upwards and then back downwards.

[33:33.390] And so if you want to make this reasonable, then what you need is the ratio of essentially this bubble time to the useful compute. The utilization, in some sense, is going to be the number of stages divided by this micro-batches. And so we need a huge batch size in order to reduce this bubble towards 0. Right, so if we want our bubble to be 0, it's going to go down as 1 over the microbatch size. This is why I said batch sizes are useful resource. Because now we can spend it in a different way, which is to pipeline more to reduce the amount of idle time in pipeline parallel. OK, so pipelines are terrible. And by folklore, people say things like, parallelization code looks reasonable and people can understand it until you implement pipeline parallel. So why do we do pipeline parallel? Well, pipelines can save memory compared to DDP.

[34:28.420] Because we've cut up layers, and you can also compose this with data parallel. So you can do both. But the other reason why this is a very important idea to learn, and you will always have a pipeline in certain kinds of topologies, is that it has very good communication properties compared to both FSDP and compared to any other kinds of things you can do. It only depends on activations, like b times s times h. It is also point to point. One layer is going to communicate to another. It's not all to all. And so depending on how you sequence these things, they can be a very efficient way to utilize your network.

[35:01.980] So in practice, what we're going to do is that pipelines are going to live on the slowest networking links. So if you have multiple data centers or multiple pods that are slower to talk to, you're going to parallelize across those two using pipelined parallel, because that's the most communication efficient parallelism that you can bring to bear here. OK. And as I said, pipeline performance is really dependent on batch size. So I'll point out this paper repeatedly, but I'll mention which one this is. Now this is the Megatron paper from NVIDIA. They have a bunch of really lovely parameter sweeps of what happens to utilization as a function of different configurations you can do. So if you have large batch sizes and you pick a large pipeline size, you can get pretty good utilization close to not doing pipeline parallel at all. But you need the big batch sizes in order to be able to keep utilization up. Otherwise, pipeline parallel is going to rapidly degrade your utilization and performance.

[36:03.472] Actually, I'm going to pause for a moment here to ask questions, because the next slide is complicated. Yes? So you mentioned on the previous slide that pipelines have good communication properties, like the FSPD. What is the specific mechanism for why that is? It's a smaller amount of stuff that's being communicated. So it's batch times sequence length times hidden. And that's almost always a smaller amount of data than attempting to communicate a whole parameter matrix, right. So people have tried to figure out much more clever ways of reducing the amount of pipeline sizes. And you can do clever things that cut up the different stages of what you're doing for different layers or different micro-batches. So you can sequence different forward pass elements in between different backwards pass elements. And by doing clever scheduling, you can further reduce the bubble size.

[37:01.920] This was taken from the DeepSeek paper. So by really cleverly manipulating when you're doing forward and backward computation, you can get a little bit better in terms of the bubble size. But I think the even more clever thing that you can do, which is called zero bubble pipelining, is not really about scheduling per se. But it's actually about thinking carefully about the structure of your backward pass. So in the backward pass, there's really two things that have to happen. If you think about it, you're going backwards in your computation graph. And as you do that, you're going to do two things. One is going to be you're going to propagate your partial derivatives further down your computation graph. So remember, back propagation, you start at the end, and you're propagating these partial derivatives back, right. So you have to do that. But you also want to compute the gradients for the current weight. Right, so you have two things that you do at every node of your computation graph. You're going to propagate backwards, and you're going to compute the derivative of your weight. The important thing about thinking about this compute graph is actually, one of these two is really important for pipelines.

[38:08.740] And the other one, you can do whenever. Propagating the partials backwards, right, that's very important. Because your next stage can't do any work until you've propagated that signal back. On the other hand, the derivative with respect to the weights, that is a leaf node, so to speak, on this graph. And I can do it whenever, right. And so the clever thing to do now is you separate these two. Right, so you got your B's over here. These B's are the backwards, propagating backwards. And then these W's, these W's are computing the derivative of the weights. And they can happen whenever. So what you do is you compute the B's as quickly as you can. Right, so those come first. And then you can defer the W's until later and do it when you have a gap in your computation. And when you do this, you can basically almost all completely fill up the pipeline.

[38:56.180] And the zero bubble pipelining thing is very complicated or much more complicated than you would normally like to deal with. But it allows you to almost deal with these pipelining issues almost entirely, depending on essentially the workload involved in the W versus the B, and so on and so forth. But I think having seen this, I was like, oh, this is so clever. Such clever things that you can do with systems. OK, so that was pipelining. Any questions about that? Good. So I think there's a natural analogy or a parallel you can think of. So pipeline parallel is like depth. So your depth, you cut up your network along the depth dimension. But we have two ways that we scale, right. We've got a depth axis and we have a width axis. So what if we start cutting up our model along width, right? So this is tensor parallel, as Percy explained before.

[39:50.580] And it's really the simple idea that matrix multiplies can be cut up into smaller matrix multiplies, which you then add the partial sums back together. This is really the same idea as tiling as well. So this core primitive just appears many, many times. And so when you do this in practice, right, what does this actually look like? So this is a GeLU that's happening. So what do you normally do? You have your x coming in as input on the very right of this slide here. And then I have a matrix A that I multiply with. I apply my non-linearity GeLU. And then I multiply with B to downproject. And then I output my z. Right, this is the normal serial computation that I do. Now in tensor parallel, what I would do is I would cut out my A into two submatrices, my B into two submatrices, and then I would run the computations in parallel, and I would gather them together at the end.

[40:45.180] Now, one thing that is important is there's a duality between the forward and backward pass. So in the forward pass, all I do with the x is I copy them twice, and then I run the computation forward. So f is the identity, this f function. And then once I get to the end of my parallel component, the g function, that's an all-reduce that brings them back together, right. But in the backwards pass, as I'm doing my gradients, this is going to flip. My g is going to be on-- sorry, my g is going to be an identity. Right, because the partial derivatives are coming backwards. And then my f are the partials that I need to sum up in order to get my partial derivatives at the start of this module. Right, so they're going to flip. And this duality is important if you're going to write tensor parallel. OK. And so you might have already noticed, right. But when we split things, some of the ways that we are going to be column-wise. So column-wise, tensor parallel is going to happen at the inputs.

[41:44.990] Right, so the inputs of the MLP, the projections of the attention are going to be cut up by the columns in each transformer block. So this is column-wise cuts. And then there's going to be row-wise cuts in the corresponding second stage. So in the downprojections of the MLP, as well as the outputs of the attention, you're going to have a row-wise cut. And then everything that's a small layer, like a layer norm or a non-linearity or routers in ML-- sorry, in MoEs, those are all fully replicated across the machines. You don't really want to bother with the overhead of cutting those guys. OK. And then, so when do we do tensor parallel? Tensor parallel is extremely computation hungry. What you're doing is every time you have a matmul, you're going to be doing some communication. Right, you have an all-reduce. Or in the backwards pass, you have all-reduce here. One in the G when you're going forward.

[42:44.330] And these are activation size, and you're doing them very frequently. So this is very communication hungry. And so you only want to do this within a node. So for GPUs, GPUs are going to be networked tightly together in a single box. And so up to eight, you might be willing to do some tensor parallel. Once you go beyond that, you're going to be going cross node. And those interconnects, those connections, are much slower. And so you're going to get a big drop in performance as soon as you hit that point of going beyond a single machine. So tensor parallel, you should remember as very communication hungry. You usually use it in the fastest interconnects that you have at hand. And this is the point at which I'm going to go back to the thing I said about TPUs. Right, remember for TPUs, you don't really have this distinction between fast eight GPUs and next machines.

[43:37.470] You just have this big mesh. And so one of the big advantages that TPU people will tell you is you can tensor parallel very large numbers compared to the GPU world, because they're able to have high bandwidth in this very regular communication pattern over this mesh. Right, so that's a big difference if you're parallelizing between TPUs versus GPUs. The amount of tensor parallelism you use versus pipeline parallel will actually be quite different. So how do we compare these two. They're both model splitting strategies. They'll both save on both parameters and potentially activations. The pros here of tensor parallel is that there's no necessarily pipeline bubble. And if your network's fast enough, you might be able to get full utilization. It's also low complexity. It's just cutting up some matmuls.

[44:25.410] But the con is that the communication is much larger. It used to be that you were doing b times s times h point to point communication. Now instead, I'm going to be doing something like a times b times s times h times roughly one all-reduce communication, right. So this is not even just point to point. This is all to all communication that needs to happen. So tensor parallel is great. Whenever we have high speed interconnects, every other time, we probably want to be using pipeline parallel. OK. So the last thing I want to talk about about-- we're talking about memory use, is to really think more carefully about memory. I think the most naive view of memory in deep learning and language modeling is something like, well, memory is just parameters. That's this green bit over here. Right, you have a certain number of bits you need to store for your parameters. That's it. I've already talked about optimizer state, which is important, which is this yellow box over here.

[45:23.713] But if you look at this plot, this is actual memory use as profiled in one of the PyTorch profilers. You'll find actually that this is not really the whole story. There's a lot of dynamic memory that needs to exist in a computation. You need to store all the activations. These are these big red bumps over here. And then as you compute your gradients, you need to store them in your backward sweep, right. And so really maximum memory usage happens a little bit after the maximum activation point. After you start sweeping backwards on your gradient, but you still need to compute a lot-- keep your activations, those are usually the maximum memory points, right. So to reduce this, maybe we need to deal with this big red hump, this activation. And if you start looking at more modern workloads for language models, you start to find that as you get to bigger and bigger parameters-- forget the column that says present work.

[46:17.950] I'll talk about that in a moment here. But if you look at just the baseline columns of this plot, as you go to bigger models, for moderately large sequence lengths, your activations are just going to dwarf the parameter memory that you need. And so any memory saving strategy has to reason about activations in order to be fully effective. We can't just deal with parameter memory. That's not really sufficient for us. So if we're storing everything-- and this is a good, I don't know, statistic or rule of thumb to start with. If we need to store everything, the amount of activations that we need is going to be something like 34 times s times b times h, and that's going to be sequence length times batch size times hidden dimension size plus 5 times as over h. as is attention heads times sequence length divided by hidden dimension size. And this sbh dependence, right, is fundamental, because we expect to have to store something for every element of the sequence.

[47:21.833] We expect to have to store something for every batch element, and we expect to have to store something for every hidden dim size. Right, so we're going to see this sbh term repeatedly as we try to go through different ways of reducing activation memory and accounting for the total amount of activation memory. And in case you're curious where the terms come from, this odd looking term, the 5 as over h is going to come from the quadratic attention terms, including the dropout terms. And we can drop this term via recomputation if we do flash attention. OK. Now what if we did tensor parallel? One of the promises that I made to you in talking about model parallelism was that we could reduce activation quite a bit. So now let's think about what happens. Well, tensor parallel is going to split out the matrix multiplies in attention and MLPs. So which parts are those?

[48:11.760] Well, the MLPs are 24 of those 34 that I talked about before. And the attention heads are 5 as over ht. Right, the second term over here. Now, I can reduce that by t, my tensor parallel size. And that's good. This is not bad. Because I've taken most of the memory, and I've made it scale linearly with the tensor parallel size t. So that's good. But it turns out that there are other things that you do. You might have a layer norm. You might have a dropout. You need to keep the inputs to attention and the MLP. Right, because the inputs to each layers need to be stored as residuals for the backwards pass. And unfortunately, these are not reduced with your tensor parallel size. Right, tensor parallel does not split the layer norm. And therefore, it does not split the activation for these guys. So this is unfortunate, because you were hoping that if you had 1,000 different GPUs in tensor parallel because you're Google, you might be able to reduce this activation dramatically.

[49:12.780] The unfortunate reality is that you're still going to suffer a 10 times sbh penalty for doing all of this. And so the last part that's usually used in composition with tensor parallel is this idea of sequence parallel. This is an extremely misleading name, because usually when I say sequence parallel, actually there's a thing called context parallel that is actually more natural for this name. But sequence parallel is the following idea. So what you're going to do is we have these remaining terms, 10 sbh. And these terms are very lightweight. They're like layer norms. They're stuff that doesn't have much computation. And what I'm going to do is I'm going to split them up somehow. And I'm going to split them up over the sequence axis rather than over the hidden axis. And this is going to then involve all-gathers and reduce-scatters before every operation where we need these. And if you think about it, this is very reminiscent of FSDP.

[50:13.475] Right, we have this thing that we need. We need attention-- sorry, we need activations at some point. We need to compute the activations going forward, and we need the activations when we sweep backwards for the derivatives, right. But we don't need them now. So we're going to store them by splitting them up across the sequence axis and materializing them on demand. So I would say this is conceptually very similar to the FSDP style idea, where we're storing them in this sharded format, and then we're gathering them as we need them. And then as I said before, there's a duality to this. So in the forward pass, g is an all-gather. g bar is a reduce-scatter. In the backwards pass, this is going to be reversed. Right, the all-gather and reduce-scatter are going to be reversed between g and g bar.

[50:57.600] So if I do this now, what do I get? Well, remember, with the no parallelism setting, the very first part, what I told you was you get this activation memory. Use 34 times sbh with an attention head component. With tensor parallel, I get to divide the MLP part and the attention part by t, but not all these other pointwise ops that I have. If I do sequence parallel, I've split these up by my accelerators as well, just in the sequence axis, which gives me now fully linear dependence. If I do activation recomputation, I can even drop this second term. Because remember, this is just the storage that I need for the softmax. So I can drop this through recomputation, which gives you sbh times 34 over t. This is nice to remember, because this is the lower bound of what you can achieve, reasonably speaking, for normal training for activation memory. And so if you want to do things like compute by hand whether your model will fit into a GPU, this is a good thing to remember.

[51:59.220] You need at least this much for activation memory. You're going to need some memory for your optimizer state, and you will need some memory for your parameters and your gradient. And once you have those things, you have a rough sense of how much memory you're going to need for your model. OK, I'm going to stop here for a moment. I have talked through both pipeline and tensor parallel at this point, and so people might be a little bit confused. And we can talk through a few things before I move on to expert parallel. Good. Oh, yes? [INAUDIBLE] 24 over t is the MLP part. The MLP part cannot be optimized through the representation? You can recompute the MLP as well. You can do a lot more computation than what is listed here. But recomputation what is listed here is very computationally expensive, like a recomputation for MLP involves running the MLP again in the backward pass, which you probably do not want to do.

[53:04.880] Recomputation for attention is generally cheaper, yeah. Because you do it time-wise as well. And also you don't want to pay the quadratic cost. That's quite expensive. The quadratic cost. That's quite painful, yeah. Cool, all right. So I'm going to now talk about the last ingredient in the standard parallelism toolbox, which is expert parallelism, right. Now that MoEs are very standard part of the toolkit, most big models are MoEs these days. And that allows us to take advantage of one of their big benefits, which is expert parallelism. And expert parallelism, I think of myself as being something analogous to tensor parallel. Right, in the sense that you're taking the MLPs, right, the FFN components of your network. And you're going to split them up in various ways, and you'll incur a communication penalty in the same way that you're maybe splitting up the matrices, except here you're splitting up the MLP across different devices.

[54:07.870] And the systems behavior of tensor parallel is-- oh, sorry, expert parallel-- is roughly like tensor parallel as well in the sense that this is a high bandwidth parallelism primitive that also reduces activation. But one important thing is that if you're doing an MoE, you almost always prefer using an expert parallel sharding strategy over tensor parallel. And I've taken this screenshot from the parallelism guidelines from Megatron, which is NVIDIA's like parallelism library. It's a really wonderful library and document. And one of the guidelines, they say is, look, if you're going to do some of parallelism, either EP or TP, you should be using EP over RP. Why is that? Well, this is actually a good list of drawbacks for tensor parallel. If you have a matrix and then you cut up the matrix too finely, well, the matrices are going to start to get small and your GPU utilization will suffer. So you want your matmuls as big as possible, and tensor parallel reduces that.

[55:10.070] For MoE layers, it's a lot easier to route sparse token activations rather than to potentially route these dense big tensor parallel matmul activations. And then you can do things like route tokens exactly to the places that they need. And you can skip some computation-- or sorry, some overhead in the cases that if you're going to have MoE anyway, you might as well have the number of experts split up over all the devices. Now, this-- sorry, yeah. This might make it seem like, all right, EP is great. Everyone should be doing expert apparel all the time. Unfortunately, expert parallel is still really complicated and still very difficult. I think everything that I've said after FSDP-- tensor parallel, pipeline parallel, systems wise is quite non-trivial to get working really well. And I'll just point to two different libraries not to tell you the details of what happens inside of them, but to tell you, this is seriously complicated business.

[56:16.110] DeepSeek has a library called DPP, which I think is one of my favorite things from the DeepSeek V3 era, which is basically their own library for doing expert parallel routing and dispatching, where they're looking at really low level GPU networking primitives in order to try to basically combine different operations or have the hardware handle certain operations in order to make MoE routing as efficient as possible. Similarly, NVIDIA has a library called Hybrid EP, which is a similar effort of having very efficient low level hardware implementations of expert parallel dispatching. And just to give you the high level thing, right, the reason why this stuff is hard is you are doing a lot of all to all dispatching. If you look at the communications pattern of MoE, much like tensor parallel, every time you have a MLP, you're going to have to route tokens all sorts of different places, and you need to do this in a very latency sensitive way. Because your computation is waiting for your tokens to arrive.

[57:20.070] So reducing the latency of this dispatch is extremely, extremely important. And one final thing that is like a fun trivia about the DPP library is the deep sea cloaks are so intense that in order to optimize, like every last bit of performance, they basically looked and found undocumented ptx, which is like GPU machine code or machine code like things. And they found how to use these undocumented instructions to further accelerate their networking communication. So really, this is the level of stuff that you need in order to get to the frontier of parallelism efficiency. Finally, up until now, I've talked about all the parallelism primitives in isolation, like data parallel and so on and so forth. With the implicit assumption that you can just put them together like LEGO blocks. And this is true, for the most part, for basically all the other parallelism strategies. It's very easy to combine, I don't know, data and tensor parallel. Expert parallel has additional constraints that you want to be somewhat respectful of.

[58:28.110] Usually, the naive way to do data and expert parallelism, which I think is in many of the old libraries, is that the replicas for DP and EP are the same. So in other words, let's say I have a data parallelism of eight. I'm going to shard my experts across those eight replicas. And then my data will be split across those eight replicas and routed in various ways. Right, so in other words, you split your GPUs according to DP, and then your EP is a subset of that split. Right, and this is a very natural thing to do because you're routing tokens. But if you do this, there's a maximum bound to how far you can parallelize with EP. And it constrains how DP and EP And so I think this is the final complexity. I apologize that expert parallel is a little bit really messy in many ways. But because of these complexities, there is actually this one interesting systems thing that I'll tell you about before I stop talking about EP,

[59:30.900] which is this following fact. Remember, MoEs only change the MLPs, right. They don't change the attention at all. Right, so because we're cutting up the MLPs, but not the attention, this means that expert parallelism kind of applies unevenly to the model. So when I parallelize at the expert level, I'm parallelizing the MLPs, but I'm not parallelizing the attention, right. And so in order to parallelize the attention, maybe I want to also use tensor parallel, right. So I want a high tensor parallel level in order to cut out my attention. But this is going to then affect my MLPs. Because if I have high tensor parallel and high expert parallel, I've cut out my matrices into really tiny pieces, and my utilization is very bad. So there's this weird tradeoff where you want high tensor parallel for the attention, but you want low tensor parallel for the MLPs. So there's this conflict. So what do you do? Well, in the last few years, people have come up with more complicated system solutions that decouple the tensor parallel in the MLPs and the attention.

[01:00:39.440] And so the attention layers get one kind of tensor parallel, and the MoE layers get another kind of tensor parallel. And so now you can fully decouple the parallelism. And this leads to more complicated but more effective combinations of EP and tensor parallel, as well as data parallel. Cool. OK. So the last thing in this space that I want to end with is context parallel. So context parallel or ring attention is an idea of basically splitting activations in a very long sequence across different accelerators. And you can do things like pass the activations to the device that's needed in this ring-like way following the mesh topology of TPUs. And so ring attention, which I think was the original paper that did this, showed that this worked really well on TPUs. Context parallel is a standard parallelization strategy that does this. Both of these are used in long context extension stages and in model serving.

[01:01:44.740] I'm not going to talk about it, mainly because I think it overlaps in concept to a lot of what we've talked about already. OK. So this is the end of the second part. And I think to put everything together, I've made a big table of parallelization strategies, including data parallel, which are the top two rows and the various model parallelism strategies below them. And I think I've colored in red what I subjectively feel like are the drawbacks of the different methods. And the reason why I've colored things in red is to try to convey to you that there is no one strictly dominant parallelization strategy. It's all a whole bunch of tradeoffs that you somehow have to manage gracefully in order to get a good outcome. So for example, FSDP is great. It's a wonderful strategy. Does not help you with activation. It consumes your global batch size in order to do parallelization. So it has limits. And to address these, maybe you want to compose them with some of these other tricks that you have at hand.

[01:02:46.892] Maybe you want to do some tensor parallel to cut down on activation memory. It also doesn't touch global batch size. But if you do this, now you need fast networking and you need high bandwidth. And also then to leverage your slower connections, maybe you want to use some of the other parallelization strategies. Right, so you see how all of these different advantages means that there's a place for one of these parallelization strategies in not every architecture, but in many large scale architectures, you use a large combination of these. Yes? [INAUDIBLE] Are they still applicable to MoEs, or is that strictly-- Yeah. So the question was, is pipeline and, what, FSDP applicable to MoEs? Yes, yes, they're absolutely applicable. If anything, the standard recommendation is expert parallel is like tensor parallel. So you want to go up to 8 or fast connections. People don't really follow that anymore.

[01:03:48.620] But even then, FSDP or data parallel in general and pipeline and tensor are used extensively in all the frontier models. Yeah. Good. OK, and you will do this or something like this as part of your assignment. But one of the things that's cool about parallelism and I think is nice is that you can do some math, right. You can actually say, OK, how much compute do I do per layer for my different sharding strategies? How much communication do I do per layer for my different communication strategies? And if I have different devices allocated to different communication strategies, how do these things scale? And then I can figure out for a combination of these strategies, how do they scale? And you can start to make plots of, if I have a certain batch size and I have a certain parallelization strategy-- so I have three different parallelization strategies here. How much am I utilizing my GPU? The ratio of compute time to communications time. And as long as my compute time is longer than communications time, in principle, you can just hide communications underneath your computation. And so you can make plots like this where this dashed line at the very top, that's the boundary of where you're efficient.

[01:05:04.520] Anything above that is fully utilizing your computation. Anything below that is on the bad part of your roofline. You're waiting for communication to do your computation. One thing that's pretty interesting here, right. So I think this should feel intuitive to you based on what I've told you. If your batch size is big enough, FSDP only is good. Right, if you have a batch size of 2,000, right, per chip, then you how far can you go? Well, you've got FSDP doing extremely well. You're fully compute bound, right. But the important thing is, as the batch size goes down, FSDP only will hit a point where you're now communication bound. Now, what do you do, right? Well, what do you do is now, in order to keep pushing efficiencies, you have to incorporate MP, in this case, is tensor parallel, to the mix. Right, now you can push that curve out a little bit more and be compute bound even into smaller batch size regimes, right.

[01:05:59.310] And you keep adding strategies to push this curve out in order to get something better. And this is what people call 3D or 4D-- I don't think there's 5D-- parallelism. Which is when you put all these strategies together to keep your compute units fully utilized under different communication topologies. There's a very simple prescription. I think up until now, I've been telling you things in generalities. So it might feel complicated. But I think actually the strategy for how you parallelize is very simple, actually. Which is, until your model fits into memory, you're going to cut up your model by whatever means necessary, right. And how are you going to cut up your model. You're going to use either tensor or expert parallel for your fast interconnect. And so if you have eight GPUs per machine, you're going to use a tensor expert parallel of eight. And then you're going to pipeline parallel or Zero-3, which is FSDP, the rest of the way to make your model fit.

[01:06:53.070] And once your model fits, then you're going to just data parallel the rest of the way. Right, then use up the rest of your GPUs with data parallel, right. If your batch size ends up being too small, you can do gradient accumulation to get better GPU utilization. Very simple strategy. And this is borne out in the actual large scale parallel training libraries that you might use. So I've added the link here to Megatron's guide to parallelizing MoEs. But this idea actually applies to the dense models as well. And you see exactly the thing that I said, but in reverse order. So guideline number one is minimize model parallelism, maximize data parallelism. Right, so use most of your GPUs to shard data if you can. Now, if you're using GPUs, stay within NVLink. So that's one machine, or so one box, right. Keep expert parallel and tensor parallel within one box.

[01:07:44.190] And then if you're going to go multi-node, use pipeline parallelism to go multi-node. And then if you're MoE, then prefer expert parallel. And then finally, if you're doing long sequences, use context parallel. Right, ring attention. So exactly what I've been trying to get across to you through the lecture, but in practitioner form that you see here. OK, the final component-- actually, maybe I can pause here, in case anyone's curious about 4D parallelization stuff. Yeah. I know you said context parallel would be a better name for decreasing parallel [INAUDIBLE] Oh, yeah, yeah. Well, I mean, sequence parallel is often just you do it with tensor parallel in order to reduce activations. It's more of an add on that you put in in order to reduce activation. It's not its own standalone thing, in many cases. Yeah? With the suspicions that this is like a loop transformer, like that, does that make things easier or harder? OK. I don't how much I buy the loop transformer mythos rumors.

[01:08:57.430] Maybe it is. It would be interesting. But to the question of how the systems structures would change, it's a good question. I mean, I guess it might mess with things like FSDP. You can't really discard the weights. It does have the advantage, though, that you don't have to keep getting weights. It is much more parameter efficient, and so maybe much of the model parallelism stuff wouldn't be as important. It would be interesting. It would be definitely different for a lot of the FSDP style things. Because the FSDP strategy is really predicated on getweights, discard weights, getweights, discard weights. And in a loop, you're not discarding anything. Good, OK. All right. So this is an old paper, but it is actually probably one of the best resources. And also, it's not like the networking fundamentals changed that much. So I think the lessons from here are very relevant today.

[01:09:53.050] And so this is a paper by NVIDIA and formerly unfortunately, Stanford folks, Matei and Deepak, who did a bunch of runs of large scale parallel training across many different configurations to basically tell you, as you're scaling up, how should your parallelization strategies change across different scales? And you see exactly this prescription that I've been talking about before, which is, you have data parallel maxed out. And then your tensor parallel increases until it hits 8, and then you stop because you don't want to go any higher. And then from that point on, your pipeline parallel increases and increases and increases, right. And with sufficiently large scale, data parallel decreases. And at the very end, you have a DP of six. And this is because you just need this much tensor and pipeline parallel to fit the stuff. Right, and you're using the rest of your budget on data parallel.

[01:10:47.130] And if you put them all together, right, even as you go to ludicrously large numbers of GPUs, your utilization stays very flat and very good. And this is why I think you see this gigantic data center build-outs. The parallelization strategies people have come up with are so effective. And the hardware for communication is so good that even with the gigantic data centers or even cross data center training, you can get really good utilization across these different regimes. They have really nice quantitative evidence showing you, for different batch sizes for tensor parallel, what is the right place to stop? And it's clearly tensor parallel size of 8. As you go beyond that, you're going to get into trouble. And pipeline parallel, you also need big batch sizes in order to be able to utilize the machines effectively. That's why there's a blue to orange gap. There's pipeline parallel.

[01:11:44.700] And the larger the pipeline parallel, the larger the gap between orange and blue. And then I think another thing that maybe you have already internalized from flash attention but maybe not is maybe you should do a bunch of activation recomputation. Because if you do activation recomputation very cleverly, then you can get a bigger batch size, and the bigger batch size then allows you to get better utilization. This is a, initially, I think, very counterintuitive observation that actually you should do more computation via recomputation in order to end up getting better utilization, because it allows you to save memory, and memory can be turned into batch size. OK. So that was a more quantitative view of how to pick parallelism strategies. The last thing I want to end with is just to look at just a very few training runs. And looking at the training runs, maybe we can get a sense of how modern parallelization is used in practice and how it's evolving.

[01:12:46.720] So Dolma by AI2 was a seven be open source model. Sorry, Olmo, which was trained on the Dolma data set. And this isn't the Dolma paper for the details. Was a 7B model trained fully with FSDP across a whole bunch of accelerators. I forget exactly how many. But this is one example of FSDP scales surprisingly well, even to many, many GPUs. As long as you're training a small model, FSDP is actually a pretty good strategy for parallelization. I think many 7D-ish models are trained purely with FSDP. DeepSeek, because of their systems skills, does do some pretty exotic and useful stuff.

[01:13:31.560] So DeepSeek V1 was trained just with data parallel ZeRO stage 1 with tensor sequence and pipeline parallel, as I talked to you before. DeepSeek V3, that's a MoE model, if you remember my MoE lecture. So they have not only pipeline parallel, but they also have expert parallel instead of tensor parallel, right. And their expert parallel actually is a little bit exotic. Because they have 64 way parallelism. So they group eight different machines together, and that's their expert parallel domain. And to enable this large expert parallel, they basically use the same tricks that they have from their pipeline parallel pipelining in order to try to make sure that expert parallel doesn't have low utilization periods. So this is actually quite a complicated trick that they do to get these very large expert parallel splits. Another example, Yi, from Chinese open weights training setting. They use, once again, ZeRO stage 1 and tensor and pipeline parallel.

[01:14:32.480] Right, this very much makes sense. It's the classic data parallel, tensor parallel, pipeline parallel combo. And once again, once you go to MoEs, you replace the tensor parallelism with expert parallelism. Right, they serve similar goals. But expert parallelism is just a little bit more efficient. On the other extreme of a really big dense model, you have Llama3. The Llama3 report, I liked quite a bit, because they're one of the few reports that have literally a full breakdown of their parallelism strategy. Not just like in the abstract, but across all the different phases. So the cool thing that you see is they have a pre pre-training stage, a little warm-up phase at the very start. You can ignore that first row. And then they have their main pre-training phase in the middle row here. And this is a very standard parallelization strategy.

[01:15:21.140] Tensor parallel of 8, one context parallel, 16 pipeline parallel, and 128 data parallel. Big data parallel splits. And then at the very end, because they want a long context model, they do long context extension. And for this, they crank up the context parallel, lower the data parallel. And that's how they parallelize this very memory hungry stage for long context extension. And notice, we're in tensor parallel, not export parallel land. Because Llama3 405B is a gigantic dense model. It's not a MoE model. And as a side note, right, one thing that I don't talk about, but if you, for example, become an infra engineer at one of these companies, you'll notice is GPUs fail all the time. Apparently during Llama3 405B training, GPUs failed 148 times. And so you need lots of not just fast parallelism, you need redundancy in order to be able to deal with all these kind of horrible things that can happen to your training.

[01:16:21.135] So there's a distributed systems challenge as well. Gamma 2, which is a Google open source model, has FSDP plus tensor parallel and sequence parallel and uses these two strategies only. And the reason why I brought up Gamma 2 is, as far as I understand, this is basically a realization of the Google claim that for TPUs, really you don't need to do pipelines. You just take a really big toroidal mesh and you tensor parallel over that big mesh network that you have. Unclear whether this can scale out forever. That has always been a little bit less clear to me. But at least the gamma scales, this is definitely the case. Finally, I took a look at some of the models that were trained this year, like Mistral 8x22B, and during my investigation, one of the things that I discovered was actually if you're interested in model training and parallelism configurations in the wild, NVIDIA has a repository called Megatron Bridge, where they release a lot of recommended

[01:17:25.590] training configurations for a lot of different model sizes and model settings. And you can take a look at what they end up saying is the right model configuration for an 8x22B MoE or DeepSeek V3 modeling. And so for Mixtral, they do expert parallel of 8, pipeline parallel of 4, and an additional tensor parallel of 4. This one for the attention layers. And you see this follows the prescription of, keep your expert parallel roughly around eight. Nemotron 3 Super follows the DeepSeek V3 model, has a bunch of expert parallel as well as context parallel, because they were doing extensions, at least in this article. And then the last one that I want to talk about is Qwen 3, which follows the DeepSeek recipe, which is a fairly large amount of expert parallel 32. And they have eight pipeline parallel and two tensor parallel to split up things like the attention matrices.

[01:18:23.790] It is actually interesting to see additionally the different kinds of benchmarking that NVIDIA folks have done for parallelizing different models. And you see, I won't go into details of what this is. This is in the NVIDIA bridge repo, if you want to take a look. There's all sorts of different configurations that you can pick even within tensor parallel that can significantly affect the performance of these. So the systems integration is quite significant. So to put everything together, I decided maybe I would build an overview. Really, the thing that's common to this is as much as you can, all the models use data parallel. They maximize the data parallel domains to the extent that they can. Tensor parallel almost always remains below eight. And expert parallel can sometimes be big now. And I think that's partially because of DeepSeek V3 and a lot of the infra that they built for large scale expert parallel training. Good. OK. So to put everything together, right. The highest level point, right, of this lecture is we need to think about multi-GPU, multi-node, maybe even multi-data center parallelism.

[01:19:31.750] And to think about efficiency in that regime, we need to really bring to bear all the approaches that we have. Because you've got fast links, you've got slow links, you've got techniques that make use of batch sizes and all these different resources that you all want to consume as much as possible. But given all of these, it turns out that there's fairly simple rules of thumb for combining all these different forms of parallelism to get in the end what is effectively full utilization of your compute hardware. Great. And next week, I think we're talking about scaling laws.

Parallelism

课程简介 ​

CS336 Lecture 8：并行化（Parallelism）——中文讲义 ​

学习目标 ​

课程主线 ​

一、并行化的动机与硬件背景 ​

1.1 两个核心瓶颈 ​

1.2 节点内 vs. 节点间 ​

1.3 TPU 与 GPU 的硬件哲学差异 ​

二、集合通信原语回顾 ​

2.1 关键原语 ​

2.2 关键恒等式（本讲反复使用） ​

三、数据并行与 ZeRO ​

3.1 朴素数据并行（DDP） ​

3.2 内存核算：为什么 DDP 糟糕 ​

3.3 ZeRO 三个阶段 ​

3.4 ZeRO-1 与 ZeRO-2 为何“免费” ​

3.5 ZeRO-3 / FSDP：按需抓取参数 ​

3.6 FSDP 的局限 ​

四、模型并行：切分层与矩阵 ​

4.1 核心区别：从通信参数到通信激活值 ​

4.2 流水线并行（Pipeline Parallelism） ​

4.3 张量并行（Tensor Parallelism） ​

4.4 激活值内存的精确核算 ​

五、专家并行（Expert Parallelism） ​

5.1 基本思想 ​

5.2 与张量并行的比较 ​

5.3 关键约束：DP 与 EP 的耦合 ​

5.4 注意力层与 MLP 层的张量并行冲突 ​

六、上下文并行（Context Parallelism） ​

七、并行策略总览与组合原则 ​

7.1 各策略的局限总结 ​

7.2 实用组合策略（简单的经验法则） ​

7.3 性能分析框架 ​

八、实践案例 ​

九、课程作业关联 ​

常见误区 ​

复习清单 ​

从单卡到数据中心：大语言模型训练中的并行化策略全景 ​

导语 ​

两个瓶颈与两种网络哲学 ​

数据并行：最朴素的想法及其局限 ​

ZeRO 与 FSDP：数据并行的内存革命 ​

模型并行：切分层、矩阵与专家 ​

流水线并行 ​

张量并行 ​

序列并行 ​

专家并行 ​

上下文并行 ​

组合策略：从乐高积木到 4D 并行 ​

现实案例：前沿模型的并行化配置 ​

结语 ​

CS336 2026 Lecture 8：完整中文逐字稿 ​

正文 ​

CS336 2026 · Lecture 8 · Parallelism ​