系统设计 Lab
LLM pretraining 受限于 GPU memory 和 GPU 间的 communication;两者都逼着你把模型切分到很多加速器上。
调整模型 parameter 数、GPU 数、global batch、sequence length、interconnect 带宽、checkpoint 间隔。设计会从单块 GPU 演进到带 all-reduce 的 data parallelism,再到为装下一个对单设备太大的模型而上的 tensor 和 pipeline parallelism,再到带 activation checkpointing 的完整 3D parallelism,最后到上千块 GPU——它们既要快的 collective,又要频繁的 checkpoint-and-recover。
分步讲解
一步一步把它想清楚
要点
常规演进场景
从左到右点击,就是预设的演示路径。每张卡片都会改变 workload 输入。
推荐形态
瓶颈
为什么会变
决策权衡
有出处支撑的规则
这些是模型背后那些经得起时间考验的 system design 论断。而 slider 的具体阈值,则被刻意标注为教学用的假设。
Tensor parallelism 把每层 split 到多块 GPU,训练对单设备太大的模型
Megatron-LM 把 attention 和 MLP 矩阵切分到多块 GPU 上,让单个 transformer 层在合起来的设备 memory 里装得下,并通过快速的节点内 interconnect 交换 activation。
Megatron-LMZeRO 把 optimizer state、gradient、parameter 都 shard 开,降低每 GPU 的 memory 下限
标准 data parallelism 会在每块 GPU 上复制完整的约 16 bytes/param 的 optimizer state;ZeRO 把这份 state 切分到各 data-parallel rank 上,让 memory 随设备数大致线性下降。
DeepSpeed ZeRO (arXiv:1910.02054)Pipeline parallelism 把层切成 stage,用 micro-batch 重叠让 GPU 忙起来
GPipe 把连续的层 stage 分配给不同加速器,并让 micro-batch 流水穿过它们,配合 re-materialization(activation recompute)来约束每个 stage 的 activation memory。
GPipe (arXiv:1811.06965)把 tensor、pipeline、data parallelism 组合起来,能把 transformer 训练扩到上千块 GPU
高效的大模型训练把这三个 parallelism 维度组合起来(3D parallelism),把高带宽的 tensor-parallel 流量放在节点内,把较低带宽的 data-parallel all-reduce 放到跨节点。
Megatron-LM scaling (arXiv:1909.08053)教学用假设
- memory 被建模成约 16 bytes/param(涵盖 weights、gradients、Adam optimizer state),再加一项随 batch 和 sequence length 增长的 activation 项;真实框架会随精度和 ZeRO sharding 变化。
- 每 GPU memory 是一个固定的教学数字(80 GB),collective/data 的阈值都是凑整的数,不是 vendor benchmark。
- failure rate 被近似成与 GPU 数成正比,所以预期损失的工作量随 fleet 规模和 checkpoint 间隔一起放大。