课程简介
CS336 Lecture 7：并行（Parallelism）中文讲义
学习目标
完成本讲学习后，你应该能够：
- 理解分布式训练中的核心集合通信原语（collective operations），包括 all-reduce、reduce-scatter、all-gather 等
- 解释 GPU 集群的硬件拓扑（NVLink、NVSwitch、InfiniBand）及其对通信带宽的影响
- 掌握三种基本并行策略：数据并行（data parallelism）、张量并行（tensor parallelism）、流水线并行（pipeline parallelism）
- 理解 DDP（Distributed Data Parallel）的实现原理：仅需在反向传播后插入一步梯度 all-reduce
- 知道如何测量通信有效带宽，并能解释 all-reduce 与 reduce-scatter 的带宽关系
课程主线
本讲从单个 GPU 扩展到多个 GPU。上周我们关注单个 GPU 内部的优化（融合、分块等），本周则讨论如何协调多个 GPU 之间的数据流动。
核心观点：计算单元与数据之间的距离是核心瓶颈。单 GPU 场景下，数据在 HBM 中；多 GPU 场景下，数据可能在另一块 GPU 上。目标都是编排计算，避免数据传输成为瓶颈。
广义层次结构
| 层级 | 连接方式 | 相对速度 |
|---|---|---|
| 单节点单 GPU | L1 缓存 / 共享内存 | 最快 |
| 单节点单 GPU | HBM | 快（上周认为慢） |
| 单节点多 GPU | NVLink / NVSwitch | 中等 |
| 多节点多 GPU | InfiniBand / Ethernet | 慢 |
上周通过融合和分块减少内存访问；本周通过**复制（replicating）和分片（sharding）**减少 GPU 间通信。
为什么需要多 GPU？
- 容量不足：参数、激活值、梯度或优化器状态无法放入单 GPU HBM。例如 1 万亿参数模型远超单卡容量。
- 加速训练：即使单卡可容纳，仍希望通过拆分利用更多 GPU 的计算能力。
注意：分散到多 GPU 需要付出通信带宽代价，需要计算权衡。
第一部分：分布式通信与计算的构建模块
集合通信操作（Collective Operations）
这些原语源于 1980 年代的并行编程文献，并非为 LLM 训练发明，但至今仍是标准工具。
基本术语：
- Rank：特定设备（本课程中即 GPU）的编号
- World size：设备总数
核心操作：
基础原语（热身）
| 操作 | 描述 | 训练中的直接用途 |
|---|---|---|
| Broadcast | 从 rank 0 复制张量到所有 rank | 初始化时广播 checkpoint |
| Scatter | rank 0 上的大张量按 rank 分割分发 | 理解 reduce-scatter 的基石 |
| Gather | 从所有 rank 收集分片到 rank 0 | 理解 all-gather 的基石 |
| Reduce | 对所有 rank 的数据应用归约（sum/min/max）到 rank 0 | 理解 all-reduce 的基石 |
训练中的核心操作
All-gather：对所有 rank 执行 gather，每个 rank 最终持有完整张量。
Reduce-scatter：对每个维度执行归约，然后将结果分发到不同 rank。
All-reduce = reduce-scatter + all-gather。输入相同，先在每个维度归约，再收集到所有 rank。
关键区别：
- All-reduce 在数据并行中用于梯度求和并复制完整参数
- 更高级方法（ZeRO、FSDP）需要将 all-reduce 拆分为 reduce-scatter + all-gather，以便更精细地管理存储
All-to-all：每个 rank 向其他每个 rank 发送指定张量。对 MoE 训练很重要——动态路由需要根据数据决定激活值送往哪个专家。当切分均衡时，all-to-all 本质上是矩阵转置。
术语记忆技巧
- Reduce：执行结合律、交换律操作（sum、min、max）
- Scatter 是 Gather 的逆操作
- All 意味着目的地是所有设备
硬件：GPU 如何连接
传统家庭式设置
- 同一节点 GPU 通过 PCIe 总线通信
- 不同节点通过以太网连接
- 这种方式不适合严肃的大规模训练
数据中心设置
- 通常每节点 8 个 GPU，通过 NVLink 连接到 NVSwitch
- 参考：NVLink 5 提供约 1.8 TB/s 总带宽；B200 HBM 带宽约 8 TB/s（慢约 4 倍）
- 多个节点组成 pod，pod 间通过 InfiniBand 连接（需经过 PCIe）
- 更大规模经以太网连接（需经过 CPU，最慢）
RDMA（Remote Direct Memory Access）
标准以太网通信需经过 CPU：GPU → CPU 内核 socket 缓冲区 → 网络数据包 → 网络接口。这会引入大量延迟。
RDMA 允许 GPU 直接读写另一 GPU 内存，无需 CPU 介入。NVLink/NVSwitch 环境天然支持 RDMA，InfiniBand 也支持。标准以太网不支持，但 RoCE（RDMA over Converged Ethernet）实现了类似能力。
值得注意的进展
- NVIDIA NVL72：8 GPU/托盘 × 9 托盘 = 72 GPU 组成单一 NVLink 域，实现高速互联
- RoCE：以太网绕过 CPU，作为 InfiniBand 的更低成本替代方案
PyTorch 分布式实现
NCCL（NVIDIA Collective Communication Library）：将集合操作翻译为 GPU 间实际发送的底层数据包。NCCL 检测硬件拓扑、优化路径、启动通信内核。
torch.distributed 提供简洁接口：
- 后端选择：GPU 用
nccl，CPU 用gloo - 支持更高级算法如 FSDP（本课程不使用）
关键示例：
# All-reduce：就地修改张量
dist.all_reduce(tensor=data, op=dist.ReduceOp.SUM)
# Reduce-scatter：输出到单独张量
dist.reduce_scatter_tensor(output=output, input=input, op=dist.ReduceOp.SUM)
# All-gather：输出到张量列表
dist.all_gather_into_tensor(output_tensor=output, input_tensor=input)同步注意：
- CUDA 操作默认异步，需要
torch.cuda.synchronize()确保内核完成 dist.barrier()等待所有进程到达同一点- 异步操作允许计算与通信重叠
通信基准测试与有效带宽
All-reduce 带宽计算：
sent_bytes = size_bytes × 2 × (world_size - 1)
total_duration = world_size × duration
bandwidth = sent_bytes / total_duration- 因子 2：发送和接收
- (world_size - 1)：归约步数
- 随 world_size 增大，(world_size-1)/world_size ≈ 1，带宽 ≈ 2 × size_bytes / duration
关键性质：
- 有效带宽与 world_size 无关（好事）
- 与拓扑无关（NCCL 决定环或树形拓扑）
- Reduce-scatter 无 2 倍因子，因为只需发送一次
由于 all-reduce = reduce-scatter + all-gather，all-reduce 移动 2 倍数据，耗时约 2 倍，但带宽相同。
第二部分：分布式训练
本部分通过极简 MLP 实现讲解三种并行策略。注意：MLP 是 Transformer 中的计算瓶颈，因此具有代表性。
数据并行（Data Parallelism）
策略：按批次维度切分数据，每个 GPU 持有完整模型参数和优化器状态。
流程：
- 每个 rank 持有
batch_size / world_size条数据 - 独立前向传播、计算损失
- 独立反向传播得到不同的梯度
- 关键步骤：对所有参数梯度执行 all-reduce（取平均）
- 更新参数
核心洞察：DDP 与标准训练的唯一区别就是在反向传播后插入一步梯度同步。
# 唯一区别
for param in params:
dist.all_reduce(tensor=param.grad, op=dist.ReduceOp.AVG)性质：
- 各 rank 损失不同（基于本地数据）
- 梯度 initially 不同，经 all-reduce 后相同
- 参数始终保持一致
约束：批次大小需 ≥ world size，最好是 world size 的倍数。
简洁性：DDP 非常模块化——不关心前向传播的具体形式，适用于任意模型。
局限：需要在内存中保存所有模型参数。当模型放不进单卡内存时，需要更聪明的方案（下一讲：FSDP/ZeRO，使用 all-gather + reduce-scatter 而非完整 all-reduce）。
张量并行（Tensor Parallelism）
策略：按宽度维度（特征维度）切分每层参数，每个 rank 持有部分参数。
列并行（Column Parallel）：
- 参数矩阵形状为
num_dim × local_num_dim（按列切分） - 每个 rank 计算部分激活值
- 通过 all-gather 收集所有激活值，拼接为完整维度
与数据并行的区别：
- 数据并行无需改动模型本身
- 张量并行必须修改模型结构（利用矩阵乘法可分解性）
前向传播：每层计算后 all-gather 激活值。
反向传播：对偶操作为 reduce-scatter 梯度。
硬件要求：
- 每层都需传输激活值（数据量大）
- 通常只在节点内部（NVLink/NVSwitch）使用
- 不会跨 NVLink 域使用张量并行
流水线并行（Pipeline Parallelism）
策略：按深度（层）切分网络，每个 rank 获得部分层。
流程：
- 数据批次拆分为微批次（micro-batches）
- rank 0 处理自己的层后，通过点对点
send发送到 rank 1 - rank 1 接收后继续处理，再发送到 rank 2
关键原语：
dist.recv(tensor=x, src=rank-1)：从上一个 rank 接收dist.send(tensor=x, dst=rank+1)：发送到下一个 rank- 异步版本（
isend/irecv）支持通信与计算重叠
微批次动机：
- 减少流水线气泡（pipeline bubble）——等待其他张量处理时的空闲时间
- 更小批次快速处理并传递，减少空闲
缺失内容：
- 通信与计算重叠（异步操作）
- 反向传播的具体实现（作业内容）
硬件要求：
- 容忍较慢互联
- 适合跨节点甚至跨地域场景（去中心化训练）
并行策略选择考量
| 策略 | 通信量 | 硬件要求 | 主要挑战 |
|---|---|---|---|
| 数据并行 | 梯度同步 | 中等 | 批大小上限（临界批大小） |
| 张量并行 | 每层激活值 | 高（NVLink） | 需修改模型 |
| 流水线并行 | 层间激活值 | 低 | 流水线气泡 |
常见组合：节点内张量并行 + 节点间数据/流水线并行。
临界批大小：数据并行中，批大小增长超过某阈值后不再提升效率，应考虑其他并行方式。
课程作业关联
作业二涉及：
- 实现集合通信操作
- 研究通信与计算重叠（特别是反向传播中梯度计算后立即开始通信）
- 实现三种并行策略
常见误区
- 认为 rank 就是 GPU：本课程中 rank 等价于 GPU，但概念上 rank 是逻辑设备编号
- 忽视 CUDA 异步性：CUDA 内核默认异步，必须
synchronize()确保完成，且注意与barrier()的顺序 - 认为 all-gather 与 gather 相同：all-gather 输出到所有 rank，gather 仅输出到指定 rank
- 混淆 reduce-scatter 与 scatter：scatter 仅分发不归约，reduce-scatter 先归约再分发
- 认为带宽与 world_size 相关：有效带宽公式中 (world_size-1)/world_size 收敛于 1，实际与规模无关
- 张量并行用于所有互联：张量并行通信频繁，只应在 NVLink 等高速互联上使用
- 忽视同步屏障的代价：过多 barrier 导致不必要等待
- DDP 计算量重复：每个 rank 都做完整前向/反向，是冗余计算，但避免传输优化器状态
关键公式与原则
有效带宽：
all-reduce: bandwidth ≈ 2 × size_bytes / duration
reduce-scatter: bandwidth ≈ size_bytes / duration集合操作组合：
all-reduce = reduce-scatter + all-gather
gather = scatter 的逆操作
all = 目的地为所有设备并行策略对比：
数据并行：按 batch 切分，同步梯度
张量并行：按 width 切分，同步激活值
流水线并行：按 depth 切分，传递激活值复习清单
- [ ] 能画出广播、散射、收集、归约、全收集、归约散射、全归约、全对全的输入输出图
- [ ] 解释为什么 all-reduce = reduce-scatter + all-gather
- [ ] 描述 GPU 集群的典型拓扑及各级带宽相对关系
- [ ] 解释 RDMA 解决的问题（绕过 CPU）
- [ ] 编写 DDP：只需在反向传播后插入梯度 all-reduce
- [ ] 解释张量并行为何需要高速互联，反向传播对应哪个操作
- [ ] 解释微批次如何减少流水线气泡
- [ ] 推导 all-reduce 有效带宽公式，解释 2 倍因子来源
- [ ] 理解异步操作（async_op=True）如何实现通信与计算重叠
从单个 GPU 到千卡集群：语言模型训练的并行化之路
导语
当我们在单块 GPU 上训练语言模型时，关注的是如何通过融合（fusion）、分块（tiling）等技巧减少内存访问。但如果目标是一万亿参数规模的模型，或者我们希望利用更多 GPU 来加速训练，问题就变成了：如何让成百上千块 GPU 协同工作？
这正是 CS336 Lecture 7 的核心议题。在上一讲中，我们深入 GPU 内部，理解了流式多处理器（streaming multiprocessors）、高带宽内存（High Bandwidth Memory）和缓存层级。现在视野要放宽——从单 GPU 扩展到多 GPU，甚至多节点集群。课程明确指出，这两种情况遵循同样的原则：计算单元与数据相距遥远，核心目标是编排计算，避免数据传输瓶颈。
本文将从集合通信原语出发，介绍 GPU 互联的硬件基础，然后逐步剖析三种基本的并行化策略：数据并行、张量并行和流水线并行。通过这些内容，你将理解分布式训练的核心计算模式，为后续学习更高级的 FSDP/ZeRO 技术打下基础。
第一部分：分布式通信的构建模块
集合通信：分布式编程的经典原语
并行编程并非为语言模型训练而发明。课程强调，集合通信操作（collective operations）的历史可以追溯到 1980 年代，它们是分布式编程的基本原语，今天依然是我们使用的核心工具。所谓"集合"，意味着你指定的是跨多个设备的通用通信模式，而不是管理某块 GPU 如何与另一块 GPU 点对点通信。这让系统能自动完成大量工作。
在描述这些操作前，需要明确两个基本概念：秩（rank） 指一个特定设备（这里是 GPU），世界大小（world size） 指设备总数。假设我们有四个秩，世界大小为 4。
课程将集合操作分为两类：前四个——广播（broadcast）、散射（scatter）、收集（gather）、归约（reduce）——是理解后续操作的基础；后四个——全收集（all-gather）、归约散射（reduce-scatter）、全归约（all-reduce）、全对全（all-to-all）——才是语言模型分布式训练中反复出现的核心操作。
广播（Broadcast） 是最简单的操作：秩 0 拥有某个张量 [0, 1, 2, 3]，广播后将这个张量复制到所有秩上。广播不直接出现在训练核心路径中，但常用于初始化——例如秩 0 加载初始检查点后广播给所有秩。这类操作只需执行一次。
散射（Scatter） 将秩 0 上的大张量按照世界大小分割，然后分发到不同秩上：秩 0 得到第 0 个分量，秩 1 得到第 1 个分量，依此类推。课程指出，散射本身不直接使用，但理解它是理解归约散射的关键。散射的直觉是：让被散射到的 GPU 各自在不同数据部分上做本地计算。收集（Gather）是散射的逆操作——把所有分片拼接回一个秩上，它是理解全收集的基石。归约（Reduce）则对分布在各秩上的数据应用某种结合律和交换律的操作（求和、最大值、最小值），结果放到秩 0。
如果对 NumPy 的广播概念感到熟悉，课程也专门澄清了这一点：概念上都是"一个东西变成很多东西"，但集合通信实现的是跨设备的通信模式，与 NumPy 的数组广播有本质区别。
真正驱动训练的核心操作
全收集（All-gather） 对所有秩执行收集操作，而不仅是秩 0。这是后面反复出现的模式：每个秩持有参数的一部分分片，真正做前向传播时，需要先全收集这些分片，得到完整参数。课程预演了训练中的典型模式：先 gather 做点事情，然后 scatter，再 gather，再 scatter。
归约散射（Reduce-scatter） 在每个维度上执行归约，然后把结果散射到不同的秩。具体来说，假设四个设备分别拥有向量，第一个维度求和得到 6 放到秩 0，第二个维度求和得到 10 放到秩 1，依此类推。课程预告了它的应用场景：反向传播后，不同 GPU 处理了不同数据，需要对梯度分片求和，然后重新分配存储。
全归约（All-reduce） 可以看作归约散射加上全收集的组合：对所有张量做归约（求和），然后把结果复制到所有秩。这是数据并行中最直接的操作。课程指出，更高级的技术如 ZeRO 或 FSDP 需要把全归约拆分成归约散射和全收集，这样才能介入中间过程，精细管理内存。但对基础版本来说，全归约就足够了。
全对全（All-to-all） 是最通用的操作。它精确指定每个秩如何向其他秩发送特定消息，本质上可以理解为矩阵转置（当切分均衡时）。全对全在混合专家模型（MoEs）中至关重要：每个秩既有数据切分，又有一组专家子集，需要根据动态路由将激活值发送到正确的专家。课程强调了负载均衡的重要性——回忆 MoE 那一讲的内容，理想目标就是让 all-to-all 看起来像均等的转置操作。
理解这些术语有个技巧：归约就是对结合律交换律的操作（求和、最值）进行归约；散射是收集的逆操作，散射是分发，收集是集中；"全"（all）意味着目的地是所有设备。
GPU 互联的硬件现实
课程从一张"老旧的经典图片"讲起传统服务器架构：CPU 通过 PCIe 总线连接 GPU 和内存，不同计算机通过以太网相连。这种拓扑下，同一节点的 GPU 走 PCIe 通信，跨节点必须经过以太网——这一切就像买了游戏 GPU 后和朋友组网来训练大模型。
认真做训练的数据中心则完全不同。典型设置是每个节点有 8 块 GPU，通过 NVIDIA 的 NVLink 连接到交换机（NVSwitch）。课程给出了参照：NVLink 5.0 能提供每秒 1.8 TB 的总带宽，而 B200 的 HBM 是每秒 8 TB——大约慢了 4 倍。这仍然很快，但绝对不如 HBM，更远不及共享内存或 L1 缓存。
从编程角度看，NVLink 连接到交换机意味着任意 GPU 可以认为能直接与任意其他 GPU 通信，硬件负责路由。但随着规模增长，节点组成 pod，pod 之间通过 InfiniBand 连接，GPU 需要经过 PCIe、InfiniBand 线缆，速度大幅下降。再往外是以太网，需要经过 PCIe 和 CPU，速度更慢。这正好呼应了内存层级的直觉——节点越多，通信越慢。
课程特别强调了一个关键概念：RDMA（远程直接内存访问）。标准以太网通信中，GPU 必须把数据拷贝到 CPU 的 socket 缓冲区，构建网络包，再拷贝到网络接口发送，引入大量延迟。RDMA 允许 GPU 直接读写另一块 GPU 的内存，完全不经过 CPU。NVLink/NVSwitch 环境天然支持 RDMA，InfiniBand 也支持。但标准以太网不支持。
两个值得一提的进展：一是 NVIDIA 的 NVL72 产品——9 个托盘、每个托盘 8 块 GPU，总共 72 块 GPU 可以通过 NVSwitch 连入一个 NVLink 域，让普通人也能获得大范围的快速互联。二是 RoCE（融合以太网上的 RDMA），以太网域中的 RDMA 方案，绕过了 CPU，是 InfiniBand 的廉价替代品，Meta 有论文展示了这一方向的探索。
在编程层面，NCCL（NVIDIA 集体通信库） 把集合操作翻译成 GPU 之间发送的实际底层数据包。课程指出，NCCL 会检测硬件拓扑，找出不同 GPU 之间的路径，启动 GPU 内核来发送接收数据——因为 GPU 上的一切最终都是内核。我们不会深入 NCCL 内部，但需要知道它存在。PyTorch 提供了 torch.distributed 库，封装了这些集合操作的简洁接口，支持不同后端：GPU 上用 NCCL，CPU 上用 gloo（并行处理在 GPU 出现之前就存在了，CPU 上也能执行集合操作）。
用代码验证概念
课程通过实际代码演示了这些操作。这里的关键设计是：直接运行课程代码会使用多进程（multiprocessing），每个进程对应一个秩；但在课堂逐步演示时，使用的是特殊的单进程模式，避免多进程的复杂性。
第一个示例是全归约。每个秩创建一个张量 [0, 1, 2, 3] 加上自己的秩号，所以秩 0 有 [0, 1, 2, 3]，秩 1 有 [1, 2, 3, 4]，依此类推。全归约后，每个秩得到的都是各列之和，结果在所有秩上完全一致。课程特别提醒：打印语句出现的顺序取决于硬件，因为进程是异步运行的。这里引入了 barrier（屏障） 同步操作——确保所有进程都到达某个执行点后再继续。过度使用屏障会带来不必要的等待。
第二个示例是归约散射。输入是各秩手中的向量，执行归约散射后，输入保持不变，但输出张量在每个秩上得到了对应分量归约后的结果。第三个示例是全收集：以上一步归约散射的输出作为输入，分配输出张量后，执行全收集，所有不同的输入被收集到所有秩上——这正好通过实例验证了"全归约 = 归约散射 + 全收集"。
测量通信带宽
课程接着演示了如何测量通信速度。这个环节延续了计算 MFU（模型浮点运算利用率）时的思路：测量有效带宽。对一亿个元素做全归约，预热后开始计时。这里需要考虑两种异步性：CUDA 内核相对于进程是异步的，进程之间也是异步的。因此既要调用 torch.cuda.synchronize() 等待 CUDA 内核完成，也要调用 barrier 等待所有进程到达。
有效带宽的计算方式是：发送的字节数除以总持续时间。对于全归约，实际发送的字节数需要一些解释。假设秩 0 加秩 1 加秩 2 加秩 3，需要迭代世界大小减 1 步（因为有世界大小减 1 次加法），因子 2 是因为既需要发送又需要接收（归约）。总持续时间是墙上时钟时间乘以世界大小——因为所有秩都在等待。课程指出一个有趣的性质：随着世界大小增加，表达式约等于 2 乘以数据字节数除以持续时间，也就是说，有效带宽与世界大小无关，也与拓扑结构无关——环形还是树形由 NCCL 决定。
对于归约散射，带宽计算没有 2 倍因子。课程强调：由于全归约等于归约散射加上全收集，全归约自然移动两倍的数据量，需要两倍的时间，但两者相抵消，得到相同的带宽。
第二部分：三种并行化策略
课程进入核心部分：通过一个简洁的 MLP 训练实现，从零开始讲解数据并行、张量并行和流水线并行。之所以选 MLP，是因为 Transformer 中 MLP 才是计算瓶颈所在，这非常能代表真实情况。这里所有的并行化都通过切分数据和参数来实现。课程给出的示意图是概念性的：数据并行切分数据，张量并行切分宽度（每层的一部分），流水线并行切分深度（层数）。
数据并行：最优雅的起点
课程用"数据并行就是切分批次维度，每个 GPU 保留所有参数，做正常的模型训练，然后同步"这句话精炼概括了核心思想。具体实现中，先从 128 行、1024 维的数据矩阵出发，按世界大小切分行——每个秩获得 32 个数据点（本地批次大小 local batch size）。在实践里，每个秩应该自己加载自己的数据，而不是通过主进程加载再分发——课程只是为了说明问题才这样做。
MLP 有若干层，每层是一个 1024×1024 的参数矩阵。训练循环是标准的：前向传播遍历层，应用激活函数；反向传播计算梯度；关键是反向传播之后多了至关重要的一步：对所有参数做全归约（ReduceOp.AVG 均值）。这是标准训练和 DDP（Distributed Data Parallel）之间唯一的区别。课程认为这非常优雅："基本上就是标准训练，把训练应用到本地批次上，但只是在反向传播之后插入了这一步。"
课程解释了数据并行的两个前提：批次大小至少要等于世界大小才有意义（实践中通常远大于），最好能整除世界大小，否则可能需要用零填充。
数据并行的一个显著优势是模块化：DDP 只关心对参数做平均，并不关心前向传播长什么样。因此对 Transformer 模型同样适用。但它的局限也很明显：需要在内存中保存所有模型参数。如果参数放不进内存，就需要更聪明的方式——这正是下一课 FSDP 和 ZeRO 的主题。
分析师们会注意到，数据并行中每个秩的损失不同（基于本地数据计算），梯度最初也不同，但通过全归约，梯度在所有秩间变得相同，因此参数也始终保持一致。这就是数据并行的收敛保证。
张量并行：切分每层，收集激活
与数据并行不同，张量并行不切分数据，而是切分每一层。每个秩获得每层的一部分参数，因此需要在每层之后通信激活值。课程采用按列切分（称为列张量并行 column tensor parallel），也可以按行切分。
具体实现中，本地维度数量（local num dim）是每个秩负责的维度子集大小。前向传播中，每个秩计算自己那部分参数对应的激活值（形状为批次大小×本地维度），然后需要把所有这些激活值放到所有秩上——这正是全收集操作的应用场景。简单来说：每层计算后，分配激活值列表，执行全收集，拼接成完整维度的数据，继续下一层。
张量并行的一个关键区别是：它必须修改模型本身，而不像数据并行那样把模型当作黑盒模块。它利用的核心事实是：矩阵乘法可以拆分成一组更小的矩阵乘法，让不同秩各自计算再收集结果。
关于反向传播，课程明确了一个要点：如果在张量并行中直接调用 .backward()，并不会自动做所需操作，因为其中没有并行性。实践中有高层封装自动处理这些，但课程刻意保持低层次，让学习者理解机制。反向传播中，梯度的分布对应着前向传播的逆操作：通过归约散射将梯度分发到所有设备。因此可以说，某种意义上的对偶关系存在于前向和反向：前向是全收集，反向是归约散射。
流水线并行：切分深度，减少气泡
流水线并行的思路是按深度切分网络：每个秩获得一部分层，但每层内部拥有所有维度。所有秩以某种形式看到所有数据。
实现的核心是**微批次（micro-batches）**概念。除拆分层外，还要把批次拆分成多个微批次。具体流程是：秩 0 获得数据，分块成微批次；对每个微批次，从前一个秩接收（receive），只在分配给本秩的层上做前向传播，然后发送（send）给下一个秩。代码中使用点对点操作 dist.recv 和 dist.send，它们比集合操作更直接但同样直观。
引入微批次的原因正是为了解决流水线气泡（pipeline bubble）问题：当等待其他张量被处理时，GPU 处于空闲状态。把数据分解成更小的批次，就能快速处理并发送到下一个秩，减少气泡数量。课程承认，这个朴素的实现没有处理通信与计算重叠的问题，而这在实际流水线并行中至关重要。理想的做法是在计算的同时接收或发送数据——给操作加上异步标志并编写更多管理代码。
课程还指出流水线并行中常见的范式：通信与计算重叠。数据并行中同样存在这种机会，一旦梯度计算完成就可以开始发送，不必等所有梯度都算完。这是作业二将探索的内容。
并行化策略的选择与缺失环节
除了上述三种，还有更多并行化类型：序列并行（切分序列，允许并行化注意力计算）、专家并行（为 MoE 并行化专家，全对全在此发挥作用），以及不同技术的组合。
并行化技术的选择强烈依赖于硬件互连能力。张量并行通信量大——每一层都需要发送相当规模的激活值，因此通常局限在节点内部使用 NVLink/NVSwitch 这类高带宽环境，不会跨越 NVLink 域。流水线并行则能容忍慢得多的互连，甚至可以在跨地域的去中心化训练中工作。典型的组合是：节点内部用张量并行，节点间用数据并行（或 FSDP），需要时再叠加流水线并行。
数据并行也有其极限：当批次大小增大到某个临界批大小后，继续增加不再有帮助，大量算力浪费在同步上，此时更好的选择是转向张量并行。这些决策需要计算和权衡。
课程最后对比了 PyTorch 与 JAX/TPU 生态的哲学差异。PyTorch 路线让你用原始集合操作从机制上看到发生了什么；JAX 路线则可以简单地定义模型和起始策略，编译器负责处理需要什么样的通信操作——你只需说明"这份数据需要放在这里、这里还有这里"，编译器施展魔法。这很吸引人，但课程带着一点幽默说："它会夺走很多从零开始构建东西的乐趣。"
结语
回到课程开始提出的问题：为什么需要多 GPU？答案是双重的——参数、激活值、梯度或优化器状态无法容纳在单 GPU 内存中；即使能容纳，拆分到更多 GPU 上也能利用更多算力训练更快。但代价是通信带宽的开销，这需要仔细计算如何并行化。
本讲建立了一个统一的视角：在单 GPU 和多 GPU 的情况下，原理相同——计算单元与数据相距遥远，核心是编排计算以避免传输瓶颈。上周用融合和分块减少内存访问，本周用复制和分片减少 GPU 间通信。这是一个从单 GPU 到多 GPU，从快速缓存到慢速网络的完整层级结构。
值得一提的是，课程最后给出了硬件发展的长期判断：尽管硬件在变快，但我们总是想要更大的模型，因此这种层级结构的思想会持续存在。正如内存层级中我们有 L1、L2、HBM、主存，"你要么重新计算，要么存储在内存中，要么存储在另一块 GPU 的内存里并通过通信获取"——这不仅是技术选择，更是设计哲学。
下一讲将由 Tatsu 深入讲解更先进的数据并行技术——FSDP 与 ZeRO。它们用全收集和归约散射的组合替代简单的全归约，在保留数据并行优雅性的同时，规避了"所有参数都在内存中"的限制。理解本讲的集合原语与三种基本并行化策略，正是迎接这些高级技术的最佳准备。
官方资源：
- 视频：CS336 Lecture 7 - Parallelism
- 讲义与代码：CS336 课程网站（Lecture 7: Parallelism）
CS336 2026 Lecture 7：完整中文逐字稿
- 视频：官方视频
- 来源：英文人工字幕提取后翻译为中文；保留段落起始时间戳，只修正明显的断句、重复和语序。
- 说明：这是逐字稿，不是讲义摘要；相邻且语义连续的片段已合并为阅读段落；术语保留必要英文。
正文
[00:05.160] 好的，我们开始吧。欢迎大家回来。今天我们要讨论并行（parallelism）。还记得上周，我们介绍了如何通过编写内核（kernels）让单个 GPU 运行得更快。我们确实深入研究了 GPU 的内部结构。而本周，我们将讨论如何利用多个 GPU 让你的代码运行得更快。所以，你脑海中应该有这样的画面。上周，我们专注于其中一个盒子，也就是你拥有的那个 GPU。记住，你有高带宽内存（High Bandwidth Memory, HBM）、L2 缓存、L1 缓存、寄存器，以及一堆流式多处理器（streaming multiprocessors）。
[00:51.680] 现在，这个画面扩展了，因为你不是只有一个 GPU，你可能有四个，也可能有 1000 个 GPU。这些 GPU 将被连接起来。稍后我会谈到这些 GPU 是如何连接的。然后你将需要弄清楚如何利用所有这些计算资源来训练模型。所以在这两种情况下，也就是在单个 GPU 和多个 GPU 的情况下，正如我们今天要讨论的，如果你把视野放宽，情况是类似的：计算单元，即算术逻辑单元、张量核心（tensor core）等等，与你的数据相距很远。对于单个 GPU 来说，“很远”意味着数据在 HBM 里，离计算单元有一段距离。现在如果是多 GPU，你这里需要的数据可能远在另一块 GPU 上，你得想办法把它传过来。但基本原则是相同的，因为核心目标都是编排计算，努力避免数据传输瓶颈。
[01:55.280] 使用大量 GPU 很容易，但有效地使用它们却很难。所以，略微自由地讲，我们可以考虑一个广义的层次结构：在最底层，紧挨着 SM（流式多处理器）的是单节点单 GPU，这里有 L1 缓存/共享内存（shared memory）。这是最快的。然后是 HBM，我们曾感叹它太慢了。但在本节课中，HBM 将被看作是快的。现在我们要考虑单节点多 GPU 的情况，其中 GPU 通过 NVLink 和 NVSwitch 连接。最后是多节点多 GPU 的情况，这时我们不得不根据你所拥有的网络，求助于 InfiniBand 和以太网（ethernet）。所以上周，我们讨论了各种减少内存访问的技巧，比如融合（fusion）和分块（tiling）。读入共享内存，尽可能多地计算，然后再写回去。
[03:05.990] 而本周，我们将讨论如何通过适当地复制（replicating）和分片（sharding）来减少 GPU 之间的通信量。那么，为什么要用多 GPU 呢？显而易见的答案是，你想要扩展规模。但更精确地说，实际上有两个原因。一是你的参数、激活值（activations）、梯度（gradients）或优化器状态（optimizer state）无法容纳在单个 GPU 的 HBM 内存中。比如 B21 有 192 GB 内存。如果你要训练一个 1 万亿参数的模型，那肯定放不进单个 GPU。另一个原因是，即使你的模型能放进单个 GPU，你可能还是希望通过拆分所有东西来利用更多 GPU，
[03:54.630] 以便更快地训练。所以有时需要做出一些决策，因为你可以把所有东西都放在 GPU 上，但那样核心数较少。但如果你把它分散开，就得付出通信带宽的代价。所以你需要做一些计算，来弄清楚如何并行化。这里有一点需要注意：到目前为止，本节课的代码是 Python 代码。你执行它，就能看到所有输出。现在，在本节课中，如果你直接运行这些代码，它会使用多进程（multiprocessing）。但当我逐步讲解时，我使用的是某种特殊的单进程模式。所以，如果你想在标准输出中看到本节课在多进程设置下的运行结果，可以点击这里。我会在讲课过程中展示这个。但请记住，当我们逐步讲解本课时，我们实际上并没有进行多进程操作，因为我们只是一行一行地执行代码。
[04:56.990] 本节课将包含两个部分。第一部分，我们将学习分布式通信和计算的构建模块，从编程模型开始。简单谈谈硬件。然后开始在 Torch 中实现一些东西，这将是你作业二的内容。第二部分，我们将看看实际的训练。我们将研究三种类型的并行：数据并行（data parallelism）、张量并行（tensor parallelism）和流水线并行（pipeline parallelism）。每种方式都会以不同的方式切分我们的模型。我们会用 MLP（多层感知机）而不是完整的 Transformer 来做这些。但这里展示的确实是核心计算。让我们开始吧。首先要讨论的是所谓的集合通信操作（collective operations）。集合通信操作是分布式编程中的原语，其历史可以追溯到 80 年代。
[06:03.200] 所以并行编程的概念非常古老。它并非为 LLM 训练而发明。而且时至今日，我们使用的仍然是这些原语。在这里，“集合”（collective）意味着你指定的是跨多个设备的通用通信模式或模板，而不是管理这个 GPU 如何与另一个 GPU 进行点对点通信。这样会容易得多，系统可以为你做更多的工作。所以这是一个非常成熟可靠的并行编程接口。一般的设置如下。这里的术语有些——我觉得有点奇怪，但这是并行编程的标准用法。核心思想是你有一堆秩（rank），其中秩对应于一个特定的设备，在我们的案例中是指 GPU。它也可能是 TPU。但关键是，假设这里有四个秩。世界大小（world size）对应于设备的数量，所以这里的世界大小是 4。
[07:13.060] 我们将介绍几个操作：广播（broadcast）、散射（scatter）、收集（gather）、归约（reduce）、全收集（all gather）、归约散射（reduce scatter）、全归约（all reduce）和全对全（all to all）。这些操作中的每一个都指定了这组秩或设备如何将一定量的数据和/或计算传输给另一组设备。前三个操作，广播、散射、收集、归约，我认为实际上只是热身。它们能让你了解这些集合通信操作是如何工作的，但它们并不是驱动大部分训练的核心操作。全收集、归约散射和全归约才是那些在语言模型分布式训练中反复出现的操作。
[08:05.080] 最后，我会在这里提一下全对全，它对混合专家模型（MOEs）很重要，但本节课我们实际上不会花太多时间在上面。好的，让我们从最简单的操作开始，即广播。在广播中，你有一个秩 0。可以是任何一个秩，但为了方便举例，我们就选秩 0，它拥有某个张量 0123。它把这个张量广播给所有的秩。所以在这个操作结束时，每个秩上都拥有了相同的张量。这应该相当直接明了。而且，这个操作同样不会真正出现在训练的核心路径中。通常，广播用于初始化，例如，你初始化或加载了一个初始检查点（checkpoint），然后将其广播到所有秩。这是那种只做一次的操作。
[09:10.260] 第二个操作是散射。散射基本上是说，我在秩 0 上有一个张量，它被按照世界大小分割开。然后我将把我的张量散射到其他秩上。所以秩 0 得到第 0 个分量。秩 1 得到这个，秩 2 得到这个，秩 3 得到那个。再强调一下，这个并不是直接使用的，但 scatter 是理解 reduce scatter 的重要基石。顾名思义，scatter 就是把一个大的张量放在一个地方，然后把它分散到多个地方去。你可以看到这为什么会很有用，因为你希望被散射到的所有 GPU 都能在不同的部分上做一些本地计算。那么 scatter 的逆操作就是 gather。所以这应该非常容易预测。输入是你有一堆分片，每个分片分别位于特定的 rank 上。然后当你做 gather 时，这是相对于某个特定 rank（比如 rank 0）而言的，它会把所有分片拼接在一起。
[10:38.170] 同样，gather 也不是直接使用的，但它是理解 all gather 的基石。接下来是 reduce。你们中做过函数式编程的人可能对 reduce 很熟悉。这里完全一样。它的思路是，起点和 gather 一样。张量被分割到多个——嗯，你在每个不同的 rank 上都有一些数据。然后你对所有这些数据应用你的归约操作，把结果放到 rank 0 上。在这个例子中，如果你用求和来做归约，那么把这些加起来，你得到 6。所以你可以把 gather 看作是一种归约，只不过操作是拼接，如果你愿意这么理解的话。当然，reduce 对理解 all reduce 很重要。让我在这里停一下。刚才那些只是热身。以防大家对什么是集合通信操作有疑问，broadcast、scatter、gather、reduce。
[11:59.850] 嗯？[听不清] NumPy，在哪里？所以问题是，这和 broadcasting 以及 NumPy 有关系吗？我觉得概念上是相同的想法，就是一个东西变成很多东西。比如在 NumPy 里，如果你有一个标量，可以广播成一个张量。但这里的具体实现是用于集合通信的。所以有点不同。让我们继续看一些更有趣的内容。all gather 就是——基本上，你对所有 rank 执行 gather，而不仅仅是对 rank 0。还记得 gather 是做什么的吧。它基本上把所有的不同分片拿过来，放到一个 rank（rank 0）上。现在 all gather 就是对每一个 rank 都做这件事。这就是 all 的含义。all 意味着把输出给到所有 rank，而 gather 是你对所有这些 rank 做的操作。
[13:12.560] 这个后面会经常出现。现在不需要精确理解这句话，但稍后我们会看到每个 rank 持有参数的一部分。然后你需要做的是 all gather 这些参数，得到完整的参数来进行完整的前向传播。总的来说，在我们训练的过程中，我们会看到很多这种先 gather 做点什么，然后 scatter，然后再 gather 再 scatter 的模式。reduce scatter 是在每个维度上执行归约，然后把结果散射出去。假设你有四个设备，每个设备都有某个向量。之前我们做 reduce 的时候，就是 0 1 2 3，然后归约成 6。但现在 reduce scatter 是说，对于这个张量的每一个分量，我要做一个归约，然后把它放到不同的 rank 上。所以第一个维度，我把这些加起来，得到 6。
[14:33.625] 现在，对于第二个维度，我把这些加起来，得到 10。对于第三个维度，我处理这些。对于第四个维度，我处理那些。这里先预告一下它会在哪里出现：在反向传播之后，当你需要求和时，你要做的是每个 GPU 处理不同的数据。你需要做的是把所有不同分片的梯度加起来。然后你要重新分配这些存储。最后是 all reduce。如果你理解了 all reduce gather 和 all gather，那基本上就是先做一个再做另一个。所以这个操作做的是 reduce scatter，输入和之前一样。记住，在 reduce scatter 中，6、10、14、18 分别位于不同的 rank 上。而 all reduce 中的 all gather 部分就是把所有东西放到同一个——所有内容放到同一个节点上。所以 all reduce 在某种程度上是最容易理解的。
[15:59.620] 你有一堆张量。你做归约，这里是求和。然后把结果复制到所有节点上。我们实际上会先遇到这个，在做数据并行时，我们对梯度求和，然后复制完整的参数。这就是我们实际要开始的地方。所以或许把注意力集中在 all reduce 上就好。后面我们会看到如何实现更高级的东西，比如 ZeRO 或 FSDP。我们需要把 all reduce 拆分成 reduce scatter 和 all gather，因为这样你才能介入其中，更精细地管理一些事情。但对于基础版本来说，all reduce 就够了。最后是 all to all。这个在某种程度上是最通用的。你基本上指定每个 rank 如何向另一个 rank 发送特定的消息。这里有一个简单的例子，输入和之前一样。
[17:09.310] 这里说的是，我想把 0 这个元素发送给 rank 0，也就是留给自己。我要把 1 发送给 rank 1，把 2 发送给 rank 2，把 3 发送给 rank 3。然后如果我是 rank 1，我想把 4 发送给 rank 0，5 发送给 rank 1，6 发送给 rank 2，7 发送给 rank 3。所以基本上，这里的位置表示哪个 rank 是最终目的地。如果你看输出，发生的是第一个 rank 会收到所有 rank 发送到第 0 列的内容，因为这些 rank 都在把这些东西发送给 rank 0。
[18:03.510] 类似地，对于 rank 1，所有 rank 都在把这一列发送给 rank 1，依此类推。这在训练 MoE 时会很有用。这里的直觉是，每个 rank 既有数据的切分，也有一组专家的子集。基本上，MoE 的关键思想是动态路由。你必须查看你的数据来确定需要把这些激活路由到哪些专家。所以最终就变成了 all to all 通信。如果一切都很均衡，也就是说每个 rank 向其他每个 rank 发送相同数量的字节，那么 all to all 可以看作本质上是一个转置。如果你把这个看作一个矩阵，你做的就是转置这个矩阵。但一般来说，all to all 也能处理不均衡的切分。我这里没有展示，但你可以配置它向任何其他 rank 发送任意数量的字节。但一般来说，你希望切分尽可能均衡。
[19:27.290] 还记得 [听不清] MoE 那节课里我们做负载均衡来确保事情尽可能均衡吗。所以从道义上讲，理想目标就是让 all to all 看起来像这样。总结一下，也许有一些有用的技巧来记住这些术语，因为我刚才讲了不少不同的操作。reduce 就是，嗯，它做归约。它对某个结合律和交换律的操作进行归约。可以是求和，可以是最大值，可以是最小值。scatter 是 gather 的逆操作。scatter 是分发，gather 是集中。而 all 只是意味着目的地是所有设备。这就解释了 all reduce 和 all gather。让我在这里停一下，看看大家对于集合通信有什么问题。嗯？对于像 gather 这样的操作，基本上 [听不清] 到 rank 0，那个 rank 0 [听不清]？它每次都是特定的某个 GPU 吗，还是说 rank 0 可以变化？
[20:47.680] 是的。所以问题是，当你做 gather 或 reduce 时，写入输出的目标位置，目前我说的是 rank 0。你稍后会在代码中看到，基本上你指定 GPU ID 或 rank，数据就会去到那里。所以它不一定需要提前很久确定，但基本上在你执行调用时必须要确定。好的。还有其他问题吗？是的。这些只是概念性的构建模块，对吧？它们实际上并不是[听不清]。这些是实际的[听不清]吗？是的，所以问题是，这些只是概念性的构建模块还是代码？我现在展示这些只是作为概念性的构建模块，但我们很快就会看到它们在代码中是如何实现的。
[21:53.460] 在进入代码之前，我想先谈谈硬件，特别是 GPU 是如何连接的，因为我们已经知道了 GPU 内部的结构。所以让我们来谈谈一般的网络。这是一张非常经典的图片。你可以从这张看起来很老旧的图片中看出，这就是计算机一般来说是如何工作的。所以你有一台服务器，然后有一堆 CPU。有一条 PCIe 总线，你用它来连接——或者说过去用来连接你的鼠标和键盘之类的东西。然后你有一堆 GPU 挂在上面，还有一些内存。然后这台计算机通过以太网连接到另一台计算机，以此类推。所以这是一个特定的设置。这是一个特定的拓扑结构。同一节点上的 GPU，你使用 PCIe 通信，不同节点上的 GPU，你必须一直通过以太网。所以这就像你买了游戏 GPU，然后和你的朋友连起来，说，我要训练一个大型模型。
[23:12.500] 那你就得这么做。但如果你真的很认真地想做训练，那么情况看起来更像是这样的。这是我在最开始展示的图片，你有 GPU，还有叫做 NVLink 和 NVSwitch 和 InfiniBand 的东西。所以典型的设置是这样的。这些数字，8 是典型的，但这个 256 有点是编的。所以通常来说，每个节点有八个 GPU。它们通过 NVIDIA 的 NVLink 连接到一个交换机。作为参照，如果你使用 NVLink 5，那么你能获得每秒 1.8 TB 的总带宽。记住，B200 的 HBM 是每秒 8 TB。所以大约慢了 4 倍。我的意思是，如果你考虑设备之间的通信，这仍然算是很快的，但显然不如高带宽内存快，而高带宽内存又比共享内存或 L1 缓存慢得多。
[24:28.450] 所以基本上，NVLink 连接到交换机，这意味着从编程的角度来看，你可以认为 GPU 连接到任何其他 GPU。你从任意 GPU 到任意其他 GPU，硬件负责将数据传输到交换机，然后交换机进行路由。所以通常你还会遇到的情况是，在某个时刻，你不能一直用 NVSwitch 和 NVLink。因为随着你的 GPU 数量增长，你需要把这些节点放进 pod 中，pod 之间通过 InfiniBand 连接。InfiniBand 的工作方式是这样的，现在会稍微多一点——现在 GPU 不再直接连接到另一个 GPU。
[25:25.390] 它必须经过 PCIe，然后通过这种特殊的 InfiniBand 线缆。你会看到速度要慢得多得多。最后，如果你用完了 InfiniBand 的容量，并且你有这些巨大的 pod，那么你需要通过以太网连接它们。在以太网中，你必须经过 PCIe，而且实际上要经过 CPU，正如我们将看到的，这甚至更慢。所以这有点类似于内存的情况。你拥有的节点越多，速度就会越慢。你不可能让一个 NVSwitch 处理 100,000 个 GPU。有一点需要注意的是，我之前提到了绕过 CPU，从硬件的角度来看这将是一个重要的事情。
[26:31.730] 所以如果你使用传统的以太网，发生的情况是 GPU 必须与 CPU 通信来拷贝数据。所以它基本上必须把数据拷贝到这个——CPU 有一个内核 socket 缓冲区。这里的内核指的不是 GPU 内核，而是 CPU 的、传统的操作系统内核概念。然后它需要构建一些网络数据包，拷贝到网络接口，然后发送出去。这通常会引入大量延迟。所以有一种叫做远程直接内存访问的技术，即 RDMA，它允许 GPU 直接写入或读取另一个 GPU 的内存，完全不需要使用 CPU。所以显然，如果你在 NVLink 和 NVSwitch 的环境中，那么你就有了 RDMA。InfiniBand 也支持 RDMA。所以如果你通过 InfiniBand 连接，那么你可以让 GPU 直接互相连接而不涉及 CPU。但标准以太网不支持。
[27:42.360] 我要提到两个值得注意的进展。NVIDIA 一直在大力推进在越来越大的 pod 上能做些什么。他们有基本上是四个 B200 和 B300 的产品。他们有一个叫 NVL72 的东西，意思是他们有这种八个 GPU 的托盘，一共九个。所以基本上，最终你有 72 个 GPU，它们都通过 NVSwitch 连接到一个 NVLink 域中。如果你还记得，NVLink 的速度非常快。所以通常来说，如果你是普通人，你会想，好吧，我有八个 GPU 可以实现非常快速的高速互联。然后超出这个范围，速度就会大幅下降。但如果你有很多钱，你可以购买这种非常高端的硬件，你可以获得高达 72 个 GPU 的快速互联。另外我要提到的另一件事是，我之前说标准以太网不支持 RDMA，但在以太网领域也有进展。
[28:55.100] 有一种叫做 RoCE 的东西，即融合以太网上的 RDMA（RDMA over converged ethernet），在这种技术中，以太网实际上绕过了 CPU。这可以说是他们对 InfiniBand 的回应。InfiniBand 通常非常昂贵，NVIDIA 的很多产品也是如此。但你可以通过使用融合以太网上的 RDMA 获得相当不错的性能。Meta 有一些论文展示了他们正在探索这一点。所以[听不清]可能训练在融合以太网上，也可能不是。这就是对硬件概况的简要介绍。你有 GPU，它们通过 NVLink 连接到 NVSwitch，处于某个域中，可能是 8 个，可能是 72 个。
[29:55.440] 然后从那里开始就是 InfiniBand。现在让我们谈谈如何编程实现这一点。在最底层，有一个叫做 NVIDIA 集体通信库的东西，即 NCCL，或者叫"nickel"，它把集体操作——all reduce、reduce broadcast——翻译成 GPU 之间发送的实际底层数据包。所以 NCCL 做的是，当你使用 NCCL 时，基本上就像是在说，我要做 all reduce。然后 NCCL 去搞清楚硬件的拓扑结构，找出不同 GPU 之间的路径，然后它实际启动 GPU 内核来发送和接收数据。因为归根结底，记住，所有在 GPU 上运行的东西都是内核。所以也有通信内核，它们实际执行与其他 GPU 的通信。我们不会深入研究 NCCL，但只需要知道它存在。
[31:04.710] 然后我们实际上要去看 PyTorch。也许在这之前，关于硬件有什么问题吗？是的。你能描述一下——有一个机架和一个托盘。我能从物理上描述一下机架和托盘吗，比如机架是什么，托盘是什么？对于 NVL72，我不是硬件专家，但机架，我的意思是，字面上就是，你见过数据中心，就是一个字面意义上的机架。每个托盘（tray）上都有——G 代表的是 [？好？]。所以上面有两个 CPU。每个 CPU 连接到四个 GPU。所以每个托盘上有八个 GPU。它们堆叠在一起，所有东西都连接到这个 NVSwitch 上。是的。[听不清] 是的，所以问题是，[听不清] 之间的区别是什么。
[32:19.000] [听不清] 是的。所以 RDMA，你可以把它更多地看作是一种需求特性。RDMA 意味着一个 GPU 可以直接从另一个 GPU 的内存中读取和写入数据。实现 RDMA 有多种方式。一种是使用 NVLink 和 NVSwitch，另一种是使用 InfiniBand。所以 InfiniBand、NVSwitch 和 NVLink 更多的是硬件层面，涉及哪些部件、哪些线缆和交换机。而 RDMA 更多是操作层面，涉及通信时发生了什么。[听不清] 是的。例如，这种 [？进步？] 融合以太网上的 RDMA（RDMA over converged ethernet）是另一种实现 RDMA 的方式。[听不清] 是吗？NCCL 是否为多节点集群做了优化？因为显然，这看起来是很直观的 [听不清]。
[33:43.350] 但对于基于 RDMA 的进程，NCCL 会是那个最优选择吗？所以问题是，NCCL 是否为多节点集群做了优化？我不太清楚他们在优化方面的具体细节。我只能说，NVIDIA 基本上一直在优化他们的整个技术栈，用于这些大型模型的推理和训练，因为他们的主要客户是主要的大型语言模型提供商。所以如果他们没考虑为这类工作负载做优化，我会感到惊讶。是吗？[听不清] 所以问题是，如果你有九个 GPU 会怎样？你如何在它们之间分配工作负载？我想这取决于在你的设置中这九个 GPU 是怎么分布的。例如，很多时候你会遇到——假设每个节点有八个 GPU。那么第九个就会在不同的节点上。如果它们之间没有 NVLink 连接，那情况会很糟糕，因为这个节点既没有提供多少计算能力，而且与之通信的成本也非常高。
[35:23.860] 但如果，比如说，所有东西都由 NVSwitch 连接，那就会合理得多。再问一个问题，然后我就要继续了。那这和 TPU 有什么不同？所以 [听不清] 之前稍微描述了一下 TPU。我的意思是，让我想想。关于这个我能快速说点什么？TPU 通常是简单得多的东西。我不太清楚这些组件各自对应什么细节，但也许我们可以在课后讨论。那么让我们来看一些实际的代码，来利用这些硬件。PyTorch 很方便地提供了一个 torch distributed 库，它为这些集合通信操作提供了简洁的接口。所以你不需要显式地去考虑 NCCL。事实上，这个库还支持针对不同硬件的不同后端。如果你在 GPU 上，那么你会使用 NCCL 后端。如果你在 CPU 上，那么有一个叫 gloo 的东西，它仍然允许你——就像我说的，并行处理在 GPU 出现之前很久就已经存在了。
[36:54.900] 你也可以在 CPU 上执行这些集合操作。这个库还支持更高级别的模型和算法，比如 SFTP，但我们课程中不会用到那些，因为我们是从零开始构建。好了。让我们来看一些集合操作的基本示例。有一个叫做 spawn 的函数，它接受另一个我即将调用的函数，并声明我将把这个函数复制运行四次，其中四就是世界大小（world size）。让我们看看这是做什么的。这实际上是我写的一个包装器，主要是为了绕开我没法在课堂上做多进程处理的问题。通常你会调用 Torch 的 multiprocessing.spawn，然后传入那个函数。但我要用这个分支来禁用分布式。那么让我们过一遍。现在我在这个函数里，它应该为每个进程异步运行。
[38:11.850] 记住，世界大小是进程的数量，而排名（rank）是从 0、1、2 一直到世界大小减一。所以有世界大小数量的这些函数，它们同时各自在一个进程上运行。我现在在排名 0 上。那我该做什么呢？有一个设置（setup）。你基本上要配置主地址（master address）和端口。注意，这实际上并不是 GPU 之间通信的方式。这更多是为了通用的元数据和协调。实际的数据是通过 NCCL 传输的。否则，速度会非常非常慢。如果你有 Cuda 可用，那么你可以使用 NCCL 后端。我在我的笔记本电脑上，所以我要用 gloo 后端。好的，现在我到这儿了。假设我有四个这样的进程在运行。这里有个 barrier 函数，它作为一个——它是一个同步屏障。
[39:20.510] 基本上，如果我看到这个，它会等待所有进程都到达这一点。基本上，你可以认为所有进程都在异步运行。我实际上无法控制。一个进程可能在另一个之前完全完成。它们可能以任何方式交错完成。所以如果我想确保某段代码在另一段代码之前执行，我就会放入这些同步屏障。那么加入更多屏障的缺点就是，你最终可能会不必要地等待。那么让我们试试 all reduce。我要创建一个张量，0, 1, 2, 3。我有我的排名。为了让例子更有趣，每个排名会有不同的张量。我要在 all reduce 之前打印出我有什么。现在我要跳到这里，看看会打印出什么。
[40:30.490] 所以排名 0 在 all reduce 之前有 0, 1, 2, 3。排名 1 有 1, 2, 3, 4，依此类推。这和我之前展示的例子一样。注意，打印语句出现的顺序取决于硬件，因为它是异步运行的。但所有的数据都在那里。那么现在如果我执行 all reduce，你传入——这是一个 PyTorch 函数。它传入这个张量。你传入归约操作（reduction operation），这里是求和（sum）。然后我要说不要异步。这个操作所做的就是，在这里调用的是 gloo，但也可能是 NCCL，如果是 NCCL，它会启动 Cuda 内核。
[41:18.020] 它会进行通信。它会为你处理所有事情。然后它基本上就地写入数据。所以在 all reduce 之后，我有了 all reduce 操作的结果，记住，这是每一列的总和，但在所有排名上复制了同样的结果。这就是 all reduce。如果你想要更花哨一点，做异步操作，那么你可以设置 async 等于 true。但那样会搞乱所有这些打印语句。所以我试着放入比平常更多的屏障。是的，有问题吗？[听不清] 所以问题是，排名是不是指 GPU。对于这门课来说，排名就是指 GPU。那让我们试试另一个例子。我要做一个 reduce scatter。这里，我要创建一个输入，是从 0 到世界大小。然后我要有输出。我要分配输出。那么在我做 reduce scatter 之前，这看起来是什么样的？
[42:49.880] 看起来，让我们看看，是这样的，输入是相同的，而输出恰好是零，但它可以是任何东西。然后我执行 reduce scatter tensor。这里不是就地写入，而是有一个输出张量和一个输入张量。我说我要做求和运算，然后之后，基本上输入不会被修改，但输出我得到的是每个分量归约后的结果，写入到各自对应的排名中。是的，有问题吗？[听不清] 所以问题是，异步执行 all reduce 是如何工作的？这意味着这是一个整体操作。你说开始做 all reduce。然后它就启动 Cuda 内核。它将要执行通信。请记住，CUDA 相对于进程来说已经是异步的了，而现在我们所有的进程也都是异步的。关键在于，这段代码会直接返回。然后你就可以去做其他事情。所以一个典型的做法，虽然我在这节课上不会展开讲，就是计算与通信的重叠。
[44:38.360] 例如，你可以先执行这个操作，然后继续为下一步加载一些其他数据，这些数据与此操作无关。然后当你想要确认操作确实完成时，就可以调用一个等待（weight）或屏障（barrier）。那么让我们来做最后一个操作，即全收集（all gather）。到了现在，我想你们已经明白这个概念了。这里我将把全规约-分散（reduce scatter）的输出作为输入。我会分配一个输出张量。所以在全收集之前，状态是这样的。我这里存放着从全规约-分散得到的结果。输出张量只是被分配了。它里面碰巧有一些值，但不用在意。然后在我执行将数据全收集到张量中的操作之后，所有不同的输入就会被收集到所有不同的秩（rank）上。你可以在这里看到，确实，通过实例证明了全规约（all reduce）等于全规约-分散加上全收集。
[46:13.270] 然后为了收尾，就像我开始时先设置一样，我最后会进行清理，清理是个好习惯。这就是你的第一个 Torch 分布式程序示例。那么让我们来做一些基准测试。这个会比较快，因为我想尽快进入第二部分。那么通信速度有多快呢？让我们做一个全规约。这里，我将对一亿个元素进行全规约。哦，抱歉，我搞错了。所以是全规约。我将创建一个具有这么多元素的张量。请记住，就像之前做基准测试时一样，我们先进行预热。在这里，我将调用 Cuda 同步（synchronize）和屏障（barrier），以确保——因为这里存在两种异步形式，即 CUDA 内核和不同的进程——确保所有操作在开始计时之前都没有在运行并且已经完成。
[47:49.430] 然后我将执行全规约。接着再次使用同步（synchronize）和屏障（barrier）等待，并停止计时。请记住，这在每个秩上都会运行。所以如果我查看输出，这里是秩 0、2、1、3，我可能会得到不同的时间，因为它们都是不同的进程。每个进程都会报告一个特定的测量值。如果你想报告一个数字，你可以取平均值，例如。现在，有一件很有用的事情可以做，类似于我们计算 MFU 时的情况，就是测量有效带宽。这里的思路是，嗯，这花了 1.6 毫秒。这算好还是坏呢？为了计算有效带宽，我们要做的是，基本上计算出在这次通信过程中应该发送了多少字节。然后如果你用这个字节数除以总时间，你就得到了有效带宽。
[49:10.460] 所以我发送的内容的大小是每个元素的大小乘以元素数量。这基本上就是这个数据张量的字节数。实际发送了多少字节呢？这需要一些解释。对于全规约，如果你想想，为了简单起见，假设你执行的是秩 0 加秩 1 加秩 2 加秩 3。你需要迭代世界大小减 1 步，因为有世界大小减 1 次加法运算。这就是这个因子的来源。有一个因子 2，因为你需要既发送又进行归约。然后你乘以有效载荷的大小。这就是发送的总字节数。总持续时间就是所花费的墙上时钟时间。然后你乘以世界大小，因为这是所有秩等待的总量。而带宽就是发送的字节数除以总持续时间。所以在这种情况下，你会得到大约每秒 400 GB 的结果。
[50:40.620] 这里有几个注意事项。第一，有效带宽，如果你看这个表达式，大小乘以 2 乘以（世界大小减 1）除以（世界大小乘以持续时间）。随着世界大小的增加，这个（世界大小减 1）除以世界大小基本上会收敛到 1。所以实际上你剩下的是 2 乘以大小字节除以持续时间。这基本上就是带宽。请注意，这与世界大小无关，这是好事。所以如果你增加 GPU 的数量，带宽不会改变。它也与拓扑结构无关，这基本上是由 NCCL 来判断你是要以环形还是树形拓扑来传递消息。所以，就是这样。这就是全规约。对于全规约-分散，这非常相似。我将创建输入和输出，进行预热，执行操作，并计时。请注意，全规约-分散有这样的时间记录。
[52:02.690] 你也可以在这里测量有效带宽，即输入中的字节数。你有了发送的字节数。这里没有 2 倍因子。然后用总持续时间去除，你就得到了带宽。所以在这里，带宽是——应该非常接近。我想有时候会有一些随机波动，但大致在 400 左右。这里有几个注意事项。记住，正如我们所说，全规约等于全规约-分散加上一个全收集。因此，全规约自然要移动两倍的数据量，因为全规约-分散有一些成本，全收集也有一些成本，而全规约做了两倍的工作。它需要两倍的时间。但这两者相互抵消，所以你得到相同的带宽。这就是第一部分的结尾。在我进入第二部分之前，有问题吗？是的。你为什么要做那个同步（针对听不清的部分）？那么我们为什么必须为 CUDA 内核做同步呢？
[53:44.050] 归根结底，我们仍然在执行 CUDA 操作。我们只是有多个进程，每个进程配有一个 GPU，在执行一些 CUDA 操作。所以如果你在执行 CUDA 操作，记住，默认情况下，它是异步的。当你在 Python 中执行到下一行代码时，那个 CUDA 操作可能还没有完成。所以我们需要总是通过同步来确保它完成。（听不清）那么你是否必须先做屏障然后再同步？我不确定。是的。（听不清）是的。我认为有个问题是，如果你先做屏障，而 CUDA 内核可能还没有运行完毕，你就直接进入了屏障。然后你仍然是在各自独立地同步不同的 CUDA 内核，这意味着你并没有真正实现同步。如果所有这些操作都只是返回了，那么屏障实际上起不到什么作用。好的，我们继续。现在让我们开始真正思考如何训练模型。
[55:12.960] 我们将从头开始看一个非常简洁的 MLP 训练实现，多层 MLP。我觉得这么说有点重复了。多层感知机本来就是多层。记住，在 Transformer 中，MLP 才是真正的计算瓶颈所在。所以这实际上非常能代表你将要看到的情况。即数据并行、张量并行和流水线并行。我希望你们脑海中有的画面是这个。这有点像示意图，所以不要太深入思考它，它更多的是一种概念化你如何切分数据和参数的方式。数据并行说的是，我将把数据分割成几部分。然后每个 GPU 负责一部分数据，并且我将只是进行正常的——我会保留所有参数并执行正常的模型训练。然后我需要同步。
[56:22.280] 那么让我来解释一下这是如何工作的。我打算生成一些示例数据。这里有一个大小为 128 的批次，维度数量是 1,024。所以这是一个批次大小乘以维度数量的矩阵，即数据矩阵。然后让我们进入这个数据并行的话题。它的工作方式是，你拥有这个数据矩阵。我要把行按照世界大小（world size）来切分，在这个例子里是 4 份。每个秩（rank）会得到一份。所以这里是维度数量、批次大小。我打算称之为本地批次大小（local batch size），基本上就是批次大小除以世界大小，也就是说每块 GPU 看到的每一行，会有 32 个数据点。
[57:29.240] 这只是索引操作。起始索引 data start 到结束索引就能得到数据的那个切片。然后我打算把它放到那个秩上。所以此时，每个 GPU 现在都有一个不同的数据张量，也就是它们各自负责的那部分。实际上，每个秩应该自己加载自己的数据，而不是通过这个瓶颈来加载，但这只是为了说明问题。那么让我们实例化这个 MLP。这里，我假设我们有若干层，以及层数（num layers）。对于每一层，我会有一个维度数量乘以维度数量的矩阵。所以我只是初始化一组随机参数。然后我把它送入优化器。这里是训练循环。在前向传播中，我取数据。记住，数据并不是全部数据，只是如果我是秩 2，那么我只得到数据中 B2 的那部分。我会遍历层数并做前向传播。然后我会做反向传播。
[59:00.470] 然后通常来说，这就结束了。但现在，记住，每个秩都有不同的数据。因此，梯度也会是不同的。所以这是让数据并行起作用的关键步骤。我们要在所有工作节点（worker）之间同步梯度。这是标准训练和 DDP 之间唯一的区别。这实际上是非常漂亮和优雅的。基本上，对于所有参数，我要对 param.grad 做一次全归约（all reduce），然后取平均。当这次全归约完成之后，这时，每个秩都有了完全相同的梯度。然后我就去更新参数。我觉得这真的很优雅，因为它基本上就是标准训练，把训练应用到你的本地批次上，但只是在反向传播之后插入了这一步。让我们通过这个全归约来平均所有梯度。
[01:00:16.370] 这是一行代码的修改，然后参数就被更新了。所以当你在训练的时候，每个秩基本上是在执行参数更新，就像它拥有全部数据一样，但它实际上只处理了一部分数据。所以这基本上就是 DDP，也就是第一种类型的数据并行。关于这个有什么问题吗？嗯。[听不清] 所以问题是，你只能在批次大小大于 1 的情况下这样做吗？是的。你的批次大小至少要等于世界大小，这样才有意义。而且通常来说，它应该比世界大小还要大不少。是吗？[听不清] 是的。所以问题是，批次大小应该是世界大小的倍数吗？那也会很好。是的。我的意思是，如果不是的话，你可以用零来填充之类的。所以是有办法的，但如果真的是倍数的话，对大家来说都更容易。[听不清]？
[01:01:39.870] 对于基于变换器（Transformer）的模型，这会是什么样子？[听不清] 是的，所以问题是，对于变换器来说，这会是什么样子？实际上基本上是一样的。DDP 有一个很好的特点，就是它非常模块化。你只需做前向传播。我的意思是，DDP 只是在这里对参数做平均。它并不关心你的前向传播长什么样。好的，让我继续。所以这就是 DDP。只是为了总结一下，各个秩之间的损失是不同的。梯度最初也是不同的。但它们通过全归约在所有秩之间变得相同。因此，参数在所有秩之间也始终保持相同。那么下一节课，Tatsu 会讲更高级的数据并行，FSDP 和 ZeRO。那里的思路是，正如我已经提到的，这里我们用的是全归约。这是一个非常简单的整体操作，但它确实要求在内存中保存所有的模型参数。
[01:02:51.380] 但如果模型参数放不进内存呢？那你就得更加聪明一点，这就是下一节课的主题。那么让我来谈谈张量并行（tensor parallelism）。这里，我们的思路是，我们要按这个方向切分，而且我们不切分数据。我们基本上要切分每一层。所以每个秩会得到每一层的一部分。一般来说，这意味着我们将不得不传输更多的数据。我们稍后会讨论这一点。那么张量并行是什么样的呢？我们只是假设我们有这些数据。为了简单起见，每个秩都拥有全部数据。在这里记住，数据是批次大小乘以维度数量。我要定义一个本地维度数量（local num dim）为——对于这个秩，我只负责维度的一个子集。这里的图景是，每个模型仍然拥有所有的层，但是对于所有的层来说，参数现在变成了维度数量乘以本地维度数量。
[01:04:24.200] 所以如果这是某一层的参数矩阵之一，我就是在按列切分。这也被称为列张量并行（column tensor parallel）。你也可以按行来做，但我们现在不讨论那个。那么前向传播是什么样子的呢？我会遍历所有的层，然后计算激活值。所以我从 x 也就是数据开始。我要访问 layer 处的参数。注意这仅仅是参数的一个切片。所以如果我在秩 1 上，我只得到矩阵的这一部分。但我仍然可以继续进行。我可以应用这个非线性函数，因为这本来就是逐元素的。但现在我要做的是，我要通信激活值。所以如果我有一个数据矩阵，秩 1 拥有这个矩阵这一部分对应的部分激活值。
[01:05:41.210] 秩 1 拥有这个矩阵对应的部分激活值，依此类推。我需要基本上把所有的激活值放到所有的秩上。但我们知道怎么做。我们介绍过全收集（all gather）这个集合通信原语。所以这是所有的激活值。这是批次大小乘以本地维度数量，也就是激活值的形状。然后我在做全收集。抱歉，这是为激活值分配内存。x 是激活值的实际部分。所以这是批次大小乘以本地维度数量。全收集的意思是每个秩都有 x，每个秩会分配一个 activations，它是一个列表，每个世界大小对应一个元素。然后在全收集之后，x 会被复制到 activations 中各自对应的位置。所以一旦我收集了所有的激活值，我就把它们拼接起来形成完整维度的 x，也就是批次大小乘以维度数量。关于列张量并行有什么问题吗？
[01:07:18.130] 每一层都要这样做。我们现在注意到与数据并行的一个区别，那就是我们现在必须去改动模型本身了。数据并行非常优雅，因为它按数据切分。模型被当作一个模块来对待。但现在我们不得不去改动模型。这在很大程度上利用了这样一个事实：如果你想做矩阵乘法，你可以把它拆分成一组更小的矩阵乘法。我们可以让不同的秩来做这些，然后我们可以收集结果。当涉及到反向传播的时候，[听不清] 梯度仍然是 [听不清]。是的。那么现在的问题是，在反向传播中会发生什么？现在，在反向传播中，你有了激活值，并且你需要通过 reduce scatter 将梯度分发到所有不同的设备上。所以在某些方面，all gather 和 reduce scatter 具有这种对偶性，在前向传播中，你是在做 all gather，而在反向传播中，你是在做 reduce scatter。
[01:08:33.490] 嗯？[听不清] 所以问题是，这是否由 autograd 自动完成？这些都不是——嗯，好的。所以如果你只是调用 .backward，它不会做这些，因为那里面没有并行性。但是 PyTorch 有所有这些自动为你完成的东西。[听不清] 我需要在多大程度上修改我的代码以及[听不清]？那么哪些是为你做的，哪些是自动的？在这里，我们相当明确地管理这些事情，这意味着我不是在做反向传播，但你需要自己管理并调用 reduce scatter。这是有意为之，因为这是 336 课程，从零开始构建语言模型。在实践中，你可能不必这样做。好的，那让我们快速讲一下流水线并行。流水线并行背后的想法是，我们要按这种方式来拆分网络。
[01:09:47.760] 所以现在每个 rank 将获得一部分层。在每一层内部，它将获得所有的维度，而且它还将获得所有的——嗯，其中一个 rank 将获得所有的数据。每个 rank 都会以某种形式看到所有的数据。所以这就是它的工作方式。我们有所有的数据，也就是批大小乘以特征维度。然后我要拆分这些层。所以局部层数将是特定 rank 要处理的层数。现在我将拥有局部参数，基本上只是——这里只有局部数量的层数那么多个。但在每一层内，我要做的是特征维度乘以特征维度。所以有一件事，也许[听不清]会在周三更多地讨论这个，就是微批次这个概念。所以也许我就先介绍一下这个，并解释我为什么要这样做。除了拆分层之外，我还要把批次拆分成许多微批次。
[01:11:14.400] 所以如果我是 rank 0，那么我获得数据，然后我会把它分块成若干个微批次。对于每个微批次，我要做的是从前一个 rank 接收它，然后只在分配给这个 rank 的层上执行前向传播。然后我会把它发送给下一个 rank。所以这里我实际上使用了这些 receive 和 send，它们是点对点操作。我之前没有讲过这些，但它们相当直观。这基本上是说，我是这个 rank，我要把这个张量发送到——抱歉，我要从 rank 减 1 接收这个张量。而这个是说，我要把张量 x 发送到 rank 加 1。所以基本上，我之所以要讲微批次是因为——而且[? Tatsu ?]周三会更多地讨论这个——在流水线并行中，你有一个 rank 获取数据。它处理一些层，然后把它发送到下一个 GPU，那个 GPU 处理一些层。
[01:12:53.930] 然后它再发送到下一个 GPU。所以这是划分深层网络的一种非常自然的方式。但问题是，你会遇到所谓的流水线气泡，也就是说当你在等待其他张量被处理的时候，你处于一种空闲等待的状态。而这最终会变得非常低效。所以微批次背后的想法是，你把它分解成更小的批次，这样你可以快速处理，然后发送到下一个。这样可以减少流水线气泡的数量。所以我要提到的另一件在这个非常朴素的版本中没有处理的事情是通信和计算重叠的概念，这对流水线并行实际上非常重要。这基本上给了你正确的结构。如果你在这些前面加一个 I，它就变成了异步的。你需要添加更多东西来管理代码。这里的思想是，你希望——当你在计算的时候，你可以同时接收数据或发送数据。
[01:14:11.270] 所以计算和通信应该重叠。这样就能减少你实际花在等待上的时间。所以这里还有一些缺失的东西，希望我们下次能补上。比如通信与计算的重叠，这在流水线并行中尤其关键。我没有提到在数据并行中，这也会发生，因为我只是做了一个前向传播，然后最后就做所有这些 all reduce。但如果你足够聪明，那么在反向传播中，一旦梯度计算完成，你就可以开始发送它了。这是作业二中将探讨的内容。而且，这再次让你能够更多地重叠通信和计算。我们有一些关于一般模型的问题。再说一次，我认为这个 MLP 基本上给了你理解基础所需的大部分内容。
[01:15:27.750] 一些更大的模型只是需要更多的簿记工作，所以更难看到核心算法。还有其他我们还没有涵盖的并行化类型。序列并行是取整个序列并把它切成碎片。这允许你并行化注意力计算。专家并行，它允许你为 MoE 并行化专家。这就是我之前提到的 all to all 发挥作用的地方。然后还有不同并行化技术的组合，这些也会出现在作业中。所以有一点是，你选择的并行化技术将强烈依赖于硬件。例如，张量并行。它有很多通信，因为对于每一层，你都需要发送所有这些激活值，这些激活值是相当大的。所以通常，张量并行发生在节点内部的 NVLink 上，或者是在你有高带宽的地方。而你不会在 NVLink 域之外做张量并行。
[01:16:46.820] 而流水线并行，你会看到人们在使用它。而且通常，它可以容忍慢得多的互连。所以一些去中心化训练的工作使用流水线并行，因为你的节点——GPU 在地球的另一端。但在那种情况下你不会想做张量并行。所以有时候当你看到这些组合时，会是在节点内部做张量并行，然后做流水线数据并行或 FSDP，如果需要的话再做流水线并行。还有其他影响，比如如果你做数据并行，你也许可以做得相当多。但然后你会开始碰到所谓的临界批大小，当你开始把批大小增加得太多时，它实际上不再有帮助了，在这种情况下你有点在浪费算力，然后你最好使用张量并行。
[01:17:56.160] 所以有这些各种各样的考量，我们会在课程进行中更多地讨论。最后一点说明。TPU 之前被提到了一点。有一点需要注意的是，我们故意使用 PyTorch，而且不仅仅是使用 PyTorch，而是以非常原始的方式使用集合操作，这样你就能从机制上看到正在发生的事情。另一种方法，特别是如果你在[听不清] TPU 领域的话，是你可以简单地定义模型和起始策略。然后编译器实际上处理了很多关于如何——基本上，你需要什么样的通信操作的决策。你基本上就是说，嗯，这份数据需要放在这里、这里还有这里。然后编译器施展一些魔法来弄清楚具体怎么放。所以这很吸引人。但显然，它会夺走很多从零开始构建东西的乐趣。简单总结一下，并行化的方式有很多种。
[01:19:14.660] 你可以按数据切分、按张量或专家切分、按流水线或序列切分。我们看了数据并行。我们只做了DDP。下次我们会讲FSDP和ZeRO。张量并行，正如我提到的，需要非常快的互联网络。流水线并行对互联的要求没那么高。但你需要非常努力地去减少这些流水线气泡。然后可能在高层次上，我们会经常看到这种模式。所以你要么重新计算，要么存储在内存中，比如我们在讨论激活检查点或者使用GPU时。或者在这种情况下，你可以把这个思路扩展为，你可以存储在不同的GPU上。从这个角度来看，你看看数据并行。从某种意义上说，你在做冗余的工作，因为每个rank实际上都在更新自己的参数并维护所有参数。但你之所以这么做，是因为你不需要在节点之间传输优化器状态。
[01:20:33.090] 所以有一件事是硬件在变得越来越快。但在某种意义上，我们总是想要更大的模型。所以这种层级结构的思想会一直存在。今天就讲到这里。下周三，Tatsu会更深入地讲解更多的并行化技术。
CS336 2026 · Lecture 7 · Parallelism
- 视频：官方视频
- 字幕：en-US · manual captions
- 处理：连续字幕合并为自然段；每段保留首条字幕的起始时间戳
- 状态：原语言字幕稿，待校对
正文
[00:05.160] OK. Let's get started. So welcome back, everyone. So today we're going to talk about parallelism. And remember in the last week, we introduced how to make a single GPU go fast by writing kernels. And we really looked inside this GPU. And this week, we're going to talk about how to leverage multiple GPUs to make your code go even faster. So the picture you should have in your head is something like this. So for the last week, we focused on one of these boxes, where you have your GPU. Remember, you have your High Bandwidth Memory, HBM, L2 cache, L1 cache, registers, and a bunch of streaming multiprocessors.
[00:51.680] So now the picture gets extended, because instead of having one GPU, you might have four, you might have 1,000 GPUs. And those are going to be connected. And I'll talk later about how those GPUs are going to get connected. And then you're going to have to figure out how to leverage all of this compute to train models. So in both cases, meaning both in the single GPU case and in the multi-GPU case, as we'll talk about today, the situation is kind of similar if you zoom out, which is that the compute, the arithmetic logic units, the tensor core, and so on, is far away from your data. And far away for a single GPU means all the way over here in HBM. And now if you have multi GPUs, the thing that you need here might be all the way on a different GPU, and you're going to have to shuffle that over somehow. But the same principles are going to be the same, because the game is to orchestrate the computation to try to avoid data transfer bottlenecks.
[01:55.280] It's very easy to use a ton of GPUs, but it's hard to use them effectively. So taking a little bit of liberty, we can think about the generalized hierarchy where at the local level next to the SMs a single node single GPU where you have an L1 cache shared memory. This was the fastest. And then you had HBM, which we lamented was so slow. But in this lecture, HBM is going to be considered fast. Now we're going to think about the single node multi-GPU setting where the GPUs are going to be connected via NVLink and NVSwitch. And then finally, the multi-node and multi-GPU, where we have to resort to InfiniBand and ethernet depending on what network you have. So last week, we talked about various tricks for in reducing memory accesses, fusion and tiling. Read into shared memory, do as much as you can, and then write it back out.
[03:05.990] And this week, we're going to talk about how you can reduce the amount of communication across GPUs by replicating and sharding appropriately. So why do you do multi-GPUs? The obvious answer is, well, you want to scale. But to put a finer point on it, there's really two reasons. One is that your parameters or activations or gradients and optimizer state don't fit on the HBM memory of a single GPU. So B21 has 192 gigabytes. If you're training a 1 trillion parameter model, that's not going to fit on a single GPU. And the other reason is that even if your model could fit on a single GPU, you might want to leverage more GPUs by splitting everything up
[03:54.630] to train faster. So sometimes there will be some decisions to be made because you could fit everything on GPUs, but you have fewer cores. But if you spread it out, then you're going to have to pay the communication bandwidth. So that's some calculation you're going to have to do to figure out how to parallelize. So just one note here is that so far this lecture is Python. You execute it and you can just show everything. Now, in this lecture, if you run this directly, it uses multiprocessing. But when I trace through it, I'm putting in some of special single process mode. So if you want to see the standard out for this lecture run in the multiprocessing setup, you can click here. And I'll show this as we go through the lecture. But just remember as we step through this lecture, we're not actually doing multiprocessing, because we're just stepping through single lines of code.
[04:56.990] So this lecture is going to include two parts. One is we're going to learn about the building blocks of distributed communication and computation, starting with the programming model. Talk a little bit about the hardware. Start to implement things in Torch, which you're going to do on your assignment two. And then the second part is we're going to look at actual training. We're going to look at three types of parallelism. Data parallelism, tensor parallelism, pipeline parallelism. Each is going to cut up our model in different ways. We're going to do this for MLPs rather than the full transformer. But it's really the core computation is going to be shown here. So let's dive in. So the first thing to talk about is these things called collective operations. So collective operations are these primitives from distributed programming that go back to the '80s.
[06:03.200] So the idea of parallel programming is very old. It wasn't invented for LLM training. And it's still the case that these primitives are the ones that we use today. And here, collective just means that you're specifying a general communication pattern or a template across multiple devices, rather than managing point to point how this GPU is going to communicate with another GPU. And this is going to be much easier, and the system can do a lot more work for you. So this is a very tried and trued interface for doing parallel programming. So the general setup is as follows. The terminology here is a little bit-- I find it a little bit strange, but this is standard in parallel programming. The idea is that you have a bunch of ranks where rank corresponds to a particular device, in our case a GPU. It could be a TPU. But the point is that you have, let's say, four ranks here. The world size corresponds to the number of devices, so the world size here is four.
[07:13.060] So there is a few operations we're going to go through. Broadcast, scatter, gather, reduce, all gather, reduce scatter, all reduce, and all to all. And each of these operations is going to specify how this set of ranks or devices is going to transfer some amount of data slash compute with it to some other set of devices. So the first three, broadcast, scatter, gather, reduce, are really just warm ups, I would say. They allow you to get a sense of how these collective operations work, but they're not really going to be the ones that are driving most of training. All gather, reduce scatter, and all reduce are the ones that are going to show up again and again for distributed training of language models.
[08:05.080] And finally, all to all I'll just mention here, which is important for MOEs, but we're not going to actually spend too much time on this this lecture. OK, so let's dive in with the simplest operation, which is broadcast. So in broadcasting, you have a rank 0. It could be any rank, but let's just say for sake of picking one, rank 0, has some tensor 0123. And it broadcasts it to all the ranks. So at the end of this operation, we have that each of the ranks has the same tensor on it. That should be pretty straightforward. And this, again, doesn't really show up in the core path of training. Generally, broadcasts are used for initialization, where let's say you initialize a load initial checkpoint and then you broadcast it to all the ranks. So something that's done like once.
[09:10.260] So the second operation is a scatter. And a scatter basically says, I have a tensor at rank 0 who is split up into the world size. And I'm going to basically scatter my tensor onto the other ranks. So rank 0 gets 0, the zeroth component. Rank 1 gets this, rank 2 gets this, and rank 3 gets that. So again, this is not directly used, but scatter is an important stepping stone to understand reduced scatter. So, as the name implies, scatter just takes a big tensor at one place and spreads it out onto multiple places. And you can see how this might be helpful, because you want all the GPUs you're scattering to do some local computation on the different parts. So the inverse of scatter is scatter. So this should be very predictable. The input is you have a bunch of pieces, each of which reside on a particular rank. And then when you do a gather, that's with respect to a particular rank, rank 0, it's going to just concatenate all the pieces together.
[10:38.170] Again, gather isn't directly used, but it's going to be a stepping stone to understand all gather. So next is reduce. So those of you who do functional programming are probably familiar with what reduce is. It's exactly the same. The idea is that you start in the same starting point as a gather. And where the tensor is split up across multiple-- well, you have some piece of data on each of the different ranks. And then you're going to apply your reduction operation to all of these and get that in, put that on rank 0. So in this case, if you do a reduction with sum, then you add these all up, you get six. So you can think about gather as a reduction where the operation is concatenation, if you will. And of course, reduce is important to understanding all reduce. So let me pause there. Those were just the warm ups. Just in case people have any questions about what a collective operation is, broadcast, scatter, gather, reduce.
[11:59.850] Yeah? [INAUDIBLE] NumPy, where is it? So the question is, is this related to broadcasting and NumPy? I mean, I think it's conceptually the same idea, where you have one thing that goes to many things. Like in NumPy, if you have a scalar, get broadcast to a tensor. But the instantiation, this is for collective communication. So it's a bit different. So let's move on to something more interesting. So all gather is what you do to-- basically, you perform gather to all ranks, not just rank 0. So remember what gather does. It basically takes all the different pieces and it just puts it on one rank, rank 0. And now all gather just does it for every single rank. That's what the all is for. All means do it for all the output to all the ranks, and gather is what you're doing to all the ranks.
[13:12.560] So this is going to come up a bunch. It's not important that you understand this statement precisely, but later we'll see that each rank holds part of the parameters. And then what you need to do is all gather the parameters to get the full parameters for the full forward pass. So in general, as we're doing training, we're going to see a lot of this gather to do something and then scatter and then gather and scatter again. So reduce scatter is performing reduce on each dimension and then scattering the results. So let's say you have four devices and each of them has some vector. And so when we did a reduce before, we just had 0123, and that got reduced to 6. But now reduce scatter says that for each component of this tensor, I'm going to do a reduction, and then I'm going to put it on a different rank. So the first dimension, I'm going to add these up, I get 6.
[14:33.625] Now, for the second dimension, I'm going to add these up. I get 10. For third dimension, I'm going to process these. For the fourth dimension, I'm going to process those. So where this is going to show up, just to foreshadow things, after the backward pass, when you sum what you're going to do is each GPU will be dealing with different data. And what you're doing is you need to sum all of the gradients from the different shards. And then you're going to redistribute this storage. And then finally, all reduce. If you understand all reduce gather and all gather, it's basically you do one and then you do the other. So what this does is reduce scatter, the same input as we had before. And remember, in reduce scatter, we had 6, 10, 14, and 18 sit on different ranks. And the all gather part of all reduce just puts them all on the same-- everything on the same node. So all reduce is in some sense the easiest to understand.
[15:59.620] You have a bunch of tensors. You reduce, in this case, some. And then you replicate them on all the nodes. So we're going to see this one actually first when we do a data parallel where we sum the gradients and then we replicate the full parameters. So that's where we're actually going to start. So maybe just focus your attention on all reduce. Later, we're going to see how to get to fancier things like 0 or FSDP. We need to break the all reduce into reduce scatter and all gather, because then you can intervene and you can manage things a bit more. But for the basic version, all reduce is fine. Finally, all to all. This one is, in some ways, the most general. You basically specify how each rank sends a particular message to another rank. And so here's a simple example where you have the same input as before.
[17:09.310] And what this is saying is that I want to send 0 to this element to rank 0, meaning keep it myself. I'm going to send one to rank 1, two to rank 2, and three to rank 3. And then if I'm rank 1, I want to send four to rank 0, five to rank 1, six to rank 2, and seven to rank 3. So basically, the position here is going to denote which rank is going to be the ultimate destination. And so if you look at the output, what happens is that the first rank is going to receive everyone who sent everything in column 0, because all these ranks are sending these things to rank 0.
[18:03.510] Similarly, for rank 1, all of the ranks are sending this column to rank 1, and so on and so forth. So this is going to be useful when for training MOEs. And the intuition here is that each rank has both a split of the data and also a subset of experts. And basically, and the key idea of the MOE is that it's dynamic routing. You have to look at your data to figure out which experts you need to route those activations to. So it ends up being a all to all communication. So if everything were balanced, meaning that every rank sent the same number of bytes to every other rank, then all to all, you can think about it as essentially a transpose. If you think about this as a matrix, all you're doing is transposing that matrix. But in general, all to all also handles unbalanced splits. And I'm not showing this here, but you can configure it to send any number of bytes to any other rank. But in general, you want the splits to be as balanced as possible.
[19:27.290] So remember [INAUDIBLE] MOE lecture where we had load balancing to make sure that things were as balanced as possible. So morally, the ideal goal is to have all to all look kind of like this. So just to summarize, maybe a few helpful tips to remember the terminology, because I just went through quite a few different operations. So reduce is, well, it's reduced. It performs on some of associative commutative operation. Could be some, could be max, could be min. Scatter is the inverse gather. Scatter distributes, gather centralizes. And all just means that the destination is all devices. So that explains all reduce and all gather. Let me pause there to take any questions about collective communications. Yeah? For operations such as gathering where we basically [INAUDIBLE] to rank 0, does that rank 0 [INAUDIBLE]? Is it a particular GPU every time or can that rank 0 change?
[20:47.680] Yeah. So the question is when you do a gather or a reduce, the target where you write the output, right now I said rank 0. You'll see later in the code, you basically specify the GPU ID or the rank and it goes there. So it doesn't have to be determined way in advance, but it has to be determined basically when you execute the call. Cool. Anything else? Yeah. These are just conceptual building blocks, right? They're not actually [INAUDIBLE]. Are these the actual [INAUDIBLE]? Yeah, so the question is, are these just conceptual building blocks or are they code? So I'm showing these right now as just conceptual building blocks, but we'll very quickly see how these are implemented in code.
[21:53.460] So before getting into the code, I want to talk a bit about the hardware, in particular how GPUs are connected, because we already what's inside a GPU. So let's talk about networking in general. So this is a very classic picture. You can tell from this very old looking image that this is how computers, I mean, generally work. So you have this server and then you have a bunch of CPUs. There's a PCIe bus, which you connect things like your-- or used to connect things like your mouse and keyboard. And then you have a bunch of GPUs sitting off of them and you have some RAM. And then this computer is connected to ethernet to another computer, and so on and so forth. So this is a particular setup. It's a particular topology. And GPUs on the same node, you use PCIe to communicate and GPUs on different nodes, you have to go all the way through ethernet. So this is like if you bought your gaming GPU and you had hooked it up with your friend and was like, I'm going to train some big model.
[23:12.500] That's what you would have to do. But if you're really serious about training, then things look more like this. This is the picture I showed in the very beginning where you have the GPU and there's something called NVLink and NVSwitch and InfiniBand. So the typical setup is this. And these numbers, eight is typical, but this 256 is kind of made up. So typically, you have eight GPUs per node. And these are connected via NVIDIA's NVLink to a switch. And just for calibration, if you use NVLink 5, then you're getting 1.8 terabytes per second of total bandwidth. And remember, HBM for B200 was 8 terabytes per second. So it's about 4x slower. So I mean, this is still pretty fast if you think about going between devices, but obviously, not as fast as high bandwidth memory, which is much slower than shared memory or L1 cache.
[24:28.450] So basically, NVLink connects to the switch, which means that from a programming perspective, you can think about GPUs as connected to any other GPU. You go GPU to any other GPU, and the hardware takes care of transmitting that to the switch, and the switch routes it. So typically, what you will also have is that at some point, you can't have NVSwitch and NVLink. Because as your number of GPUs grows, then you're going to have to put these nodes into pods, which are connected by InfiniBand. And the way InfiniBand works is that now there's a bit more the-- so now the GPU doesn't connect directly to another GPU.
[25:25.390] It has to go through PCIe and goes through this kind of special InfiniBand cable. And you see that the speeds are much, much lower. And then finally, if you run out of InfiniBand and you have these huge pods, then you need to connect them via ethernet. And in ethernet, you have to go through PCIe and that actually goes through the CPU, which as we'll see, is even slower. So it's kind of analogous to the kind of memory situation. The more nodes you have, then the slower it's going to be. You can't have an NVSwitch handling 100,000 GPUs. So one note is that I alluded to this bypassing the CPU, which is going to be an important thing from a hardware perspective.
[26:31.730] So if you have traditional ethernet, what happens is that the GPU has to talk to the CPU to get its data copied. So it basically has to copy the data to this. The CPU has a kernel socket buffer. Here kernel means not the GPU kernel, but the CPU, traditional notion of a kernel. And then it has to build some network packets, copy to the network interface, and then ship it over. So this generally introduces a lot of latency. And so there's this technology called Remote Direct Memory Access, RDMA, which allows a GPU to directly write or read from another GPU's memory without using the CPU at all. So obviously, if you're in NVLink land and NVSwitch land, then you have RDMA. InfiniBand also supports RDMA. So if you're connected via InfiniBand, then you can directly have GPUs connect to each other without involving the CPU. But standard ethernet does not.
[27:42.360] There's two notable advancements I will mention. So NVIDIA has been really pushing the limits on what you can do with larger and larger pods. So they have four of basically the B200's and the B300's. They have something called NVL72, which means that they have these trays of eight GPUs but have nine of them. And so basically, at the end of the day, you have 72 GPUs that are all kind of NV switched into one NVLink domain. And if you remember, the NV link speeds are very fast. So normally, if you're mortal, you think, well OK, I have eight GPUs that are interlinked really fast. And then outside of that, then things get slowed down a lot. But if you have a lot of money, you can buy this really fancy hardware and you can get really fast interconnects up to 72 GPUs. And the other thing I'll mention is that I said standard ethernet doesn't support RDMA, but there has been progress on the ethernet front as well.
[28:55.100] So there's something called RoCE, RDMA over converged ethernet, where the ethernet actually bypasses the CPU. And this is sort of their answer to InfiniBand. So InfiniBand generally is very expensive, as is a lot of NVIDIA products. But you can get pretty good performance by using RDMA over converged ethernet. And Meta had some paper showing that they were exploring this. So [INAUDIBLE] may or may not have been trained over converged ethernet. So that's just a brief overview of what the hardware looks like. You have GPUs. They're connected via NVLink to NVSwitch over some domain, maybe 8, maybe 72.
[29:55.440] And then it's InfiniBand from there. And now let's talk about how do you program this. So at the very lowest level, there's something called the NVIDIA Collective Communications Library, NCCL or "nickel," which translates the collective operations, the all reduce, reduce broadcasts, into the actual low level packets that are sent between GPUs. So what a NCCL does is when you use NCCL, it's basically like saying, I want to all reduce. And then NCCL goes and figures out what is the topology of the hardware, figures out the path between different GPUs, and then it actually launches the GPU kernels to send and receive data. Because at the end of the day, remember, everything that runs on a GPU is a kernel. So there are communication kernels as well that actually do communication with other GPUs. So we're not going to look too much more into NCCL, but just know that it exists.
[31:04.710] And then we're going to actually go to PyTorch. So maybe before that, any questions about hardware? Yeah. Can you describe-- there's a rack and a tray. Can I describe physically a rack and a tray, like what a rack is and what a tray is? So for the NV72, so I'm not a hardware expert, but a rack, I mean, is literally, I mean, you've seen data centers, literally a rack. And each tray is something that has-- so the G stands for a [? great. ?] So there's two CPUs. And each CPU is connected to four GPUs. So each tray has eight GPUs on it. And they're stacked and everything is connected to this NVSwitch. Yeah. [INAUDIBLE] Yeah, so the question is, what is the difference between [INAUDIBLE].
[32:19.000] [INAUDIBLE] Yeah. So RDMA, you can think about it as more of a desiderata. RDMA means that one GPU can read and write from another GPU's memory. And there's multiple ways to do RDMA. One is to use NVLink and NVSwitch and another way is to use InfiniBand. So InfiniBand and NVSwitch and NVLink are more the hardware, what pieces, what cables and switches are there. And RDMA is more operation, what happens when you're communicating. [INAUDIBLE] Yeah. For example, this [? advancement ?] RDMA over converged ethernet is another way to do RDMA. [INAUDIBLE] Yeah? Is NCCL optimized for multi-node clusters? Because obviously, it seems intuitive [INAUDIBLE].
[33:43.350] But for RDMA based processes, would NCCL be optimal for that? So the question is, is NCCL optimized for multi-node clusters? So I don't know the details of how they have or have not optimized it. All I can say is that NVIDIA has been basically optimizing their entire stack for inference and training of these large models, because their main customers are the major providers of language models. So I would be surprised if they haven't thought of optimizing for those type of workloads. Yeah? [INAUDIBLE] So the question is, what happens if you have nine GPUs? How do you distribute the workload across those? I guess it kind of depends on how the nine lands in your setup. For example, a lot of times you'll have-- let's say you have eight GPUs per node. So then the ninth one would be on a different node. And if you don't have NVLink connecting them, then that's going to be really bad, because that's going to be one node which is not providing that much compute and also very expensive to communicate with.
[35:23.860] But if you, let's say, had everything were connected by NVSwitch, then it would be much more reasonable. One more question and I'm going to move on. So how is this different from TPUs? So [INAUDIBLE] described TPUs a bit more. So I mean, let's see. What can I say quickly about this? So TPUs are generally much simpler objects. I'm not too familiar with the details of what each of these components corresponds to, but maybe we can talk about it offline. So let's actually get down to some code to take advantage of this hardware. So PyTorch conveniently has a torch distributed library that provides a clean interface into these collective operations. So you don't have to explicitly think about NCCL. And in fact, this library also supports different backends for different hardware. So if you are on GPUs, then you would use the NCCL backend. And if you were on CPUs, then there's something called gloo which still allows you to-- like I said, parallel processing has been around for a long time before GPUs.
[36:54.900] And you can do these collect operations on CPUs as well. This library also supports higher level models and algorithms such as SFTP, but we're not going to use those in the course because we're building things from scratch. All right. So let's walk through some basic examples of collective operations. So there's this function called spawn, which takes another function I'm going to call, and it says I'm going to run this replicated four times where four is the world size. So let's see what this does. This is actually a wrapper I wrote just to hack around the fact that I can't do multiprocessing in this lecture. So normally what you would do is you call Torch's multiprocessing.spawn and you call the function. But I'm going to do this branch which disables distributed. So let's just go through that. So now I'm in this function that it's supposed to be running asynchronously for each process.
[38:11.850] So remember, world size is the number of processes and the rank is either 0 or 1 or 2 all the way up to world size minus 1. And so there are world size number of these functions that are each running on a process at the same time. So I'm on rank 0 right now. So what do I do? There's a setup. You basically configure the master address and port. Notice that this is not actually how the GPUs are going to communicate. This is more for general metadata and coordination. The actual data goes through NCCL. Otherwise, it would be very, very slow. And if you have Cuda available, then you can use the NCCL backend. I'm on my laptop, so I'm going to use the gloo backend. OK, so now I'm here. Pretend I have four of these processes running. So there's this barrier function, which is useful as-- it's a synchronization barrier.
[39:20.510] So basically, if I see this, then it waits for all the processes to get to this point. So basically, you can think about all the processes are running asynchronously. So I don't really control. One could completely finish before the other. They might finish interleaved in any way. So if I want to make sure that there's some code that is executed before other code, I put these synchronization barriers in. So now the downside of putting more barriers in is that, well, you end up kind of waiting potentially unnecessarily. So let's try an all reduce. So I'm going to create this tensor, 0, 1, 2, 3. And I have my rank. Just to make it more interesting, each rank is going to have a different tensor. I'm going to print out what I have before the all reduce. So now I'm going to skip over here and see what gets printed out.
[40:30.490] So rank 0 before all reduce has 0, 1, 2, 3. Rank 1 has 1, 2, 3, 4, and so on. It's the same example as I showed before. Notice that the print statements are coming in whatever order the hardware feels like it, because it's running async. But all the data is there. So now if I do all reduce, and this you pass in-- this is a PyTorch function. It passes in this tensor. You pass in the reduction operation, which is a sum. And I'm going to say don't do async. And what this does is it calls, in this case, it would be gloo, but it could be NCCL, in which case, it would spin up the Cuda kernels.
[41:18.020] It would do the communication. It takes care of everything for you. And then it basically writes in place to the data. So after the all reduce, I have the all reduce operation, remember, which is the sum of each of these columns but replicated across all the ranks. So that's all reduce. And if you wanted to be fancier and do async, then you could say async equals true. But then it would screw up all these print statements. So I'm trying to put more barriers than I would normally would. Yeah, question? [INAUDIBLE] So the question is the rank the GPU. For this class, the rank is the GPU. So let's try another example here. So I'm going to do a reduce scatter. So here, I'm going to create an input, which is 0 through the world size. And I'm going to have the output. I'm going to allocate the output. And so what does this look like before I do the reduce scatter?
[42:49.880] It looks like, let's see, this, where it's the same input and the output happens to be zeros, but it could be just anything. And then I do the reduce scatter tensor. Here instead of writing in place, I have an output tensor and an input tensor. I say I want to do the sum and then afterwards, I get basically the input is not touched, but the output I get the reduction of each component written into the respective ranks. So yeah, question? [INAUDIBLE] So the question is, how does it work to do the all reduces asynchronous? So what this means is that this is like a monolithic operation. You say go do the all reduce. And it's spinning up Cuda kernels. It's going to do the communication. And remember, Cuda is already kind of async with respect to the processes, and now we have all the processes being async. And the point is that this code just would return. And then you could do other things. So a typical thing, which I'm not going to talk about this class, is overlapping computation and communication.
[44:38.360] So for example, you can do this operation, and then you can go ahead and load some other data for the next step, which is independent of this operation. And then when you want to make sure that you actually are done, then you can call a weight or a barrier. So let's do the final one, which is all gather. So by now, I think you get the idea. Here I'm going to set as the input the output of the reduce scatter. I'm going to allocate an output. So before the all gather, it looks like this. I have my results from the reduce scatter here. The output is just allocated. It happens to have some values in it, but don't worry about it. And then after I do the all gather into tensor, and then I will have all the different inputs gathered onto all the different ranks. And you can see here that, indeed, proof via example that all reduce is equal to reduce scatter plus all gather.
[46:13.270] And then just to wrap things up, just as I started by setup, then I clean up, which it's good practice to clean up. So that was your first example of a Torch distributed program. So let's do some benchmarking. This will be quick, since I want to actually move on to part two. So how fast does communication happen? So let's do an all reduce. So here, I'm going to all reduce with 100 million elements. Oops. Sorry, I messed this up. So all reduce. So I'm going to create this tensor with this number of elements. And remember, just like before when we do benchmarking, we warm up first. And here, I'm going to call the Cuda synchronize and also the barrier, just to make sure that because there's two forms of asynchrony here, the Cuda kernels and the different processes. And just want to make sure everything is kind of not running and is done before I start the time.
[47:49.430] And I'm going to do the all reduce. And then wait again with the synchronize in the barrier and stop the time. So remember, this is running for every single rank. So if I look at the output, here for rank 0, 2, 1, 3, I have a different time potentially, because they're all different processes. Each of them is going to report a certain measurement. And if you want to report one number, you can take the average, for example. So now one thing that is useful to do, which is analogous to when we were computing MFU, is to measure the effective bandwidth. And the idea here is that, well, this took 1.6 milliseconds. How much-- is that good or bad? So to compute the effective bandwidth, what we're going to do is to compute essentially how many bytes were sent and should be sent during this computation. And then if you divide by the total time, then you get the effective bandwidth.
[49:10.460] So the size of what I'm sending around is the size of each element times the number of elements. So that's basically the number of bytes of this data tensor. How many bytes get actually sent? So this needs some unpacking. So for all reduce, if you think about, let's just say for simplicity, you do a rank 0 plus rank 1 plus rank 2 plus rank 3. You need to iterate this world size minus 1 steps, because there's a world size minus 1 addition operations. So that's this factor. There's a 2, because you need to both send and reduce. And then you multiply by the size of the payload. So that's the total number of sent bytes. And the total duration is the wall clock time that it took. And then you multiply by the world size, because it's the total amount that all the ranks have waited. And the bandwidth is the bytes sent divided by the total duration. So in this case, you get something like about 400 gigabytes per second.
[50:40.620] So a few notes here. One is that the effective bandwidth, if you look at this expression, size times 2 times world size minus 1 divided by world size times duration. So as world size increases, this world size minus 1 divided by world size essentially converges to 1. So you're effectively left with 2 times the size bytes over the duration. So this is essentially the bandwidth. Notice that this is independent of the world size, which is good. So if you grow the number of GPUs you have, the bandwidth doesn't change. It is also independent of the topology, which is something that kind of NCCL figures out whether you're going to pass the messages in a kind of a ring or a tree topology. So yeah. So that is all reduce. And for reduce scatter, this is very similar. So I'm going to create the inputs and the outputs, warm up, perform the operation, and time it. So notice that the reduce scatter has these timings.
[52:02.690] And you can also measure the effective bandwidth here, the number of bytes that were in the input. You have the number of bytes that were sent. Here there is no 2x here. And then the total duration, you divide by that, you get the bandwidth. So here, the bandwidth is-- it should be very similar. I guess sometimes there's some stochasticity, but it's in the kind of 400s. So a few notes here. All reduce is, remember, as we stated, reduce scatter plus an all gather. And so all reduce, naturally, is moving twice the amount of data, because reduce scatter has some cost, all gather has some cost, and all reduce is doing twice as much work. And it takes twice the amount of time. But the two cancel out, so you get the same kind of bandwidth. So that is the end of part one. Maybe any questions before I move to part two? Yeah. Why did you do that synchronized for [INAUDIBLE]? So why do we have to do the synchronized for the Cuda kernel?
[53:44.050] So at the end of the day, we are still doing Cuda operations. We have just multiple processes, each with a GPU, doing some Cuda operations. So if you're doing a Cuda operation, remember, by default, it's async. So when you reach the next line in Python, that Cuda operation might not be done. So we need to wait always to make sure it's done by synchronizing. [INAUDIBLE] So do you have to do barrier first and then synchronize? I'm not sure. Yeah. [INAUDIBLE] Yeah. I think one problem is if you barrier first and then the Cuda might not be done running, and you just immediately go to the barrier. And then you are still each independently synchronizing the different Cuda kernels, which means that you're not really synchronized. If all those operations just return, the barrier doesn't really do anything. OK, let me move on. So now let's actually start thinking about how you train models.
[55:12.960] So we're just going to walk through a very bare bones implementation of training MLPs, multilayer MLPs. I guess that's redundant. It's multilayer perceptrons already. And remember that MLPs are the ones that are the actual compute bottleneck in the transformer. So this is actually pretty representative of what you'll see. So data parallelism, tensor parallelism, and pipeline parallelism. And the picture I want you to have in your head is this picture. So this is a little bit of a schematic, so don't think too deeply about this, but it's more of a way to conceptualize how you're cutting your data and parameters. So data parallelism says, I'm going to split the data into pieces. And then each of the GPUs is going to be responsible for part of the data, and I'm going to just do normal-- I'm going to keep track of all the parameters and do normal model training. And then I need to synchronize.
[56:22.280] So let me explain how this works. So I'm going to generate some sample data. So there's a batch size of 128. Number of dimensions, 1,024. So this is a batch size by num dim matrix, data matrix. And let's jump into this data parallelism. So the way that it's going to work is that you have this data matrix. And I'm going to break up the rows into world size pieces. And in this case, four. And each rank is going to get a piece. So the number of dimensions, the batch size. So I'm going to call the local batch size, basically the batch size divided by the world size, which is every row that GPU sees, it's going to have 32 data points.
[57:29.240] And this is just indexing. Start index data start to end gets you the slice of that data. And I'm going to just put it on that rank. So at this point, each GPU has now a distinct data tensor, which is the part that they're responsible for. Now, in practice, each rank should probably load its own data rather than have this bottleneck, but this is just for illustrative purposes. So let's instantiate the MLP. So here, I assume we have layers. And num layers. And for each of the layers, I'm going to have just a num dim by num dim matrix. So I'm just going to initialize a random set of parameters. And then I'm going to feed that into optimizer. So here's the training loop. So in the forward pass, I take the data. So remember, data is not all the data. It's just if I'm rank 2, then I only get B2 part of the data. I'm going to just go through the number of layers and do a forward pass. And then I'm going to do a backward pass.
[59:00.470] And then normally, this would be it. But now, remember, every rank has different data. Therefore, the gradients are going to be different as well. So this is the key step that makes data parallelism work. We're going to synchronize the gradients across all the workers. This is the only difference between standard training and DDP. It's actually pretty nice and elegant. So basically, for all the parameters, I'm going to do an all reduce of parem.grad. And I'm going to average. And then after this all reduce is done, then now at this point, each of the ranks has the exact same gradients. And then I'm just going to update the parameters. So it's kind of really elegant, I find, because it's basically standard training where you apply it to your local batch, but you just insert this after the backward pass. Let's just average all the gradients via this all reduce.
[01:00:16.370] It's a one line code change, and then the parameters get updated. So as you're training, each rank is basically performing parameter updates as if it had all the data on it, but it's only actually processing a part of the data. So that's basically DDP or the first type of data parallel. Any questions about this? Yeah. [INAUDIBLE] So the question is, can you only do this with batch size greater than 1? Yes. So your batch size has to be at least world size for this to really make sense. And usually, it should probably be quite a bit larger. Yeah? [INAUDIBLE] Yeah. So the question is, should the batch size be a multiple of world size? And that also would be nice. Yes. I mean, if it's not, then you can pad it with zeros or something. So there's ways, but it's just easier for everyone if it is. [INAUDIBLE]?
[01:01:39.870] What would it look like for a transform-based [INAUDIBLE]? Yeah, so the question is, what would this look like for a transformer? It would actually be basically the same. The DDP has the nice thing that is very modular. You do the forward pass. I mean, DDP just averages the parameters here. It doesn't care what your forward pass looks like. OK, let me move on. So that's DDP. So just to summarize, the losses are different across the ranks. The gradients are also initially different. But they are all reduced to be the same across the ranks. And therefore, the parameters all remain the same across ranks. So next lecture, [? Tatsu ?] is going to talk about fancier data parallelism, FSDP and 0. And the idea there is, as I've alluded to, here we use all reduce. It's a very simple monolithic operation, but it does require holding all the models, parameters in memory.
[01:02:51.380] But what if the model parameters don't fit in memory? Then you're going to have to be more clever, and that's the topic for the next class. So let me talk about tensor parallelism. So here, the idea is we're going to cut this way and we're not going to cut the data. We're going to cut essentially each layer. And so each rank is going to get part of each layer. And generally, this means that we're going to have to transfer a lot more data. We'll discuss this a bit later. So what does tensor parallel look like? So we're just going to assume we have this data. Every rank has all the data just for simplicity. And here, remember, the data is batch size times num dim. And I'm going to define a local num dim dim to be-- for this rank, I only am responsible for a subset of the dimensions. So the picture here is that each model still has all the layers here for all the layers, but the parameters now are num dim times local num dim.
[01:04:24.200] So if this were one of the parameter matrices for one layer, I would be splitting down the columns. So this is also known as column tensor parallel. You can also do it by rows, but we're not going to talk about that right now. So now what does the forward pass look like? So I'll go through all the layers and I'm going to compute activations. So I start with x the data. And I'm going to access the parameters at layer. And notice that this is only a slice of the parameters. So if I'm on rank 1, I only get this part of the matrix. But I can still proceed. I can apply this non-linearity, because this is element wise anyway. But now what I'm going to do is I'm going to communicate activations. So if I have a data matrix, rank 1 has activations for part of the activations for this part of the matrix.
[01:05:41.210] Rank 1 has part of the activations for this matrix and so on and so forth. And I need to basically put all the activations on all the ranks. But we know how to do that. We introduced all gather as a collective primitive. So this is all the activations. So this is the batch size times local num dim, which is the shape of the activations. And then I'm doing an all gather. Sorry, so this is the allocating memory for the activations. So x is the actual part of the activation. So this is batch size times local num dim. And all gather says each rank has x and each rank is going to allocate activations, which is a list, one for each world size. And then after the all gather, x is going to be copied into each of the respective location activations. So then once I gather all the activations, I concatenate them to form the full dimensional x, which is batch size times num dim. Any questions about column tensor parallel?
[01:07:18.130] So this is done for every layer. We now notice one difference between data parallel is now we have to muck around with the model. Data parallel is very elegant, because it's splitting by data. The model is treated as a module. But now we have to muck around with the model. And this is strongly leveraging the fact that if you want to do a matrix multiplication, you can split it up into a set of smaller matrix multiplications. We can do those on different ranks, and then we can gather the results. When it comes to the back propagation, [INAUDIBLE] gradients are still [INAUDIBLE]. Yeah. So now the question is, what happens in backprop? Now, in the backprop, you have your activations and you have to reduce scatter to all the different gradients. So in some ways, all gather and reduce scatter have this kind of duality, where in forward, you're all gathering, in the backward, you're reduce scattering.
[01:08:33.490] Yeah? [INAUDIBLE] So the question is, is that done automatically by autograd? So none of this-- well, OK. So if you just call .backward, it's not going to do it, because there's no parallelism in that. But PyTorch has all these things that are done automatically for you. [INAUDIBLE] To what extent do I need to change my code and [INAUDIBLE]? So what is done for you and versus automatically? Here we're managing things fairly explicitly, which means I'm not doing the backward pass, but you would have to manage and call the reduce scatter yourself. And that's by design, because this is 336, building language models from scratch. In practice, you probably wouldn't have to do that. OK, so let's do pipeline parallelism quickly. So the idea behind pipeline parallelism is we're going to split the network this way.
[01:09:47.760] So each rank is going to get a subset of the layers now. Within each layer, it's going to get all the dimensions, and it's also going to get all the-- well, one of the ranks is going to get all the data. Every rank is going to see all the data in some form. So this is the way it's going to work. So we have all the data, which is, again, the batch size times num dim. And I'm going to split up the layers. So local num layers is going to be the number of layers that a particular rank is going to handle. And so now I'm going to have local params, which is basically only-- there's local number of layers number of them. But within each layer, I'm going to do the num dim by num dim. So one thing, maybe [INAUDIBLE] will talk more about this on Wednesday, is this idea of micro-batches. So maybe I'll just present this and explain why I'm doing things this way. So in addition to splitting up the layers, I'm also going to split up the batch into a bunch of micro-batches.
[01:11:14.400] So if I'm rank 0, then I get the data, and I'm going to chunk it up into a number of micro-batches. And for each micro-batch, what I'm going to do is receive it from the previous rank, and then do the feed forward pass only on the layers that are assigned to this rank. And then I'm going to send to the next rank. So here I'm actually using these receive and send, which are pointwise operations. So I didn't cover those before, but they're fairly explanatory. This basically says I'm rank and I'm going to send this tensor over to-- sorry, I'm going to receive this tensor from rank minus 1. And this says I'm going to send tensor x to rank plus 1. So basically, the reason why I'm talking about micro-batches is-- and [? Tatsu ?] will talk more about this on Wednesday-- is that in pipeline parallelism, so you have one rank that gets the data. It processes some of the layers, and then it sends it to the next GPU, and it process some of the layers.
[01:12:53.930] And it then sends it to the next GPU. So this is a very natural way of dividing a deep network. But the problem is that you get these what are called pipeline bubbles, where while you're not processing, you're kind of waiting around for other tensors to process. And this ends up being quite inefficient. So the idea behind micro-batches is that you break it up into smaller batches, so you can process it quickly, send it on to the next one. So this can reduce the number of pipeline bubbles. So the other thing I will mention that is not handled in this very kind of naive version is the idea of overlapping communication and computation, which is actually very important to pipeline parallelism. This basically gives you the right structure. If you put an I before these, then it becomes kind of async. You have to add more things to manage the code. And the idea here is that you want to be-- while you're computing here, you can be receiving data or sending data.
[01:14:11.270] So computation and communication should overlap. So that reduces the amount of time you actually spend waiting. So a few things that are kind of missing here, which we'll hopefully fill in next time. So the communication versus computation overlap, which is especially crucial in pipeline parallelism. I didn't mention in data parallelism, this also happens, because I just did a forward pass and then at the end, I'm just doing all these all reduces. But if you're clever, then on the backward pass, as soon as the gradients are done, you can start sending that. And that's something that will be explored in the assignment two. And this, again, allows you to just overlap communication and computation more. We had some questions about what about general models. Again, I think this MLP gives you essentially most of what you need for understanding the basics.
[01:15:27.750] Some of the larger models just require a lot more bookkeeping, so it's harder to see the core algorithms. There are other types of parallelism that we haven't covered. So sequence parallelism takes a whole sequence and chops it up into pieces. And that allows you to parallelize the attention computation, expert parallelism, which allows you to parallelize the experts for MOEs. And this is where the all to all that I mentioned comes in. And then also different combinations of different parallelization techniques, which also will show up in the assignment. So one thing to parallelism technique you choose is going to be strongly dependent on the hardware. For example, tensor parallelism. There's a lot of communication, because for every layer, you need to send all these activations, which are fairly big. So generally, tensor parallelism happens within a node on NVLink or where you have high bandwidth. Whereas you wouldn't do tensor parallelism past an NVLink domain.
[01:16:46.820] Whereas a pipeline parallelism, you'll see people using it. And generally, this can tolerate much slower interconnects. So some of the decentralized training work uses pipeline parallel, because your nodes are-- GPUs are halfway across the world. But you wouldn't want to do tensor parallel in that setting. So sometimes when you look at these combinations, it will be tensor parallel within a node and then pipe data parallel or FSDP and then pipeline parallelism if you need it. There's other effects, such as if you do data parallel, you might be able to do data parallel quite a bit. But then you start hitting to something called the critical batch size, where if you start increasing the batch size too much, it doesn't actually help you, in which case you're kind of wasting your compute and then you're better off using tensor parallel.
[01:17:56.160] So there's a bunch of these considerations, which we'll talk more about as we go through the class. So final note. So TPUs came up a little bit. One thing to note is that on purpose, we are using PyTorch and not just using PyTorch, but really using the collective operations in a very primitive way, so you can see mechanically what's happening. Another approach, especially if you're in [INAUDIBLE] TPU land, is that you can simply define the model and the starting strategy. And the compiler actually handles a lot of the decision of how to-- basically, what kind of communication operations you need. You basically say, well, this piece of data needs to be here and here and here. And then the compiler does some magic to figure out what. So that's appealing. But obviously, it would take a lot of the joy out of actually building things from scratch. Just to summarize, so there's many ways to parallelize.
[01:19:14.660] You can cut by data, cut by tensor or expert, cut by pipeline or sequence. We looked at data parallelism. We only did DDP. Next time we'll do FSDP and 0. Tensor parallelism, as I mentioned, requires very fast interconnects. Pipeline less so. But you need to really work hard to reduce these pipeline bubbles. And then maybe at a high level, we see this pattern come up a lot. So you can either recompute or store in memory when we were talking about things like activation checkpointing or when we're working with GPUs. Or in this case, you can think about an extension of this as that you can store on a different GPU. From that perspective, you look at the data parallel. You're doing redundant work in some sense, because every rank is actually updating its parameters and keeping track of all the parameters. But the reason you're doing that is that you don't have to move the optimizer state across.
[01:20:33.090] So one thing is that hardware is getting faster. But in some sense, we'll always want bigger models. So this idea of having a hierarchical structure will always be there. So that's it for today. So next Wednesday, [? Tatsu ?] will do more of a deep dive on more parallelism techniques.