DeepSpeed ZeRO-Offload与ZeRO-Infinity：单卡训练1000亿参数大模型的完整指南

DeepSpeed ZeRO-Offload与ZeRO-Infinity：单卡训练1000亿参数大模型的完整指南

【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed

🎯 本文带你完整解析DeepSpeed（微软开源的深度学习优化库）的两大核心内存优化技术——ZeRO-Offload与ZeRO-Infinity。它们让大模型分布式训练和推理变得简单、高效：零代码修改，仅修改几行 JSON 配置，即可用单张 GPU 训练 100B（1000亿）参数的 GPT-2 大模型。文章面向新手，图解原理 + 配置实操 + 进阶路线，一文掌握。

一、为什么大模型训练会"爆显存"？

训练一个 1000 亿参数的大模型，显存里装的不只是权重。以 FP16 训练 + Adam 优化器为例，每个参数背后还藏着 32 位权重副本、动量和二阶矩等状态：

一张 32GB 的 V100 显然装不下。传统方案是模型并行——但需要大改模型代码，复杂且通信开销高。而 DeepSpeed 的ZeRO（Zero Redundancy Optimizer，零冗余优化器）思路完全不同：

数据并行的每张卡上都有一份冗余的优化状态？那就把它们切分到所有设备上，各存各的、各算各的，用时再聚合。

这样不仅省显存，还不需要修改任何模型代码——只需在 DeepSpeed 的 JSON 配置中加几行参数即可，这是 ZeRO 相比模型并行的最大吸引力。

二、ZeRO 三级进阶：从 Stage 1 到 Stage 3

ZeRO 以渐进阶段（stage）实现优化，后期阶段包含前期全部优化（详见官方教程 docs/_tutorials/zero.md）：

• Stage 1 — 切分优化器状态：每张卡只保存并更新自己那一份优化器状态。官方实测：1.5B 参数 GPT-2 的 Adam 状态从 18GB 降到 2.25GB（8 卡）；
• Stage 2 — 再切分梯度：16 位梯度也按卡切分，每卡只保留自己负责部分的梯度；
• Stage 3 — 连参数也切分：16 位模型权重同样切分到所有卡上，前向/反向传播时自动聚合与释放，单卡显存占用降至极限。

实现位于 deepspeed/runtime/zero/ 模块，其中 stage_1_and_2.py 负责前两级，stage3.py 负责第三级的参数切分与通信。

三、ZeRO-Offload：把优化器"搬"到 CPU 上

单卡训练时没有"其他卡"可分摊，怎么办？ZeRO-Offload 的答案是：利用主机 CPU 的内存和算力。

它把优化器的状态与计算从 GPU 卸载到 CPU，由 DeepSpeed 高度优化的 CPU 版 Adam（DeepSpeedCPUAdam，源码见 csrc/adam/cpu_adam.cpp）接管更新——比 PyTorch 标准实现快 5~7 倍，避免 CPU 成为瓶颈。

官方教程实测：单张 32GB V100 + 32GB 内存的机器，成功训练100 亿参数 GPT-2。下面是真实训练日志截图：

配置方法极其简单，在 DeepSpeed JSON 配置中开启zero_optimization并将优化器设备指定为cpu：

{ "zero_optimization": { "stage": 2, "offload_optimizer": { "device": "cpu" }, "contiguous_gradients": true, "overlap_comm": true } }

就这几行，无需任何代码修改。多卡训练时还可在deepspeed启动命令中加--bind_cores_to_rank，把物理 CPU 核均匀分配给各 rank，让 CPU 优化器性能接近最优。完整实操步骤见 docs/_tutorials/zero-offload.md。

四、ZeRO-Infinity：GPU → CPU → NVMe 三级存储，突破内存天花板

如果 CPU 内存也不够呢？ZeRO-3 内置了infinity offload 引擎（实现见 deepspeed/runtime/zero/offload_config.py），构成ZeRO-Infinity：

它把模型状态扩展到GPU 显存 → CPU 内存 → NVMe SSD三级存储层级，通过**智能调度（Smart Scheduling）决定哪些状态驻留在哪一层，并支持预取（Prefetch）**提前把下一层数据搬上来、内存池管理保证连续大块显存、通信优化隐藏传输延迟。最终结果：内存需求无限逼近"理论下限"，模型规模不再受任何单一硬件限制——这正是官方宣传的"解锁前所未有的模型规模"。

配合 blogs/deepspeed4science/ 中介绍的科学大模型实践，以及 DeepNVMe（教程见 docs/_tutorials/deepnvme.md）基于 NVMe SSD + libaio + GPUDirect Storage 的高吞吐异步 I/O 加速，NVMe 层不再是性能短板。

💡 想更进一步？DeepSpeed 还推出了ZeRO-PP（流水线-参数并行双切分，详见 blogs/zeropp/）与DeepSpeed-FP6（用 FP6 低精度压缩状态），可继续把单卡可训练的模型规模推高。

五、快速上手三步走

1. 安装：pip install deepspeed（含 CPU Adam、异步 I/O 等算子，可用ds_report检查环境）；
2. 包装模型：调用deepspeed.initialize(model=..., config=配置json)，几行代码接入任意 PyTorch 模型（入门教程见 docs/_tutorials/getting-started.md）；
3. 调配置：按上文 JSON 模板设置stage与offload_optimizer/offload_param（ZeRO-Infinity 需 stage 3 + 参数卸载），跑起来。

六、总结：选型建议一览

核心关键词回顾：DeepSpeed、ZeRO-Offload、ZeRO-Infinity、ZeRO Stage 3、大模型分布式训练、单卡训练、GPU 显存优化、CPU 卸载、NVMe 卸载。

从零冗余思想到三级存储层级，DeepSpeed 用"切分 + 卸载 + 预取"三板斧，把 1000 亿参数大模型的训练门槛从"超算集群"拉低到"一张显卡"。修改几行 JSON 即可上车的易用性，让它成为新手进阶大模型训练的首选工具。🚀

【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed