Appearance
Q128 · 在训练一个百或千亿参数级别的 LLM 时，你会面临哪些主要的工程和算法挑战？
假设团队要从头训练一个 100B（1000 亿）参数的文本大模型，使用 BF16 参与主要计算、Adam 类优化器更新权重。刚把模型权重放进显存，单张 GPU 就已经装不下；把训练分到许多卡上，又发现卡与卡之间等待通信、数据读取跟不上、某一步损失突然变成异常值。即使前面几天都正常，一次机器故障也可能迫使训练重来。
所以这不是单独解决“显存”就完结的题。训练要同时满足：模型状态和中间结果放得下、多卡能高效协作、数据持续喂得上、数值更新保持稳定、故障后能续训，而且最终质量可检验。这篇先搭建全局判断框架；Q129 会细讲单卡装不下的原因，Q130 专讲显存优化技术，Q131 专讲长时间训练的容错。下面所有硬件容量与训练配置均是教学假设，不是对某个集群可行性的承诺。ZeRO 原论文 · Megatron-LM 大规模训练论文
术语与符号
| 术语或符号 | 先用日常话解释 | 100B 例子 |
|---|---|---|
| LLM | 从大量文本中学习词语关系、可生成文本的大语言模型 | 要训练的 100B 文本模型 |
| 参数 / 权重 | 模型在训练中不断调整的数字 | 共有 N = 100,000,000,000 个 |
| BF16 | 每个数用 16 位，即 2 字节存放的一种浮点格式；精度与数值范围有取舍 | 计算用权重按 2 字节/参数估算 |
| FP32 | 每个数用 32 位，即 4 字节存放的浮点格式 | 本例的主权重和 Adam 状态 |
| 梯度 | 当前一批训练数据指出“权重应往哪边改”的数值 | 反向传播后每个参数的更新信号 |
| Adam 优化器状态 | 记录梯度的两种历史统计量，便于控制更新方向与幅度 | 每参数两份 FP32 状态，共 8 字节 |
| 主权重副本 | 为更新保留的一份高精度参数；并非所有实现都这样存 | 本例按每参数 4 字节计 |
| 激活值 | 前向计算过程中的中间结果，反向传播常需用它 | 会随序列长度、批量和网络结构变化 |
| KV | 注意力层中 Key、Value 的缩写；训练时相关张量属于计算中间量 | 不把推理服务的持久 KV 缓存另算进本例静态状态 |
| 显存 / HBM | GPU 上容量有限、速度很快的内存 | 一张假设为 80 GB 的卡 |
| 数据并行 | 多卡各处理不同训练样本，再同步梯度 | 扩大吞吐，但普通做法会复制模型状态 |
| 张量并行 | 把同一层的大矩阵计算拆到多卡 | 单层太大时按维度切分 |
| 流水线并行 | 把不同网络层分给不同卡组 | 前几层和后几层分别运行 |
| ZeRO / FSDP | 把权重、梯度、优化器状态的副本按策略分片的技术家族 | 减少每卡重复存储，但引入通信与临时聚合 |
| 检查点 | 定期保存可恢复的训练状态 | 故障后从第 30 万步附近继续 |
| 损失 / loss | 训练预测与目标的差异量；越低通常表示拟合训练目标更好 | 监控是否突然升高、变成 NaN |
| NaN / Inf | “不是一个有效数字” / “无限大”的计算异常 | 表示某些数值更新已失控 |
这里的 100B 是参数数量，不是 100 GB 显存。同一参数在训练中可能以多个形态出现；训练还要保留某些中间结果。图示只画“模型状态切分与通信”的主关系，数据、数值稳定和故障恢复在正文展开。
图里的四张卡只代表“多卡切分”这个机制，不意味着四张卡就能训练 100B 模型；实际卡数取决于显存容量、切分方式、激活、通信缓冲和余量。
第一关：先把训练显存账算完整
为了能复述数量级，明确采用一个特定的混合精度 Adam 配置：计算用 BF16 权重每参数 2 字节；保留 BF16 梯度每参数 2 字节；优化器持有 FP32 主权重 4 字节，另有 Adam 一阶、二阶历史统计量各 4 字节。合计 2 + 2 + 4 + 4 + 4 = 16 字节/参数。这个口径与 ZeRO 论文对 16 位参数、16 位梯度、32 位主权重和两份 Adam 状态的推算相同；论文原例用 FP16，本题把同样的字节宽度假设用于 BF16。ZeRO 论文第 3.1 节
| 静态模型状态 | 每参数字节 | 100B 总量，按十进制 |
|---|---|---|
| BF16 权重 | 2 | 200 GB |
| BF16 梯度 | 2 | 200 GB |
| FP32 主权重 | 4 | 400 GB |
| Adam 一阶状态 | 4 | 400 GB |
| Adam 二阶状态 | 4 | 400 GB |
| 合计 | 16 | 1,600 GB = 1.6 TB |
这些是全模型静态状态的理论数量级，不包括激活、注意力中间张量、通信缓冲、算子工作区、显存碎片和运行时对象。也不代表每张卡都必须存 1.6 TB：分片后可分摊；相反，简单数据并行若每卡保留完整副本，只增加卡数也不能让单卡从根本上装下模型。ZeRO 论文明确将“模型状态”和“剩余内存消耗，如激活”分开分析。ZeRO 论文第 3 节
假设都是 80 GB 显存的卡，1,600 ÷ 80 = 20 只是静态字节总量的算术下界：20 张卡即使理想平均切分也已占满，没有一字节留给激活和缓冲，不能据此说“20 卡可训”。若理想地把 1.6 TB 静态状态均匀切成 32 份，每卡是 50 GB，仍要验证临时聚合、最长序列和实际峰值是否装得下。有的实现把梯度累积为 FP32，或不保留本例的主权重副本，状态总量就会改变；必须按真实框架配置测峰值，不能把 16 字节/参数当普适常数。Megatron Core：Distributed Optimizer
训练中的 Key、Value 和其他注意力张量会占用激活/中间计算显存。它们跟序列长度、每卡微批量和是否重算有关。在线推理常说的“KV 缓存”是为了生成后续 token 而跨解码步保留的状态，不能把推理服务的缓存公式又加一次到上述训练静态状态里。若激活成为瓶颈，可以缩短序列、减小每卡微批量，或用激活重计算：前向少存一部分中间量，反向需要时重新计算，交换显存与额外算力。PyTorch：Activation checkpointing techniques
第二关：多卡切分解决容量，也带来通信账
数据并行让不同卡读不同样本，吞吐容易扩展；传统方式每卡留完整权重、梯度和优化器状态，并在反向后同步梯度。因此它主要解决“更多样本并行处理”，未必解决“单卡放不下 100B”。ZeRO/FSDP 则在数据并行组里进一步切分状态：可从只分优化器状态，逐步扩展到梯度和参数。每卡存得更少，但某层计算需要完整或局部权重时，要按需聚合；更新后又要分发。PyTorch 对 FSDP 的说明就是将参数、梯度和优化器状态在数据并行工作进程间分片。ZeRO 原论文 · PyTorch：FSDP
张量并行把单层矩阵拆到多卡，流水线并行把不同层分给不同卡组；Megatron-LM 展示了数据、张量、流水线三种并行组合训练更大模型。切分边界决定通信：张量并行通常要在层内频繁交换部分结果；流水线并行要在阶段间传激活和梯度，并可能有某些阶段空等的“气泡”；状态分片常见梯度 reduce-scatter 和参数 all-gather。在这里，前一个词指“把多个卡的梯度汇总后按份分给各卡”，后一个指“把各卡手里的参数份额收齐给需要它们的计算”。Megatron-LM 论文 · Megatron Core：Parallelism guide · Megatron Core：Distributed Optimizer
所以方案不是“卡越多越快”。若卡间链路慢，计算每层 20 毫秒却等通信 30 毫秒，新增 GPU 可能让更多资源空等。需要按真实拓扑测每步时间：数据读取、前向、反向、同步、优化器更新分别占多少；尽量让高频交换发生在高速互联的卡组内，调节切分粒度和微批量，并在正确性验证后再尝试通信与计算重叠。优化通信不能牺牲必需的梯度同步，否则各卡参数会悄悄分叉。
第三关：GPU 等数据和“数据够不够好”是两件事
即使模型和通信都配置好，训练数据若读盘、解压、分词或跨机器传输太慢，GPU 仍会空等。分词是把文本转成模型接受的 token 序列；数据加载器负责读样本、混合来源、切成训练长度并送到卡上。大规模运行要预处理与索引数据、为不同进程分片、控制缓存和读取并发，记录样本位置以便故障恢复。Megatron Core 的数据加载指南专门讨论多节点规模下索引、缓存、预取等瓶颈。Megatron Core：Data loading at scale
数据吞吐快也不等于训练有效：重复或低质量文本会浪费算力，测试集混入训练集会让评估虚高。训练前要制定数据配比、去重和质量过滤，保留独立的验证集；训练中监控每类数据的采样比例、实际处理的 token 数和验证效果。Chinchilla 研究说明，在固定算力预算下，模型规模与训练 token 数量需要一起权衡；不能只把参数堆到 100B，却用不足或失衡的数据训练，然后怪罪分布式系统。Hoffmann 等：Training Compute-Optimal Large Language Models
第四关：数值稳定靠观察与定位，不能只调一个开关
训练每一步都用梯度修改大量参数。学习率过高、某批异常数据、精度范围/舍入、错误的归一化或并行同步，都可能让损失突然飙升、梯度范数异常，甚至出现 NaN/Inf。BF16 的数值范围有优势，但不是“用了 BF16 就不会炸”。Megatron Core 提供 BF16/FP16 配置、全局梯度范数裁剪等选项；这些是可用手段，具体阈值需要在目标模型上测，不该照搬一个数。Megatron Core：Optimizer config · Megatron Core：Gradient clipping
正常路径是：先在较小模型和短序列上跑通数据、正反向与检查点；逐步扩大并行规模，监控训练/验证损失、梯度范数、有效 token 吞吐、各卡显存峰值和异常值。出现尖峰时，先定位是特定数据批次、学习率阶段、某个算子还是通信后各卡结果不一致；必要时停更、回到健康检查点，验证根因后调整学习率、裁剪或精度配置。若只在 loss 曲线上看到一个尖峰就盲目“调低学习率”，可能掩盖坏数据或实现 bug，也可能拖慢收敛。论文研究过大模型预训练中的 loss spike 现象，说明稳定性是独立于存储容量的挑战。Spike No More: Stabilizing the Pre-training of Large Language Models
第五关：长跑需要能恢复，还要知道是否学到了东西
100B 训练可能持续很久，机器、网络、存储任何一环都会出故障。检查点至少要让权重和优化器状态恢复；为了尽量接续同一训练过程，还应记录学习率调度器阶段、全局步数、随机数状态、数据读取位置、数据版本和并行配置。大规模分片检查点需要校验所有分片完整，再将其标为可恢复；不能在只写完一半时覆盖最后一个好版本。恢复演练要真的执行一次：从已存检查点继续跑几步，核对损失、步数和样本进度。PyTorch 的分布式检查点 API 支持面向分片模型和优化器状态的保存/加载；保存格式与可移植性仍取决于配置。PyTorch：Distributed checkpoint
仅能续训也不够。团队应在固定版本的验证集上观察损失与下游任务表现；验证集与训练集分离，避免数据泄漏。记录训练 token、有效 GPU 时间和每次评估的配置，才能判断增加规模或改变数据配比有没有带来质量提升。训练损失下降不保证对真实用户更有帮助，速度提高也不能替代质量评测。
把一次正常运行和一次故障连起来
在教学假设中，团队先算出 1.6 TB 静态状态，选择足够多卡的状态分片与模型并行组合，并通过小规模试运行测得各卡还有激活余量。数据管线每步及时交付已分词文本，验证损失缓慢下降，梯度范数没有持续异常，检查点周期性写入且可恢复。此时“能跑”才逐渐接近“值得继续跑”。
反过来，若最长序列批次突然 OOM，不能只引用“平均每卡 50 GB 静态状态”解释安全：峰值激活与临时参数聚合可能超过余量。若换了分片设置后 OOM 消失，却每步大半时间都在通信，吞吐仍不达标。若某次恢复后 loss 跳变，要检查优化器状态、调度器和数据进度是否也恢复，不能只加载权重就宣称无损续训。这几类失败分别对应容量、通信和状态一致性，处理顺序不同。
面试时怎样回答
我会先做显存与吞吐预算。以 100B 参数、16 位权重和梯度、32 位主权重与 Adam 两份状态为例，静态模型状态约每参数 16 字节，总计 1.6 TB，激活和通信缓冲另算；所以要用状态分片，并按层或张量结合模型并行，实际卡数靠峰值测量确定。切分后再看跨卡同步、流水线空等和数据加载是否拖慢每步。算法侧要监控 loss、梯度范数与 NaN，排查学习率、精度、异常数据和实现错误。长时间训练还要保存可恢复的模型、优化器和数据进度，实际做恢复演练；最后用独立验证集判断质量，而非只看训练速度。并行、显存、通信、稳定性和数据质量需要一起权衡。
追问“为什么不是 100B × 2 字节就够”时，答权重只是训练状态的一部分，还有梯度、优化器统计量、可能的主权重，以及与批量和序列有关的激活。追问“多加 GPU 为什么还会慢”时，答切分产生梯度和参数同步，链路与流水线空等可能吞掉计算收益。追问“怎样验证方案”时，答先用短跑量显存峰值和分步耗时，再测试异常批次、长序列与检查点恢复，最后看验证集质量。
参考资料
- ZeRO: Memory Optimizations Toward Training Trillion Parameter Models：混合精度 Adam 状态内存及分片。
- Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM：数据、张量和流水线并行组合。
- Megatron Core：Parallelism Strategies Guide
- Megatron Core：Distributed Optimizer
- PyTorch：FSDP与Activation Checkpointing
- Megatron Core：Data Loading at Scale
- PyTorch：Distributed Checkpoint
- Training Compute-Optimal Large Language Models