In recent years, we are seeing better results on many NLP tasks with larger pre-trained language models. How to train large and deep neural networks is challenging, as it demands a large amount of GPU memory and a long horizon of training time. There are several parallelism paradigms to enable model training across multiple GPUs, as well as a variety of model architecture and memory saving designs.
近年来，更大的预训练语言模型在许多 NLP 任务上带来更好结果。如何训练大型深度神经网络是巨大挑战——它需要海量 GPU 内存和漫长的训练时间。存在多种跨 GPU 训练的并行范式，以及各种模型架构和省内存设计。本文是分布式训练的奠基性综述——数据并行、模型并行、流水线并行、张量并行，至今仍是训练万卡集群的骨架。
The main bottleneck for training very large models is the intense demand for GPU memory, way above what a single GPU can host. Besides model weights (tens of billions of floats), it's usually even more expensive to store intermediate outputs like gradients and optimizer states (e.g. Adam's momentums and variances). As a result, parallelism is necessary. Parallelism can happen at different dimensions: data, model architecture, and tensor operation.
训练超大模型的主要瓶颈是对 GPU 内存的强烈需求，远超单卡容量。除了模型权重（数百亿浮点数），存储梯度、优化器状态（如 Adam 的动量和方差）等中间输出通常更昂贵。因此并行是必需的。并行可在不同维度发生：数据、模型架构、张量运算。
"优化器状态比权重还贵"是训练大模型时容易被忽视的事实。以 Adam 为例，每个参数要存权重、梯度、动量、方差四份——瞬时内存翻 4 倍。这就是 ZeRO 优化的核心动机。理解内存构成是理解一切分布式训练的前提：权重 + 梯度 + 优化器状态 + 激活值，这四座大山决定了你能否把模型装下。所有并行策略（DP/MP/PP/TP）本质上都是在回答"如何把这四座大山分摊到多张卡上"。
The naive way copies the same weights into multiple workers and assigns a fraction of data to each. Two main synchronization approaches: Bulk Synchronous Parallel (BSP) syncs at the end of every minibatch—prevents staleness and good learning efficiency but each machine has to halt and wait. Asynchronous Parallel (ASP) processes data with no waiting—can lead to stale weights and lower statistical efficiency. Gradient accumulation (in DDP since PyTorch v1.5) syncs gradients globally once every x iterations.
朴素做法是把相同权重复制到多个 worker，各分配一部分数据。两种主要同步方式：批量同步并行（BSP）每个小批次结束同步——防止权重陈旧、学习效率好，但每台机器要停下来等；异步并行（ASP）不等待地处理数据——会导致权重陈旧、统计效率下降。梯度累积（PyTorch v1.5 起的 DDP）每 x 次迭代全局同步一次梯度。
BSP vs ASP 是分布式训练最经典的权衡。BSP（同步）数学等价于单机训练，结果可靠，但慢——所有卡要等最慢的那张（木桶效应）。ASP（异步）快但有 staleness 问题——某卡用旧权重算的梯度可能覆盖新权重的进步。现代主流几乎都选 BSP/近同步，因为 staleness 在大模型上更致命（梯度噪声大）。梯度累积是关键工程技巧：显存装不下大 batch 时，多次小 batch 前向+反向后再一起更新，等价于大 batch——这在显存受限时几乎是必备手段。
Model parallelism solves the case when weights cannot fit into a single node. Since deep networks usually contain a stack of vertical layers, it feels straightforward to split by layer. However, a naive implementation leads to big bubbles of waiting time and severe under-utilization: data is processed by one worker at a time due to sequential dependency.
模型并行解决权重装不下单节点的情况。由于深度网络通常是层堆叠，按层切分看似直接。但朴素实现会导致巨大的等待时间"气泡"和严重的算力浪费：由于顺序依赖，数据一次只能由一个 worker 处理。
Pipeline parallelism combines model parallelism with data parallelism to reduce "bubbles." The main idea is to split one minibatch into multiple microbatches and enable each stage worker to process one microbatch simultaneously. Given m microbatches and d partitions, the bubble fraction is (d-1)/(m+d-1). GPipe observed bubbles are almost negligible if microbatches > 4× partitions (m ≥ 4d).
流水线并行结合模型并行与数据并行来减少"气泡"。核心思想是把一个小批次切成多个微批次，让每个阶段的 worker 同时处理一个微批次。给定 m 个微批次和 d 个分区，气泡占比为 (d-1)/(m+d-1)。GPipe 发现当微批次数 > 4 倍分区数（m ≥ 4d）时气泡几乎可忽略。
PipeDream schedules each worker to alternately process forward and backward passes (1F1B). PipeDream-flush adds a globally synchronized pipeline flush periodically to reduce memory footprint (only one weight version maintained). PipeDream-2BW maintains only two versions of weights ("double-buffered weights").
PipeDream 让每个 worker 交替处理前向和反向（1F1B 调度）。PipeDream-flush 周期性加入全局同步的流水线刷新以减少内存（只维护一份权重版本）。PipeDream-2BW 只维护两份权重版本（"双缓冲权重"）。
流水线并行是训练超大模型的关键拼图之一，它的精妙之处在于用"微批次流水"填满了模型并行留下的气泡。直觉：把模型竖着切成 4 段分给 4 张卡，朴素做法是数据走完卡1→卡2→卡3→卡4，期间其他卡闲着；流水线做法是同时喂入多个微批次，像工厂流水线一样让所有卡同时干活。气泡公式 (d-1)/(m+d-1) 是核心——分区越多气泡越大，微批次越多气泡越小。GPipe 的"m≥4d"是经验法则。1F1B 的进一步优化在于尽早启动反向传播来释放激活值内存，在吞吐和内存间取平衡。
Both model and pipeline parallelisms split a model vertically. Tensor parallelism horizontally partitions the computation for one tensor operation across multiple devices. Megatron-LM adopts a simple way to parallelize intra-layer computation: split the MLP weight matrix A by column [A1, A2], so GeLU(XA) splits into [GeLU(XA1), GeLU(XA2)]. The attention block runs Q/K/V GEMM in parallel per the partitioning. PTD-P (Narayanan et al. 2021) combines pipeline, tensor and data parallelism with interleaved 1F1B scheduling.
模型并行和流水线并行都是竖着切模型。张量并行横向地把单个张量运算的计算分到多个设备。Megatron-LM 用简单方式并行层内计算：把 MLP 权重矩阵 A 按列切分 [A1, A2]，于是 GeLU(XA) 拆成 [GeLU(XA1), GeLU(XA2)]。注意力块按此分区并行运行 Q/K/V 的 GEMM。PTD-P（Narayanan 等 2021）结合流水线、张量、数据并行及交错 1F1B 调度。
张量并行是 Megatron-LM 的招牌，也是当今训练万亿参数模型的标配。核心洞察：矩阵乘法可以按列（或按行）拆分后并行，结果再合并——因为 GeLU 等非线性激活对列拆分是逐元素无关的。TP 的优势是单层内就能并行，延迟低；代价是通信密集（每层都要 AllReduce）。现代大模型训练通常是 3D 并行 = DP × TP × PP：数据并行铺节点，张量并行铺节点内 GPU（靠 NVLink 高带宽），流水线并行铺跨节点。选哪个并行放哪个维度，完全由网络拓扑和带宽决定——这是分布式训练工程师的核心手艺。
The core idea is ensemble learning: combination of multiple weak learners gives a strong learner. A MoE layer contains n feed-forward experts and a trainable gating network G that routes traffic to a few selected experts via noisy top-k gating. GShard scales MoE transformer to 600 billion parameters with sharding. Switch Transformer scales to trillions by routing each input to only one expert (top-1). Expert Choice (EC) routing instead lets each expert select the top-k tokens, guaranteeing fixed capacity and achieving perfect load balancing.
核心思想是集成学习：多个弱学习器组合成强学习器。MoE 层包含 n 个前馈专家和一个可训练门控网络 G，通过带噪 top-k 门控把流量路由到少数选中的专家。GShard 把 MoE Transformer 扩展到 6000 亿参数。Switch Transformer 通过把每个输入只路由到一个专家（top-1）扩展到万亿参数。专家选择（EC）路由反过来让每个专家选 top-k token，保证固定容量并实现完美负载均衡。
MoE 的精髓在于"条件计算"——总参数巨大但每个 token 只激活一小部分。门控路由是 MoE 的命门，它面临两个对立难题：负载均衡（不能让少数专家过载而其他闲置）和容量浪费（专家满了的 token 要丢弃）。三代演进很清晰：原始 MoE（top-2 + 噪声 + 辅助损失）→ Switch（top-1，极简但丢弃多）→ Expert Choice（让专家选 token，完美均衡但需看到未来 token，无法用于自回归生成）。这就是为什么自回归 LLM 仍主要用 token-choice 路由——EC 的完美均衡在生成场景不适用。带噪门控的"噪声"不是 bug 而是 feature：它打破赢家通吃，迫使所有专家都被训练到。
Activation recomputation (checkpointing) reduces the memory cost of training an ℓ-layer net to O(√ℓ), at the cost of one extra forward pass per batch. We evenly divide the network into d partitions; only activations at partition boundaries are saved, and intra-partition activations are recomputed during backward passes. The minimum cost is O(√ℓ) at d=√ℓ.
激活值重计算（检查点）把训练 ℓ 层网络的内存成本降到 O(√ℓ)，代价是每批次多一次前向。把网络均分成 d 个分区，只保存分区边界的激活值，分区内激活值在反向时重算。在 d=√ℓ 时成本最低，为 O(√ℓ)。
激活值重计算是"用算力换内存"的经典操作，几乎所有大模型训练都在用。直觉：前向时不存所有中间激活值，反向需要时再重算一遍。朴素版（全重算）省内存但翻倍算力；分层版（只在检查点存）通过数学优化把额外开销降到约一次前向的 1/√ℓ。O(√ℓ) 的最优解是个优雅结论：分区数 d 太少则省内存有限，太多则重算开销大，d=√ℓ 恰好平衡。实践要点：①通常配合 FlashAttention 等不存注意力的方案；②选择性地对部分层做检查点（如只对注意力层）可进一步减少开销。这是显存不够时的"免费午餐"。
Narang & Micikevicius et al. (2018) train with half-precision (FP16) without losing accuracy, via three techniques: Full-precision master copy of weights (FP32 copy accumulates gradients, since each update may be too small for FP16 range); Loss scaling (scale up loss to handle small-magnitude gradients); Arithmetic precision (accumulate partial results in FP32, save final output as FP16).
Narang & Micikevicius 等（2018）用半精度（FP16）训练而不损失精度，靠三项技术：全精度权重主副本（FP32 副本累积梯度，因为每次更新可能小到超出 FP16 范围）；损失缩放（放大损失以处理小幅梯度）；算术精度（部分结果在 FP32 累积，最终输出存为 FP16）。
Adam needs to maintain momentums and variances—4× the memory of weights. Adafactor only tracks per-row and per-column sums to estimate second moments. ZeRO (Zero Redundancy Optimizer; Rajbhandari et al. 2019) partitions optimizer state, gradients and parameters across data-parallel processes via a dynamic communication schedule, combined with ZeRO-R optimizing residual states (partitioned activations, constant buffers, defragmentation).
Adam 要维护动量和方差——4 倍权重的内存。Adafactor 只跟踪行和列的和来估计二阶矩。ZeRO（零冗余优化器；Rajbhandari 等 2019）通过动态通信调度把优化器状态、梯度、参数分区到数据并行进程上，并结合优化残余状态的 ZeRO-R（分区激活值、恒定缓冲区、碎片整理）。
ZeRO 是大模型训练内存优化的里程碑，它揭示了数据并行中隐藏的"冗余"：朴素 DP 下每张卡都存了完整的优化器状态/梯度/权重——但既然梯度要 AllReduce，何不每张卡只存一份切片？ZeRO 分三阶：ZeRO-1 分优化器状态（省 4 倍）、ZeRO-2 再分梯度、ZeRO-3 再分权重（最省但通信最重）。配合 CPU 卸载（ZeRO-Offload）甚至能在单卡训练 10B+ 模型。ZeRO 与 3D 并行结合，是当今训练万亿参数的事实标准。它的核心哲学：DP 不必复制一切，把"冗余"消除掉就能装下大模型。
把本文所有技术串起来，就是一张"如何把大模型装进集群"的完整地图：四种并行（DP/MP/PP/TP）切分计算与内存，MoE 用条件计算增加容量，激活重计算/混合精度/ZeRO 压缩内存足迹。现代训练框架（Megatron-DeepSpeed、FSDP）本质上是把这些技术的组合自动化。记住一个心法：并行策略由网络拓扑决定——TP 放 NVLink 机内，DP 放跨机，PP 填补气泡。理解这些 2021 年的奠基原理，就能看懂今天万卡训练 GPT-4、Llama 的工程底层。