
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 训练 100B1000亿参数的 GPT-2 大模型。文章面向新手图解原理 配置实操 进阶路线一文掌握。一、为什么大模型训练会爆显存训练一个 1000 亿参数的大模型显存里装的不只是权重。以 FP16 训练 Adam 优化器为例每个参数背后还藏着 32 位权重副本、动量和二阶矩等状态状态单参数占用100B 模型总量模型参数FP162 字节~200 GB梯度FP162 字节~200 GB优化器状态FP32 权重 m/v12 字节~1200 GB合计~16 字节~1.6 TB一张 32GB 的 V100 显然装不下。传统方案是模型并行——但需要大改模型代码复杂且通信开销高。而 DeepSpeed 的ZeROZero Redundancy Optimizer零冗余优化器思路完全不同数据并行的每张卡上都有一份冗余的优化状态那就把它们切分到所有设备上各存各的、各算各的用时再聚合。这样不仅省显存还不需要修改任何模型代码——只需在 DeepSpeed 的 JSON 配置中加几行参数即可这是 ZeRO 相比模型并行的最大吸引力。二、ZeRO 三级进阶从 Stage 1 到 Stage 3ZeRO 以渐进阶段stage实现优化后期阶段包含前期全部优化详见官方教程 docs/_tutorials/zero.mdStage 1 — 切分优化器状态每张卡只保存并更新自己那一份优化器状态。官方实测1.5B 参数 GPT-2 的 Adam 状态从 18GB 降到 2.25GB8 卡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 版 AdamDeepSpeedCPUAdam源码见 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-InfinityGPU → 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 低精度压缩状态可继续把单卡可训练的模型规模推高。五、快速上手三步走安装pip install deepspeed含 CPU Adam、异步 I/O 等算子可用ds_report检查环境包装模型调用deepspeed.initialize(model..., config配置json)几行代码接入任意 PyTorch 模型入门教程见 docs/_tutorials/getting-started.md调配置按上文 JSON 模板设置stage与offload_optimizer/offload_paramZeRO-Infinity 需 stage 3 参数卸载跑起来。六、总结选型建议一览场景推荐方案单卡可训规模参考多卡、显存中等紧张ZeRO Stage 1/210B单卡或 CPU 内存充裕ZeRO-Offload本教程核心10B~数十 B超大规模、显存内存都吃紧ZeRO-InfinityStage 3 CPU/NVMe 卸载100B~1T核心关键词回顾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创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考