分布式训练（二): DeepSpeed 知识总结

目录

• 一、DeepSpeed 简介
  • 1.1 核心特性
  • 1.2 典型使用场景
  • 1.3 如果已经有 PP、TP，为什么还需要 DeepSpeed
    • 1.3.1 PP / TP 解决的与未解决的
    • 1.3.2 DeepSpeed 在其中的作用
    • 1.3.3 关键结论
  • 1.4 混合精度训练详解
    • 1.4.1 不同张量的精度分工
    • 1.4.2 训练流程示意
• 二、ZeRO 原理与阶段对比
  • 2.1 ZeRO 定义
  • 2.2 ZeRO-1：分片优化器状态
  • 2.3 ZeRO-2：分片优化器状态 + 梯度
  • 2.4 ZeRO-3：分片优化器状态 + 梯度 + 参数
  • 2.5 显存对比（以 7.5B 模型、64 卡为例）
  • 2.6 扩展能力
  • 2.7 三个阶段优缺点与使用场景对比
  • 2.8 选择建议
• 三、DeepSpeed 配置示例
  • 3.1 ZeRO-1 配置
  • 3.2 ZeRO-2 配置（推荐中等规模模型）
  • 3.3 ZeRO-3 配置（大模型）
  • 3.4 配置参数说明
  • 3.5 启动命令

一、DeepSpeed 简介

DeepSpeed 是微软开源的深度学习优化库，主要用于大规模模型训练。

1.1 核心特性

• ZeRO (Zero Redundancy Optimizer)：将优化器状态、梯度、参数分片到多个 GPU，大幅降低单卡显存占用，可训练千亿/万亿参数模型
• 3D 并行：原生支持数据并行 (DP) 与流水线并行 (PP)；张量并行 (TP) 主要通过与 Megatron-LM 集成（Megatron-DeepSpeed）实现，从而组合出完整的 3D 并行
• 混合精度训练(1.4 详细介绍)：训练中同时使用低精度 (FP16/BF16) 和 FP32——前向/反向计算、激活和梯度用低精度以省显存、加速计算，而 master 权重和优化器状态保留 FP32 以保证参数更新的数值精度
• CPU/NVMe Offload：将参数和优化器状态卸载到 CPU 内存或 NVMe，进一步突破显存限制
• DeepSpeed-MoE：稀疏专家模型支持

1.2 典型使用场景

• 大模型训练：训练 LLM、视觉大模型、MoE 模型等（如 BLOOM、MT-NLG）
• 显存受限训练：单卡/多卡显存不足时，借助 ZeRO 与 Offload 训练更大模型
• 框架集成：可与 HuggingFace Transformers、PyTorch Lightning、Accelerate 等无缝配合

1.3 如果已经有 PP、TP，为什么还需要 DeepSpeed

PP（流水线并行）和 TP（张量并行）只是「怎么切模型」的两种方法，而 DeepSpeed 是「怎么训模型」的完整训练引擎。两者层次不同，不是替代关系。

1.3.1 PP / TP 解决的与未解决的

• PP / TP 解决：模型切分（按层切 / 按矩阵切），让超大模型能放进多卡
• PP / TP 未解决：DP 维度冗余、显存进一步压缩（Offload）、混合精度、训练工程细节、通信优化、分布式 checkpoint、并行调度策略等

1.3.2 DeepSpeed 在其中的作用

1.3.3 关键结论

• PP / TP 是切分策略，DeepSpeed 是训练框架
• 真实训练几乎都是TP + PP + DP 组合，DP 维度需要 ZeRO 去冗余
• DeepSpeed 让用户只写 config，不必手撸分布式细节

1.4 混合精度训练详解

"混合精度"指的是训练中同时使用 FP16/BF16 和 FP32：低精度负责计算，FP32 负责关键的累加和参数更新。

1.4.1 不同张量的精度分工

1.4.2 训练流程示意

FP32 master weights │ 转 FP16/BF16 ▼ 低精度权重 ─► forward ─► loss ─► backward ─► 低精度梯度 │ 转 FP32 ▼ 优化器更新 (FP32) │ ▼ FP32 master weights（更新）

补充：为什么必须保留 FP32 主权重？FP16 只有约 3-4 位有效数字，若直接累加小的更新量（如W=1.0、lr*grad=0.0001）会被精度"吞掉"导致更新丢失，所以必须在 FP32 上做累加，再把更新后的权重转回低精度用于下一轮 forward。这里的"FP32 主权重"是训练时的内存策略，与模型保存到磁盘的格式无关——训练 checkpoint 通常存 FP32 便于恢复训练，而对外发布的权重（如 LLaMA、Qwen）大多是 BF16。

二、ZeRO 原理与阶段对比

2.1 ZeRO 定义

ZeRO全称Zero Redundancy Optimizer（零冗余优化器），是 DeepSpeed 中的核心显存优化技术，由微软在论文《ZeRO: Memory Optimizations Toward Training Trillion Parameter Models》中提出。

在传统的数据并行（Data Parallelism, DP）训练中，每张 GPU 都会保存一份完整的模型参数、梯度和优化器状态，存在大量冗余。ZeRO 的核心思想是：将这些训练状态在数据并行的 GPU 之间进行分片（partition），每张卡只保存自己负责的那一部分，从而消除冗余、大幅降低单卡显存占用，使得在有限显存下能训练更大的模型。

训练状态分为三类：

• 优化器状态 (Optimizer States)：如 Adam 的 momentum、variance，通常占用最多显存
• 梯度 (Gradients)：反向传播得到的梯度
• 参数 (Parameters)：模型权重本身

根据分片的范围不同，ZeRO 划分为三个递进的阶段：ZeRO-1、ZeRO-2、ZeRO-3。

2.2 ZeRO-1：分片优化器状态

• 分片内容：仅优化器状态（如 Adam 的 momentum、variance）
• 显存节省：约 4 倍（相比传统 DP）
• 通信开销：与传统 DP 基本相同
• 原理：每张卡只保存 1/N 的优化器状态，参数更新后通过 all-gather 同步

2.3 ZeRO-2：分片优化器状态 + 梯度

• 分片内容：优化器状态 + 梯度
• 显存节省：约 8 倍
• 通信开销：与传统 DP 基本相同（reduce-scatter 替代 all-reduce）
• 原理：每张卡只保留自己负责更新的那部分参数的梯度

2.4 ZeRO-3：分片优化器状态 + 梯度 + 参数

• 分片内容：全部三类状态
• 显存节省：与 GPU 数量 N 成线性关系（理论上 N 倍）
• 通信开销：约为传统 DP 的 1.5 倍
• 原理：参数也分片存储，前向/反向计算时按需 all-gather，用完即释放

2.5 显存对比（以 7.5B 模型、64 卡为例）

2.6 扩展能力

• ZeRO-Offload：在 ZeRO-2 基础上把优化器状态和梯度卸载到 CPU
• ZeRO-Infinity：在 ZeRO-3 基础上把参数卸载到 CPU/NVMe，可训练万亿参数模型

2.7 三个阶段优缺点与使用场景对比

2.8 选择建议

• 显存够用→ ZeRO-1 或 ZeRO-2（通信少，速度快）
• 显存不够→ ZeRO-3
• 极端显存受限→ ZeRO-3 + Offload（CPU/NVMe）
• 建议优先级：ZeRO-2 → ZeRO-3 → ZeRO-3 + Offload（按需逐级升级，避免过度牺牲速度）

三、DeepSpeed 配置示例

3.1 ZeRO-1 配置

{"train_batch_size":32,"optimizer":{"type":"AdamW","params":{"lr":1e-4,"weight_decay":0.01}},"bf16":{"enabled":true},"zero_optimization":{"stage":1}}

3.2 ZeRO-2 配置（推荐中等规模模型）

{"train_batch_size":32,"optimizer":{"type":"AdamW","params":{"lr":1e-4,"weight_decay":0.01}},"bf16":{"enabled":true},"zero_optimization":{"stage":2,"overlap_comm":true,"contiguous_gradients":true}}

3.3 ZeRO-3 配置（大模型）

{"train_batch_size":32,"optimizer":{"type":"AdamW","params":{"lr":1e-4,"weight_decay":0.01}},"bf16":{"enabled":true},"zero_optimization":{"stage":3,"overlap_comm":true,"contiguous_gradients":true,"stage3_gather_16bit_weights_on_model_save":true}}

3.4 配置参数说明

顶层参数

zero_optimization参数

3.5 启动命令

deepspeed--num_gpus=8train.py--deepspeed--deepspeed_configds_config.json