[CS336-07] Parallelism Basics

这篇是 Stanford CS336 Spring 2025 第 7 讲 Course Materials 的学习笔记。主题是大模型训练里的并行化基础。

1. 为什么必须多卡多机#

单张 GPU 的算力和显存都有上限，而现代 LM 的参数、activation、optimizer states、训练 tokens 都在增长。因此大模型训练需要把计算和状态切到多张 GPU，甚至多个机柜。

目标有两个：

• memory scaling：模型和训练状态能随 GPU 数增加而放得下；
• compute scaling：总训练 FLOPs 能接近随 GPU 数线性增长。

难点是通信。GPU 之间同步数据比 GPU 内部计算慢得多，因此并行策略的核心是减少、隐藏或重叠通信。

2. Collective communication 是基本积木#

Lecture 7 先讲 collective：broadcast、reduce、all-reduce、all-gather、reduce-scatter 等。

这些操作是分布式训练的底层语言。例如 data parallel 需要 all-reduce gradients；ZeRO 可以用 reduce-scatter 和 all-gather 分片 optimizer states 和参数；tensor parallel 需要在层内通信 partial results。

一个重要事实是：all-reduce 经常可以分解成 reduce-scatter 加 all-gather。在带宽受限场景下，这种分解不仅等价，还帮助理解 ZeRO 的设计。

3. Data Parallelism 和 ZeRO#

Naive data parallel 的做法是每张 GPU 保存完整模型，处理不同 batch shard，然后同步梯度。

它的 compute scaling 好，但 memory scaling 差，因为每张 GPU 都保存完整参数、梯度、optimizer states。

ZeRO 的核心是把这些状态切分：

• ZeRO-1：shard optimizer states；
• ZeRO-2：再 shard gradients；
• ZeRO-3：进一步 shard parameters。

越往后 memory saving 越强，但通信和实现复杂度也越高。

4. Model Parallelism#

当模型本身单卡放不下，或者单层计算太大时，需要 model parallel。

Pipeline parallel 把不同层放到不同设备上，像流水线一样处理 micro-batches。它能切分模型深度，但会有 pipeline bubble，并且需要处理 activation 传递。

Tensor parallel 把单层矩阵乘法按维度切开，让多个 GPU 一起算同一层。它适合单层很大的 Transformer，但层内会引入频繁 collective。

实际大规模训练常常组合 data parallel、tensor parallel、pipeline parallel，有时还会加 sequence parallel。

5. 拓扑决定并行策略#

并行策略不能脱离硬件。单机内 NVLink/NVSwitch 的带宽和延迟，与跨机网络完全不同。TPU mesh、GPU all-to-all fabric、InfiniBand/RDMA 也会带来不同通信模式。

因此工程上经常会让 tensor parallel 留在高速互联的单机内，把 data parallel 放到跨机层面。这样可以把高频通信放在快链路上，把低频大通信放到更外层。

6. takeaway#

第 7 讲的核心是：大模型训练不是简单“多开几张卡”。

• data parallel 解决吞吐，但不解决完整模型复制；
• ZeRO 用 state sharding 改善 memory scaling；
• tensor parallel 和 pipeline parallel 切模型，但引入不同通信；
• 硬件拓扑决定哪种并行组合更合理。

后面第 8 讲会用 PyTorch distributed 把这些 collective 和训练流程写成代码。

参考#

• Stanford CS336 Spring 2025 Course Materials: lecture 7.pdf