想象你们班组织期末复习互助小组：老师发下三大本资料——课本、错题本、演算纸。每本都很厚，任何一个人的书包都装不下。于是组长设计了一套规则：三个人分别保管三种资料，谁要用就互相借阅，用完归还。DeepSpeed 就是大模型训练世界里的这个组长，而它的核心方案叫 ZeRO：不是把资料多复印几份，而是让每一张 GPU 卡只保管一部分，需要时再通过通信“借”过来。 这套思路听起来简单，落地却有大量细节：显存到底被谁吃掉了？配置文件怎么写？训练命令怎么启动？为什么卡多了反而更慢？本文全程用“全班共用一套复习资料”的故事，带你把 DeepSpeed 的架构、ZeRO 优化、配置、启动、提速技巧和避坑方法一次讲清。

岗位关联

掌握 DeepSpeed 与 ZeRO，是分布式训练工程师和大模型算法工程师的必备硬技能。它能让你遇到“一张 GPU 放不下模型”的问题时，不再束手无策，而是用多卡协作把模型“搬”起来。

目录
  1. 可视化演示
  2. 核心架构介绍
  3. ZeRO内存优化
  4. 配置文件撰写
  5. 启动训练方式
  6. 训稳提速技巧
  7. 工程避坑指南
  8. AI 应用场景
  9. Python 演示
  10. 常见误区
  11. 面试题

可视化演示

展示DeepSpeed的ZeRO三阶段内存优化策略，通过GPU集群分区可视化理解分布式训练中的显存节省机制。

本节配有交互动效（进入视口后自动加载，约 3KB）。

ZeRO-1 优化器分区 优化器状态均匀切分到4个GPU 显存降低约4倍 通信开销最小

核心架构介绍

期末复习资料堆得比课桌还高，老师发下三大本：课本、错题本、演算纸。任何一个同学的书包都装不下。于是班级决定成立互助小组：大家不各买一套新资料，而是把现有的三套资料分工保管，谁需要用就去借。DeepSpeed 就像这个互助小组的组长，它本身不带来新的知识，只负责设计“怎么分、怎么借、怎么还”的规则，让整个班集体能把大模型训练起来。

放到真实世界里，DeepSpeed 是微软开源的大模型分布式训练框架。它的定位不是一个独立模型，而是 PyTorch 的“总控引擎”。你原本用 model 和 optimizer 训练，现在把这两个对象交给 deepspeed.initialize()，它会返回一个 model_engine。此后前向传播、反向传播、梯度同步和参数更新都由这个引擎统一调度。你的模型结构几乎不用改，DeepSpeed 像一个外挂教务系统，为训练过程排课、点名、收作业。

为什么需要这样一个组长？因为大模型训练时，单张 GPU 卡要同时保存三份“复习资料”：模型参数、梯度、优化器状态。如果使用混合精度（fp16）训练，一个参数就要为这三样东西付出约 16 字节的代价。参数量一上去，单卡立刻爆显存。

所以 DeepSpeed 的价值不是“发明了一种新优化器”，而是当组长，重新分配这些资料。它不会改变模型本身的数学含义，只改变资料放在谁手里。

M单卡≈(2+2+12)Ψ=16Ψ\displaystyle M_{\text{单卡}} \approx (2 + 2 + 12)\Psi = 16\Psi
这里的 \Psi 是模型参数量，16\Psi 表示一个模型会在单张卡上形成 16 倍参数量的显存负担。这就是单卡放不下大模型的直接原因。
  1. 参数 \(2\Psi\)模型权重在半精度下，每个参数约占 2 字节，也就是课本本身。
    2Ψ2\Psi
  2. 梯度 \(2\Psi\)反向传播需要保存与参数同尺寸的梯度，相当于演算纸。
    2Ψ2\Psi
  3. 优化器状态 \(12\Psi\)Adam 需要额外保存 FP32 主权重、动量项和方差项，三项各占 4 字节，像错题本里三种颜色笔记。
    12Ψ12\Psi
  4. 总量 \(16\Psi\)三份资料同时压在一个人书包里，自然背不动。
    16Ψ16\Psi

ZeRO内存优化

组长把全班同学召集起来，提出三步走方案。第一步，4 个人每人只保管错题本的四分之一；第二步，演算纸也按每人四分之一分配；第三步，课本也拆开轮流保管。这样一来，每人书包里的重量大幅下降，而全班加起来的资料总量并没有变，只是少了大量重复复制。

这三步在 ZeRO 里对应三个等级。ZeRO-1 只分区优化器状态，梯度仍然每张卡都完整保留；ZeRO-2 把优化器状态和梯度都按 GPU 数量平分；ZeRO-3 更进一步，把模型参数本身也分到各张卡上。注意，这不是模型并行。ZeRO 依然是数据并行：每张卡都跑完整的网络计算流程，只是在计算某个参数前，需要通过通信把对应参数从其他卡上“借”过来，算完再还回去。

以 ZeRO-3 为例，如果我们有 N 张卡，理想情况下每张卡只需要承担原来 1/N 的模型状态显存：

这就是 ZeRO 的核心哲学：宁可多跑几步路去借，也不让每个人手里都攥着整份复制品。

MZeRO-3(N)=2ΨN+2ΨN+12ΨN=16ΨN\displaystyle M_{\text{ZeRO-3}}(N)=\frac{2\Psi}{N}+\frac{2\Psi}{N}+\frac{12\Psi}{N}=\frac{16\Psi}{N}
当 N=8 时，单卡显存理论上降到原来的八分之一。但代价是：每次计算前要增加通信，卡越多，借资料时在路上花的时间也可能越多。
  1. 参数 \(2\Psi/N\)模型参数平均分到 N 张卡，每张只带一小部分课本。
    2ΨN\frac{2\Psi}{N}
  2. 梯度 \(2\Psi/N\)反向传播产生的梯度也只在本地保存自己的那一片演算纸。
    2ΨN\frac{2\Psi}{N}
  3. 优化器状态 \(12\Psi/N\)Adam 的三份错题笔记同样按 N 平分。
    12ΨN\frac{12\Psi}{N}
  4. 加总 \(16\Psi/N\)三份资料都分出去后，单卡负担从 16Ψ 降为 16Ψ/N。
    16ΨN\frac{16\Psi}{N}

配置文件撰写

方案再好，也要落实到一张“排班表”上。DeepSpeed 的排班表就是 ds_config.json，训练开始前，组长先和白板上的规则对齐，全班才知道自己该带哪部分资料。

在配置文件中，最常用的键包括：zero_optimization 用来设置 ZeRO 等级；fp16 用来开启混合精度；offload_optimizer 可以把优化器状态搬到 CPU 内存；optimizer 指定优化器类型；train_batch_size、gradient_accumulation_steps 共同控制全局 batch。一个典型的配置长这样：

这里最容易误解的是“全局 batch 不等于单卡 batch”。DeepSpeed 的 train_batch_size 是真正意义上一次参数更新看到的样本数，它由三个因数相乘得到：

排班表写的越清楚，训练启动后就越少出现“某某同学不知道自己该干嘛”的情况。

Bglobal=bmicro×g×N\displaystyle B_{\text{global}} = b_{\text{micro}} \times g \times N
一句话总结：全局批量 = 每张卡每次前向的样本数 × 梯度累积步数 × GPU 数量。写配置前先算清楚这三项，否则容易造成实际训练不收敛或资源浪费。
  1. \(b_{\text{micro}}\)单张卡单次前向传播用的微批量，相当于每次小测验的题数。
    bmicrob_{\text{micro}}
  2. \(g\)梯度累积步数，攒多少次小测验才统一改一次笔记。
    gg
  3. \(N\)GPU 数量，就是参与分工的同学人数。
    NN
  4. 三者相乘得到全局 batch，真正决定每一次参数更新的样本规模。
    Bglobal=bmicro×g×NB_{\text{global}} = b_{\text{micro}} \times g \times N
json8 行
{  "train_batch_size": 16,  "gradient_accumulation_steps": 4,  "optimizer": {"type": "AdamW", "params": {"lr": 1e-4}},  "fp16": {"enabled": true},  "zero_optimization": {"stage": 2},  "zero_allow_untested_optimizer": true}

启动训练方式

有了排班表，组长还要喊集合口令。启动 DeepSpeed 训练与普通 PyTorch 训练的区别，就像从“各自在家自习”变成“全班统一到教室集合”。多出来的是启动器、配置文件和统一调度。

代码改造通常只有三步：第一，用 deepspeed.initialize 包装模型；第二，用返回的 model_engine 替换原来的 model；第三，把 loss.backward() 换成 model_engine.backward(loss)，把 optimizer.step() 换成 model_engine.step()。启动命令则从 python train.py 换成：

--num_gpus 指定单机用几张卡，--deepspeed_config 指向排班表；多节点训练时再加 --hostfile，DeepSpeed 会按 hostfile 里的主机列表自动分配角色。--master_port 则类似班级对讲机的统一频道，确保所有节点能找到彼此。

需要注意的是，虽然 DeepSpeed 对 PyTorch 代码侵入很小，但它要求训练脚本能通过命令行参数接收配置。实际项目中，常见做法是把配置路径放到 argparse 里，再用启动器传入，这样代码既能在单卡上调试，也能在多卡上无缝切换。

bash1 行
deepspeed --num_gpus 8 train.py --deepspeed --deepspeed_config ds_config.json

训稳提速技巧

分工完成后，训练效率还取决于同学们的“复习技巧”。省显存只是第一步，训练稳定和速度一样重要。DeepSpeed 最常用的两个技巧，一个是激活检查点，一个是梯度累积。

激活检查点很像“用完草稿纸就扔，需要时再重算”。神经网络训练时，前向传播会保存大量中间激活值，供反向传播使用。如果把这些中间结果全部保留，显存会迅速被占满。激活检查点会舍弃一部分中间结果，等到反向传播真正需要时再重新算一遍，从而用更多计算换回显存。梯度累积则像把几次小测验的错题攒到一起，再统一订正一次笔记，它能让全局 batch 变大，同时减少卡与卡之间的同步次数。参数更新的数学形式可以写成：

实际训练中，还应开启日志监控和吞吐量统计。DeepSpeed 支持打印每步耗时、显存占用和通信占比，配合 TensorBoard 能快速发现瓶颈。如果某一步显存下降不明显，优先检查激活值占用是否仍然过高。

Wt+1=Wt−η×1G∑i=1G∇Li\displaystyle W_{t+1} = W_t - \eta \times \frac{1}{G}\sum_{i=1}^{G} \nabla L_i
这里 G 是累积步数，\nabla L_i 是第 i 个微批量的梯度。程序并不每步都更新参数，而是等 G 个微批量的梯度平均后再更新，等效于用一个更大的全局 batch 训练。
  1. \(W_t\)当前这一步的模型参数，也就是同学们此时手里的核心答案。
    WtW_t
  2. \(\eta\)学习率，控制每次订正笔记的改动幅度。
    η\eta
  3. \(\frac{1}{G}\sum_{i=1}^{G} \nabla L_i\)对 G 个小批量梯度取平均，把零散错题汇总成一份综合报告。
    1G∑i=1G∇Li\frac{1}{G}\sum_{i=1}^{G} \nabla L_i
  4. \(W_{t+1}\)更新后的新参数，全班据此进入下一轮复习。
    Wt+1W_{t+1}

工程避坑指南

互助小组运行一段时间后，总会出现一些新问题：明明把资料分出去了，书包怎么还是那么重？卡从 4 张加到 8 张，速度怎么反而更慢了？这些都是 DeepSpeed 工程落地里最常见的坑。

第一，显存不降反升。可能原因有两个：一是模型本身很小，ZeRO 带来的通信和框架开销盖过了省下的显存；二是只开了 ZeRO，却没开激活检查点，中间激活值仍然吃掉大量显存。显存占用不能只看参数，要全面体检。第二，通信耗时过高。ZeRO 的每个等级都要付出“借资料”的通信成本，尤其是 ZeRO-3，前向和反向都要先 all-gather 参数。卡越多，通信链路越长，小模型上甚至可能出现“借书时间比看书时间还长”的现象。第三，版本兼容问题。DeepSpeed 与 PyTorch、Transformers 库各有版本要求，盲目升级常常导致启动失败。生产环境建议固定版本组合，先在单机双卡上跑通再上多机。

最后，评估参数量时要格外谨慎。配置里的 train_batch_size、gradient_accumulation_steps、zero_optimization.stage 是一套联动的组合拳，只调一个参数很难得到理想效果。像组长把排班表改了却忘记通知同学一样，训练很容易在卡与卡之间产生隐性调度冲突。

---

AI 应用场景

千亿模型多卡预训练

阶段：分布式训练 / 预训练

解决的问题：单张 A100 连 7B 模型都放不下，更不用说千亿参数。传统做法要么用昂贵的多卡张量并行，要么干脆放弃大模型。DeepSpeed 的 ZeRO-3 把参数、梯度、优化器状态全部分区存放，让普通 8 卡服务器也能训练超大模型。

具体案例：某团队要在 8 张 A100 上预训练 30B 参数模型。他们在 ds_config.json 中开启 zero_optimization stage 3，并配合 fp16 与激活检查点，成功把单卡显存压到可用范围，训练吞吐量也比普通数据并行提升明显。

单卡低成本微调

阶段：微调 / 大模型适配

解决的问题：没有多卡集群的开发者，想在 24GB 显存的消费级 GPU 上微调 7B 模型。全部模型状态放在 GPU 上依然不够，最迫切的需求是“借”到 GPU 外面的内存。

具体案例：开发者使用 ZeRO-Offload，把优化器状态和梯度卸载到 CPU 内存，GPU 只负责前向和反向计算。开启后，7B 模型在单张 RTX 3090 上完成 LoRA 微调，速度虽然不如 A100，但成本大幅降低。

多节点大规模协作训练

阶段：训练资源调度

解决的问题：单机 8 卡不够用，需要多台服务器组成训练集群。手动在多台机器上分头启动训练命令，既容易出错，又难统一配置。

具体案例：团队准备一份 hostfile，写入四台机器 IP 和各自 GPU 数量，然后在主节点执行 deepspeed --hostfile hostfile --master_port 29500 train.py --deepspeed_config ds_config.json。DeepSpeed 自动完成节点间通信初始化与资源分配，训练日志也统一汇总。

Python 演示

python59 行
#!/usr/bin/env python3# -*- coding: utf-8 -*-"""ZeRO memory estimator: 按阶段估算大模型单卡显存并生成配置。"""import argparseimport json PARAM_BYTES = 2GRAD_BYTES = 2OPT_BYTES = 12  def stage_memory(psi, num_gpus, stage):    if stage >= 1:        opt = OPT_BYTES * psi / num_gpus    else:        opt = OPT_BYTES * psi     if stage >= 2:        grad = GRAD_BYTES * psi / num_gpus    else:        grad = GRAD_BYTES * psi     if stage >= 3:        params = PARAM_BYTES * psi / num_gpus    else:        params = PARAM_BYTES * psi     return params + grad + opt  def make_config(stage):    return {        "train_batch_size": 16,        "gradient_accumulation_steps": 4,        "optimizer": {"type": "AdamW", "params": {"lr": 1e-4}},        "fp16": {"enabled": True},        "zero_allow_untested_optimizer": True,        "zero_optimization": {"stage": stage},    }  def main():    parser = argparse.ArgumentParser(description="estimate ZeRO memory")    parser.add_argument("--psi", type=float, required=True, help="model size in billions")    parser.add_argument("--gpus", type=int, default=8)    parser.add_argument("--stage", type=int, choices=[0, 1, 2, 3], default=2)    args = parser.parse_args()     psi_bytes = args.psi * 1e9    for stage in [0, 1, 2, 3]:        gb = stage_memory(psi_bytes, args.gpus, stage) / 1e9        print("stage={}: {:.2f} GB per GPU".format(stage, gb))     config = make_config(args.stage)    print(json.dumps(config, indent=2))  if __name__ == "__main__":    main()
目的：输入模型参数量和 GPU 数量，估算 ZeRO 各阶段的单卡显存，并生成一份可直接交给 DeepSpeed 使用的 ds_config.json 配置。
  1. PARAM_BYTES = 2表示半精度参数每个元素占 2 字节。
  2. if stage >= 1: opt = OPT_BYTES * psi / num_gpusZeRO-1 开始把 Adam 优化器状态按 GPU 数分摊。
  3. if stage >= 2: grad = GRAD_BYTES * psi / num_gpusZeRO-2 在此基础上继续分摊梯度。
  4. if stage >= 3: params = PARAM_BYTES * psi / num_gpusZeRO-3 把参数也分片保管。
  5. return params + grad + opt返回当前 stage 下单卡理论显存。
  6. config = make_config(args.stage)生成与当前 ZeRO 等级匹配的 JSON 配置。

常见误区

误区一直以为 ZeRO 是把神经网络切成 N 段，分给 N 张卡分别执行，跟张量并行是一回事。

正解ZeRO 本质上是数据并行下的模型状态分区。每一张卡都跑完整的前向与反向计算，只是参数、梯度和优化器状态被分片存放，按需通信取回。更准确地说，全班同学都要做同一套题，只是每个人只保管一部分参考答案，而不是把一道题切成四段。

为何常见：两者都涉及“把东西分到多张卡”，名字里又带有并行含义，很容易混淆。

误区认为只要把 ZeRO 等级开到 stage 3，再加几张卡，训练速度就会持续上涨。

正解ZeRO-3 每次前向和反向前都要 all-gather 参数，通信开销会随着卡数增长。小模型或小批量下，通信时间可能超过计算时间，最终出现“越多越慢”。是否加卡，要看吞吐量的实测曲线。

为何常见：直觉上认为“人多力量大”，却忽略了交换资料也需要时间。

误区认为配置文件里写上 `"zero_optimization": {"stage": 3}`，显存就会自动降下来。

正解ZeRO 管理的只是模型参数、梯度和优化器状态，中间激活值仍然可能占据大量显存。如果模型很深、序列很长，激活值才是真正的“显存大户”。通常需要配合激活检查点、梯度累积和混合精度，才能看到明显收益。

为何常见：大家习惯只关注 ZeRO 这个明星功能，忽略了其他显存来源。

误区CPU 内存便宜且大，干脆把所有优化器状态都放到 CPU，这样 GPU 就不怕 OOM 了。

正解CPU 与 GPU 之间的 PCIe 带宽比显存带宽低得多，频繁搬数据会导致训练速度骤降。Offload 适合 GPU 显存确实不够、且能接受速度损失的低成本微调场景；如果显存还有空间，优先保留在 GPU。

为何常见：只看“省显存”的好处，却忽略了跨设备搬家的时间成本。

误区觉得引入 DeepSpeed 就要把模型改成特殊结构，像换一套全新框架一样伤筋动骨。

正解DeepSpeed 以 PyTorch 为基础，多数情况下只需要把 model 和 optimizer 传给 deepspeed.initialize，再把 loss.backward() 和 optimizer.step() 换成 engine.backward() 与 engine.step()，模型内部结构无需大规模改动。

为何常见：看到框架名以为是另一个“新世界”，其实它更像是给 PyTorch 装了一个分布式外挂。

面试题

Q1. DeepSpeed 中大模型训练最核心的显存优化手段是什么？

单选题 · easy

  • A.ZeRO优化← 正确答案
  • B.张量并行
  • C.流水线并行
  • D.数据增强

答案：A。DeepSpeed 的核心杀手锏是 ZeRO，它把模型参数、梯度和优化器状态分区存放。张量并行与流水线并行是另一类并行策略，数据增强和显存优化无关。

Q2. ZeRO-1 阶段只把哪一类状态按 GPU 数量平分？

单选题 · easy

  • A.优化器状态← 正确答案
  • B.模型梯度
  • C.模型参数
  • D.训练数据

答案：A。ZeRO-1 只分区优化器状态；ZeRO-2 加入梯度分区；ZeRO-3 才分区模型参数。训练数据不属于模型状态。

Q3. ZeRO-3 阶段中，每个 GPU 上实际保有的模型参数是多少？

单选题 · mid

  • A.参数总量的1/N← 正确答案
  • B.整份完整参数
  • C.仅优化器状态
  • D.仅记录损失值

答案：A。ZeRO-3 把模型参数也按 GPU 数量 N 平分。计算前需要通过通信临时聚合所需参数，计算完成后再释放，而不是让每张卡都保存完整参数。

Q4. 下列哪些 DeepSpeed 配置项可以进一步降低单卡显存？

多选题 · mid

  • A.开启ZeRO优化← 正确答案
  • B.开启混合精度← 正确答案
  • C.优化器放CPU← 正确答案
  • D.加大训练批量

答案：ABC。ZeRO 分区模型状态、混合精度减少参数占用、offload 把优化器搬到 CPU 都能省显存。加大训练批量反而会增加显存占用。

Q5. 训练脚本改造完成后，最常用的启动命令是下列哪一个？

单选题 · easy

  • A.deepspeed train.py← 正确答案
  • B.python train.py
  • C.torchrun train.py
  • D.bash train.py

答案：A。DeepSpeed 启动器是 deepspeed 命令，后面可接 --num_gpus、--hostfile、--deepspeed_config 等参数。普通 python 启动缺少分布式调度逻辑。

Q6. deepspeed.initialize 返回的 engine 对象主要负责什么？

单选题 · mid

  • A.统一前反向与参数更新← 正确答案
  • B.自动清洗训练数据
  • C.生成模型结构图
  • D.编写训练日志文件

答案：A。engine 封装了模型、优化器和 ZeRO 调度逻辑，开发者用 engine.backward(loss) 和 engine.step() 替代原生 PyTorch 流程，数据清洗、绘图、日志都不属于它的核心职责。

Q7. ZeRO-3 在每次前向或反向计算前需要做什么？

单选题 · hard

  • A.临时聚合被分区的参数← 正确答案
  • B.随机丢弃一部分梯度
  • C.临时降低学习率
  • D.重新打乱训练数据集

答案：A。ZeRO-3 把参数分散在多卡上，计算某一层前需要先通过 all-gather 收集该层参数，计算完再丢弃本地不需要的其他层参数。随机丢弃梯度和调整学习率都不是 ZeRO 的机制。

Q8. 激活检查点（Activation Checkpointing）为什么能节省显存？

单选题 · mid

  • A.丢弃中间激活值并按需重算← 正确答案
  • B.删除部分神经网络层
  • C.降低模型参数精度
  • D.关闭反向传播计算

答案：A。激活检查点不删除层，也不改变参数精度，它在反向传播需要中间激活值时重新计算一次，用计算换显存。反向传播仍然正常进行。

Q9. 关于梯度累积，下列说法正确的有？

多选题 · hard

  • A.可扩大全局批量← 正确答案
  • B.可降低通信频率← 正确答案
  • C.能提升单卡显存
  • D.能减少模型层数

答案：AB。梯度累积把多个微批量的梯度累加后再统一更新，等效于扩大全局 batch，同时减少卡间同步次数。它不会提升单卡显存，也不会改变模型结构。

Q10. 在 8 张 A100 上训练千亿参数模型，最稳妥的配置是？

单选题 · extreme

  • A.ZeRO-3加offload← 正确答案
  • B.ZeRO-1只分优化器
  • C.只开启梯度裁剪
  • D.只增加训练轮次

答案：A。千亿参数下仅 ZeRO-1 或仅 fp16 通常不够，ZeRO-3 把参数、梯度、优化器全部分区，必要时配合 offload 和激活检查点。梯度裁剪和增加轮次解决的是训练稳定性或效果问题，不是显存问题。