DeepSpeed

一句话定位

DeepSpeed 是围绕大模型训练的执行与内存优化引擎，重点包括训练状态切分、并行和大规模训练效率。

它负责

  • ZeRO 等参数、梯度和优化器状态切分
  • 与 PyTorch 模型和 optimizer 集成
  • 混合精度、offload、pipeline 等训练能力

它不负责

  • 不定义 Transformer 数学结构
  • 不替代 PyTorch 的 Tensor 与 autograd
  • 不等同于 Ray 的 task/actor 和角色编排
  • 不等同于 LoRA 这类参数高效算法

配置选择的因果链

1
2
3
4
5
显存账本与模型规模
→ 选择 DDP / ZeRO-1/2/3 / offload
→ 观察通信与 step time
→ 调整 bucket、overlap、micro-batch
→ 用 checkpoint 与恢复流程验收

DeepSpeed 是执行引擎，真正的训练目标仍由 PyTorch、数据管线和上层算法决定。配置文件能启动不代表有效：要同时记录有效 batch、吞吐、峰值显存、通信占比和 loss 数值一致性。若只是单机小模型，额外引入 engine 可能增加调试面而没有收益。

相关节点：ZeRO、Training Memory Accounting。

使用边界

源仓库的 DeepSpeedEngine.initialize 示例不应照搬；实际入口和配置需要按当前版本核验。单机多卡不必为了“分布式”强行引入 Ray。

ZeRO 配置为什么会改变 step time

ZeRO 不是简单的“把模型切小”：它在不同阶段切分不同训练状态。ZeRO-1 主要切 optimizer state，ZeRO-2 再切 gradient，ZeRO-3 连参数也按需分片/聚合。切得越深，单卡常驻显存越低，但 forward/backward 周围的参数 gather、gradient reduce-scatter 和 checkpoint 处理越复杂；当互联或网络较慢时，省下的显存会被通信等待抵消。

1
2
3
4
5
显存不够
→ 从 ZeRO-1 提升到 ZeRO-2/3 或启用 offload
→ 记录每 step 的通信/计算重叠
→ 调整 micro-batch、gradient accumulation、bucket
→ 在相同有效 batch 下比较吞吐和 loss

比较时必须保持有效 batch、序列长度、梯度累积和 optimizer 不变。只把 micro-batch 调大后报告吞吐，可能同时改变了优化轨迹；只看显存下降，也可能隐藏通信和 CPU offload 的等待。

故障与恢复不是附属功能

分片训练的 checkpoint 不只是一个完整的 state_dict。保存时要记录 world size、ZeRO stage、optimizer、scheduler、随机数状态、数据游标和参数配置；恢复测试要验证 loss 曲线连续、optimizer step 没回退、所有 rank 都能重新 join。常见错误是“训练能跑但恢复后模型不一致”，根因可能是只保存了参数、数据 sampler 没恢复或版本改变了 partition 规则。

如果显存只差一点，先比较 activation checkpointing、序列并行、gradient accumulation 与 ZeRO/offload 的代价；DeepSpeed 是手段，不会替上层算法决定该牺牲吞吐、通信还是精度。