中文 EN

系统设计 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 输入。

Workload

这些是输入,不是预设好的架构阶段。

推荐形态

当前架构路径
LLM pretraining 基础设施图 大规模 LLM pretraining 的白板风格架构图:放在 object storage 里的训练 dataset、sharded data loader、由 data/tensor/pipeline 组构成的 3D-parallel training mesh、collective-communication fabric 和 checkpoint store,以及异步的 metrics 和 fault recovery。 Dataset Data loaders Training mesh Comms + checkpoint Ops Token corpus object storage Data loaders stream + shuffle Data-parallel replicas Tensor-parallel split layers Pipeline-parallel layer stages Collective fabric all-reduce Checkpoint store sharded state Metrics loss + throughput Fault recovery restart mesh
Dataset
Token corpus 存在 object storage 里的 tokenized 训练 shard,run 期间被 stream 出来
Data loaders
Data loaders prefetch、shuffle 并打包 token batch,让 GPU 永远不会等输入
Training mesh
Data-parallel 每个 replica 在 batch 的一片上训练,每步对 gradient 做 all-reduce
Tensor-parallel 把每层矩阵 split 到多块 GPU 上,它们在节点内交换 activation
Pipeline-parallel 把连续的层 stage 分配给各 GPU,在它们之间 stream micro-batch
Comms + checkpoint
Collective fabric NVLink 和 InfiniBand 承载 gradient all-reduce 和 tensor-parallel 交换
Checkpoint store 持久化 sharded 的 weights 和 optimizer state,好让 run 在故障后能恢复
Ops
Metrics 在训练热路径之外记录 loss、throughput 和硬件健康状况
Fault recovery 检测挂掉的节点,并从最新的 checkpoint 重启 mesh

瓶颈

每 GPU memory 压力

Collective comms 负载

Data pipeline 负载

Activation memory

Checkpoint / 故障暴露

为什么会变

    决策权衡

    Parallelism 策略

    Gradient sync / collectives

    Activation memory

    Checkpoint + fault tolerance

    Data pipeline throughput

    Interconnect topology

    有出处支撑的规则

    这些是模型背后那些经得起时间考验的 system design 论断。而 slider 的具体阈值,则被刻意标注为教学用的假设。

    已验证的规则

    Tensor parallelism 把每层 split 到多块 GPU,训练对单设备太大的模型

    Megatron-LM 把 attention 和 MLP 矩阵切分到多块 GPU 上,让单个 transformer 层在合起来的设备 memory 里装得下,并通过快速的节点内 interconnect 交换 activation。

    Megatron-LM
    已验证的规则

    ZeRO 把 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 间隔一起放大。