2、大模型训练基础：模型规模与计算需求、内存瓶颈、分布式训练概述
在正式接触 DeepSpeed 之前，必须先建立对大模型训练“成本结构”的直觉。很多初学者一上来就研究 ZeRO 的 Stage 划分，却说不清显存到底被什么吃掉了，结果调参时只能凭感觉。本节把训练一个 Transformer 大模型所需的算力与内存拆开来看，理解瓶颈在哪里，才能明白 DeepSpeed 究竟在优化什么。
一、模型规模：参数量、层数与隐藏维度
模型规模通常用参数量衡量，而参数量由结构超参决定。以标准 Decoder-only Transformer 为例，主要超参包括层数 L、隐藏维度 d、注意力头数 h、词表大小 V 与序列长度 S。粗略估算：
每层注意力参数 ≈ 4 * d * d
每层FFN参数 ≈ 8 * d * d (中间维度取 4d)
嵌入参数 ≈ V * d
总参数 ≈ L * 12 * d^2 + V * d可以看到，参数量对隐藏维度是平方级增长。d 从 4096 翻到 8192，单层参数就变成 4 倍。这就是为什么 7B、13B、70B、175B 之间的资源需求不是线性上升，而是急剧攀升。
二、计算需求：FLOPs 与训练时间估算
训练一步的计算量可用经验公式估算：前向加反向约等于 6 * N * T，其中 N 是参数量，T 是 token 数。若训练总 token 数为 D，则总计算量约为 6 * N * D FLOPs。
举例：一个 7B 模型，训练 1T token，总计算量约 6 × 7e9 × 1e12 = 4.2e22 FLOPs。若单卡有效算力为 300 TFLOPS，理想情况下需要约 1.4e8 秒，即约 4.4 年。这显然不可接受，必须依靠数据并行、张量并行等多卡协同，把时间压缩到可接受范围。
三、内存瓶颈：训练显存到底被谁占用
训练显存主要由四部分构成，这是理解 DeepSpeed 优化的核心：
- 模型参数（Parameters）：FP16 下每个参数 2 字节，7B 模型约 14GB。
- 梯度（Gradients）：与参数同量级，FP16 约 14GB。
- 优化器状态（Optimizer States）：Adam 需保存一阶、二阶动量，FP32 下每个参数 8 字节，7B 模型约 56GB。
- 激活值（Activations）：与 batch size、序列长度、层数强相关，长序列时可能超过参数本身。
把上述相加，7B 模型混合精度训练仅静态部分就接近 84GB，远超单张 80GB 显卡容量。这就是“单卡放不下”的真实原因，也是 ZeRO 要切分优化器状态、梯度、参数的动机。
四、分布式训练概述：三条并行路线
当单卡放不下、算力不够时，分布式训练是唯一出路。主流并行策略有三类：
| 并行方式 | 切分对象 | 通信特点 | 典型适用 |
|---|---|---|---|
| 数据并行 DP | 切分数据 batch | 每步梯度 AllReduce | 模型能放进单卡 |
| 张量并行 TP | 切分单层权重矩阵 | 层内高频 AllReduce | 单层过大，节点内 NVLink |
| 流水并行 PP | 切分模型层到不同设备 | 阶段间点对点传输 | 层数多，跨节点 |
实际训练中常采用“3D 并行”：TP 放在节点内，PP 跨节点，DP 在最外层。但 3D 并行配置复杂、通信开销大，对中小团队并不友好。DeepSpeed 的 ZeRO 提供了一条更平滑的路径：在数据并行的框架下，把参数、梯度、优化器状态分片到各卡，从而在不改变模型代码结构的前提下大幅降低单卡显存。
五、常见误区与复盘
- 误区一：只看参数量估算显存。 优化器状态往往是参数量的 4 倍以上，忽略它会严重低估需求。
- 误区二：认为开混合精度就能省一半。 FP16 只压缩参数与梯度，优化器状态通常仍为 FP32。
- 误区三：盲目上 3D 并行。 通信与调试成本高，很多场景 ZeRO-2/ZeRO-3 已足够。
- 误区四：忽略激活值。 长序列、大 batch 时激活值可能成为新的瓶颈，需要激活重计算配合。