文章目录
- 说明
- 一 DeepSpeed
- 二 DeepSpeed 的核心技术
-
- 2.1 ZeRO（零冗余优化器）显存优化技术
-
- 2.1.1 ZeRO-1: 优化器状态分片 (Optimizer State Partitioning)
- 2.1.2 ZeRO-2: 梯度分片 (Gradient Partitioning)
- 2.1.3 ZeRO-3: 参数分片 (Parameter Partitioning)
- 2.1.4 ZeRO-Offload
- 2.2 3D 并行策略
- 2.3 混合精度训练
- 2.4 DeepSpeed-Chat: 端到端RLHF训练
- 三 DeepSpeed使用和参数配置
-
- 3.1 DeepSpeed简单使用
- 3.2 DeepSeed详细参数配置
- 四 总结
说明
- 文中部分资料整理参考自互联网，仅供学习和交流使用。
一 DeepSpeed
- DeepSpeed是微软开发的深度学习优化库，专门解决大模型训练中的显存和计算问题。
- DeepSpeed 的根本目的是解决训练超大规模模型时面临的显存瓶颈和计算效率问题，其作用主要体现在以下几个方面：
| 🚀 极致显存优化 | 单卡显存无法容纳庞大的模型参数、优化器状态、梯度等 | 能在单卡或少量卡上训练比原来大得多的模型（例如，单卡V100用ZeRO-Offload可训练130亿参数模型） |
| ⚡ 大幅提升训练速度 | 传统数据并行在扩展时通信开销大、GPU利用率低 | 通过通信优化、混合精度等技术，实现近乎线性的加速比，训练速度可提升数倍至十数倍 |
| 🌐 支持超大规模分布式训练 | 训练万亿参数模型需要数千张卡协同工作，复杂度高 | 提供3D并行（数据并行+模型并行+流水线并行）策略，支持跨千卡集群的高效扩展 |
| 🛠️ 降低训练门槛与成本 | 昂贵的计算集群和复杂的分布式配置限制了大众参与 | 让研究人员和中小企业能用消费级GPU（如RTX 3090/4090）或更少的云资源进行大模型训练，大幅降低经济成本 |
DeepSpeed 的强大功能使其在大模型领域得到了广泛应用：
- 🔬 超大规模模型训练：训练参数量达到千亿、万亿级别的模型（如Megatron-Turing NLG 530B参数），这是单卡或简单并行无法实现的。
- 🏢 企业级与学术研究：帮助企业和研究机构用更少的GPU、更低的成本训练出性能更好的模型，加速大模型技术的普及和创新。
- 🎯 高效微调与适配：利用ZeRO-Offload和LoRA等技术，可以在消费级显卡上高效地对大模型进行微调，适配特定任务或领域，极大降低了个人开发者的门槛。
- 💬 开发ChatGPT类应用：DeepSpeed-Chat 让训练和部署类似ChatGPT的对话模型变得简单、快速且经济，加速了通用人工智能助手的发展。
二 DeepSpeed 的核心技术
- DeepSpeed 的强大功能源于其一系列创新技术，其中最核心的是 ZeRO (Zero Redundancy Optimizer) 优化技术和多种并行策略。
2.1 ZeRO（零冗余优化器）显存优化技术
-
在传统的分布式数据并行中，每个GPU都会维护一份完整的模型参数、梯度和优化器状态，这造成了巨大的显存冗余。
-
ZeRO通过切分（Sharding） 这些状态，在多个GPU之间进行存储和计算，从而消除冗余，极大节省显存。
ZeRO有三个逐步深入的优化阶段，逐步节省显存：
ZeRO–1: 只分割优化器状态
├── 模型参数：每个GPU都有完整副本
├── 梯度：每个GPU都有完整副本
└── 优化器状态：分割到不同GPU ✂️
ZeRO–2: 分割优化器状态 + 梯度
├── 模型参数：每个GPU都有完整副本
├── 梯度：分割到不同GPU ✂️
└── 优化器状态：分割到不同GPU ✂️
ZeRO–3: 分割所有内容
├── 模型参数：分割到不同GPU ✂️
├── 梯度：分割到不同GPU ✂️
└── 优化器状态：分割到不同GPU ✂️
- ZeRO-1：只切分优化器状态（如Adam的动量、方差）。4x 显存节省。
- ZeRO-2：在ZeRO-1基础上，再切分梯度。8x 显存节省。
- ZeRO-3：在ZeRO-2基础上，连模型参数也进行切分。这是最极致的显存优化，Nx 显存节省（N为GPU数量），允许训练远超单卡显存容量的模型。
- ZeRO-Offload是ZeRO的重要扩展，它允许将优化器状态和梯度卸载到CPU内存甚至NVMe硬盘上，从而进一步解放GPU显存，让消费级显卡也能尝试训练更大的模型。
DeepSpeed与传统PyTorch分布式数据并行（DDP）的区别：
| 显存占用 | 高冗余：每卡存储完整的模型参数、梯度、优化器状态 | 低冗余：通过ZeRO技术切分这些状态，显存占用可减少4x至Nx倍 |
| 可支持的模型规模 | 受限于单卡显存 | 远超单卡显存，能训练超大规模模型 |
| 通信效率 | 通信随GPU数量增加，可能成为瓶颈 | 通过梯度分片、通信重叠等优化，显著减少通信量，提升扩展效率 |
| 资源利用 | 多卡训练时，部分GPU显存可能闲置 | 极致利用所有GPU显存，支持CPU卸载，榨干硬件性能 |
| 使用复杂度 | 相对简单，集成度好 | 配置选项更多，但提供了丰富的模板和示例，官方文档和社区支持良好 |
2.1.1 ZeRO-1: 优化器状态分片 (Optimizer State Partitioning)
# 4个GPU的情况下，ZeRO-1的显存重新分配
传统方式 (每个GPU重复存储):
GPU 0: [完整模型14GB] + [完整梯度14GB] + [完整优化器56GB] = 84GB ❌
GPU 1: [完整模型14GB] + [完整梯度14GB] + [完整优化器56GB] = 84GB ❌
GPU 2: [完整模型14GB] + [完整梯度14GB] + [完整优化器56GB] = 84GB ❌
GPU 3: [完整模型14GB] + [完整梯度14GB] + [完整优化器56GB] = 84GB ❌
ZeRO–1 (优化器状态分片):
GPU 0: [完整模型14GB] + [完整梯度14GB] + [优化器1/4: 14GB] = 42GB ✅
GPU 1: [完整模型14GB] + [完整梯度14GB] + [优化器2/4: 14GB] = 42GB ✅
GPU 2: [完整模型14GB] + [完整梯度14GB] + [优化器3/4: 14GB] = 42GB ✅
GPU 3: [完整模型14GB] + [完整梯度14GB] + [优化器4/4: 14GB] = 42GB ✅
关键机制：
通信模式：
- 优势：通信开销最小、训练速度最快、实现简单，稳定性最好。
- 劣势：显存节省有限 (只节省优化器显存) ，仍需要大量GPU显存存储完整模型和梯度。
- 适用场景：GPU显存充足 (40GB+) ，追求最高训练速度，模型不是特别大 (7B-13B)。
2.1.2 ZeRO-2: 梯度分片 (Gradient Partitioning)
# ZeRO-2进一步分片梯度
GPU 0: [完整模型14GB] + [梯度1/4: 3.5GB] + [优化器1/4: 14GB] = 31.5GB ✅
GPU 1: [完整模型14GB] + [梯度2/4: 3.5GB] + [优化器2/4: 14GB] = 31.5GB ✅
GPU 2: [完整模型14GB] + [梯度3/4: 3.5GB] + [优化器3/4: 14GB] = 31.5GB ✅
GPU 3: [完整模型14GB] + [梯度4/4: 3.5GB] + [优化器4/4: 14GB] = 31.5GB ✅
通信机制：
通信模式
优势：显存节省显著 (节省75%的梯度显存) ，通信开销适中，性能和显存的良好平衡。 劣势：反向传播时需要额外通信，仍需要完整模型参数
适用场景：GPU显存中等 (24-40GB) ，平衡性能和显存使用，中等大小模型 (13B-30B)。
2.1.3 ZeRO-3: 参数分片 (Parameter Partitioning)
- ZeRO-3是最激进的分片策略。
GPU 0: [模型1/4: 3.5GB] + [梯度1/4: 3.5GB] + [优化器1/4: 14GB] = 21GB ✅
GPU 1: [模型2/4: 3.5GB] + [梯度2/4: 3.5GB] + [优化器2/4: 14GB] = 21GB ✅
GPU 2: [模型3/4: 3.5GB] + [梯度3/4: 3.5GB] + [优化器3/4: 14GB] = 21GB ✅
GPU 3: [模型4/4: 3.5GB] + [梯度4/4: 3.5GB] + [优化器4/4: 14GB] = 21GB ✅
通信模式 (最复杂)
# ZeRO-3前向传播的详细通信过程
时间轴: ───→
GPU 0: │参数1/4│ All–Gather │完整参数│ 计算 │释放│ All–Gather │...
GPU 1: │参数2/4│ ↓ │完整参数│ 计算 │释放│ ↓ │...
GPU 2: │参数3/4│ 收集参数 │完整参数│ 计算 │释放│ 收集参数 │...
GPU 3: │参数4/4│ ↑ │完整参数│ 计算 │释放│ ↑ │...
步骤详解:
优势：显存节省最大 (可训练超大模型) ，支持无限大的模型，可以在小显存GPU上训练大模型。 劣势：通信开销最大 (每层都需要通信) ，训练速度最慢，实现复杂，可能不稳定。 适用场景：GPU显存有限 (16-24GB) ，超大模型 (30B+) ，愿意牺牲速度换取显存。
显存节省效果对比
| 传统 | 14GB | 14GB | 56GB | 84GB | 0% |
| ZeRO-1 | 14GB | 14GB | 14GB | 42GB | 50% |
| ZeRO-2 | 14GB | 3.5GB | 14GB | 31.5GB | 62.5% |
| ZeRO-3 | 3.5GB | 3.5GB | 14GB | 21GB | 75% |
2.1.4 ZeRO-Offload
- ZeRO-Offload允许将优化器状态和梯度卸载到CPU内存甚至NVMe硬盘上，从而进一步解放GPU显存，让消费级显卡也能尝试训练更大的模型。
不用Offload：所有东西都放在GPU显存里
┌─────────────────┐
│ GPU 显存 │ ← 模型+梯度+优化器都在这里
│ (40GB) │
└─────────────────┘
用CPU Offload：把一些东西放到CPU内存里
┌─────────────────┐ ┌─────────────────┐
│ GPU 显存 │ │ CPU 内存 │
│ (40GB) │ │ (128GB) │
│ 正在计算的数据 │ │ 暂时不用的数据 │
└─────────────────┘ └─────────────────┘
NVMe Offload:
┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐
│ GPU 显存 │ │ CPU 内存 │ │ NVMe 存储 │
│ 当前计算数据 │ │ 缓冲数据 │ │ 冷数据 │
└─────────────────┘ └─────────────────┘ └─────────────────┘
Offload 可以选择卸载不同的组件：
DeepSpeed Offload 卸载策略对比
| CPU 内存 | "cpu" | 中等(比GPU慢10-20倍) | 大(通常64GB-512GB) | 中等 |
| NVMe SSD | "nvme" | 慢(比GPU慢100-200倍) | 很大(1TB-8TB) | 低 |
| 普通硬盘 | "disk" | 很慢(比GPU慢1000倍+) | 最大(10TB+) | 最低 |
不同 Offload 方式对比
| 无 Offload（纯 GPU） | 速度最快 | 显存占用最多 | 显存充足的情况 |
| CPU Offload | 节省显存，速度还可以 | 比纯 GPU 慢一些 | 显存不够但 CPU 内存充足 |
| NVMe Offload（硬盘） | 节省最多显存 | 速度最慢 | 显存和 CPU 内存都不够的极端情况 |
2.2 3D 并行策略
- 当模型参数量巨大（如千亿、万亿级），单卡甚至单机都无法容纳时，DeepSpeed 提供 3D 并行策略，将模型和数据在三个维度上进行切分：
| 数据并行 (DP) | 复制模型到每个GPU，拆分数据 | 加速计算，处理大规模数据集 | 模型能放进单卡，追求训练速度 | All-Reduce：通信频率较低，但随GPU数增加，通信量可能成为瓶颈 |
| 模型并行 (MP)（包括张量并行TP） | 层内切分：将同一层内的矩阵运算拆分到多GPU | 解决模型单层参数过大，无法存入单卡显存 | 模型宽度大（隐藏层维度大），计算密集 | All-Reduce/All-Gather：通信频繁，数据量大，依赖节点内高速互联（如NVLink） |
| 流水线并行 (PP) | 层间切分：将模型的不同层分配到不同GPU，流水线式处理 | 解决模型层数过多，单卡无法容纳 | 模型很深（层数多），每个stage计算耗时相近 | P2P通信：相邻阶段间传递激活值，通信频率较高但单次数据量小 |
张量并行和流水线并行是模型并行的两种核心形式，它们切分模型的方式和解决的问题完全不同：
- 张量并行（TP）：水平切分，就像切蛋糕。把一个大的矩阵乘法拆分成多个小的矩阵乘法，在不同的GPU上并行计算，最后合并结果。它主要针对模型单层参数量过大的情况（如Transformer的MLP层）。通信非常频繁（几乎每个操作后都要同步），因此极度依赖GPU之间的高速互联（如NVIDIA NVLink），通常只适用于单机内的多卡并行。
- 流水线并行（PP）：纵向切分，就像工厂的流水线。把模型的不同层分配到不同的GPU上。数据像流水一样经过每个GPU进行计算，每个GPU只负责模型的一部分。它主要解决模型总层数过多，单卡显存无法容纳整个模型的问题。通信发生在相邻的“工位”（GPU）之间，传递中间结果（激活值）。其挑战在于如何合理安排“微批次”（micro-batches）来减少“气泡”（GPU空闲等待时间），提高利用率。
- 3D并行就是根据模型的特点和硬件条件，灵活组合这三种策略，例如在单机内用TP+DP，跨节点间用PP，从而实现极致的扩展性。
2.3 混合精度训练
DeepSpeed 支持 FP16 和 BF16（Bfloat16）混合精度训练。这意味着：
- 计算：使用低精度格式（FP16/BF16）进行矩阵乘法等计算，速度更快。
- 存储：模型参数和梯度用低精度存储，节省显存。
- 精度：对于某些关键步骤（如损失值缩放、优化器状态更新），仍保持较高精度（FP32）以保证数值稳定性和模型精度。
- 通过动态损失缩放（Loss Scaling） 技术，DeepSpeed还能有效避免FP16计算中可能出现的数值下溢问题。
2.4 DeepSpeed-Chat: 端到端RLHF训练
- DeepSpeed 还提供了 DeepSpeed-Chat 模块，这是一个一键式训练类似ChatGPT模型的系统框架。它完美复现了InstructGPT论文中的RLHF（基于人类反馈的强化学习） 训练流程，包括：
- 监督微调（SFT）
- 奖励模型（RM）微调
- 基于人类反馈的强化学习（RLHF）
它通过 DeepSpeed Hybrid Engine 将高效训练和推理能力结合，使得训练ChatGPT类模型的速度比现有系统快15倍，成本大幅降低。
三 DeepSpeed使用和参数配置
3.1 DeepSpeed简单使用
- 使用 DeepSpeed 通常非常简单，它提供了与 PyTorch 和 Hugging Face Transformers 无缝集成的接口。基本流程如下：
"train_batch_size": 32,
"gradient_accumulation_steps": 4,
"optimizer": {
"type": "AdamW",
"params": {
"lr": 5e-5
}
},
"fp16": {
"enabled": true
},
"zero_optimization": {
"stage": 2,
"offload_optimizer": {
"device": "cpu"
},
"allgather_partitions": true,
"overlap_comm": true
}
}
# … 定义的模型、optimizer、dataloader等 …
model_engine, optimizer, _, _ = deepspeed.initialize(
model=model,
model_parameters=model.parameters(),
config="ds_config.json"
)
# 训练循环中，用 model_engine 替代原来的 model
for batch in dataloader:
loss = model_engine(batch)
model_engine.backward(loss)
model_engine.step()
deepspeed –num_gpus=4 your_training_script.py –deepspeed ds_config.json
# 多节点多卡（需要创建hostfile等）
deepspeed –num_gpus=8 –num_nodes=2 –hostfile=hostfile your_training_script.py –deepspeed ds_config.json
💡 小贴士：DeepSpeed 官方提供了丰富的示例和文档，Hugging Face Transformers 也提供了对DeepSpeed的原生支持，只需在 TrainingArguments 中指定 deepspeed="ds_config.json" 即可，非常方便。
3.2 DeepSeed详细参数配置
| 基础训练参数 | train_micro_batch_size_per_gpu | 设置每个GPU上的微批次大小 | 影响：线性影响GPU显存使用。调优：更大通常训练更稳定，但受限于显存。避免与HF的per_device_train_batch_size冲突。 |
| train_batch_size | 全局批次大小（所有GPU总和） | 计算：= micro_batch_size × gradient_accumulation_steps × num_gpus。影响：影响学习率调度和收敛行为。 | |
| gradient_accumulation_steps | 梯度累积步数，模拟大batch size | 优势：不增加显存使用。劣势：增加计算时间。作用：提高训练稳定性。 | |
| ZeRO 优化参数 | stage | 选择ZeRO优化级别 (1, 2, 3) | 权衡：Stage越高（如3），显存节省越多，但通信开销越大。需在内存效率与训练速度间平衡。 |
| overlap_comm | 通信与计算重叠 | 作用：减少通信等待时间。效果：提升训练速度（约5-15%）。要求：需支持异步通信的硬件。 | |
| contiguous_gradients | 将梯度存储在连续内存中 | 作用：减少内存碎片，提高内存访问效率，降低内存分配开销。 | |
| ZeRO-3 专有参数 | stage3_max_live_parameters | 控制同时在GPU上的最大参数数量 | 调优：显存紧张时应减少此值，强制更频繁地释放参数。 |
| stage3_max_reuse_distance | 控制参数重用的最大距离 | 调优：减少此值可节省显存（参数更快被换出），但会增加通信频率。 | |
| stage3_prefetch_bucket_size | 控制参数预取的桶大小 | 权衡：更大的桶减少传输次数但占用更多内存；小的桶节省内存但可能传输效率低。 | |
| Offload 参数 | device (offload_optimizer/param) | 指定卸载目标 (cpu 或 nvme) | cpu：速度中等，容量大。nvme：速度较慢，容量极大。 |
| pin_memory | 使用页锁定内存加速传输 | 效果：传输速度提高2-3倍，但占用更多系统内存。建议：系统内存充足时开启。 | |
| buffer_count | 并行传输的缓冲区数量 | 建议：4-8个通常是最佳选择，更多缓冲区占用更多内存但提高并行度。 | |
| buffer_size | 每个缓冲区的大小（字节） | 建议：1e8 (100MB) 通常是好的起点，影响传输粒度和内存占用。 | |
| max_in_cpu | CPU中允许的最大数据量 | 作用：限制CPU内存使用，超出时数据进一步卸载到磁盘。根据系统内存大小设置。 | |
| offload_param | 卸载模型参数 (仅限 ZeRO-3) | 特点：最激进的内存节省方式，但速度最慢。 | |
| 通信优化参数 | allgather_bucket_size | AllGather操作的桶大小 | 作用：影响参数收集效率和网络带宽利用率。网络带宽高时可增大。 |
| reduce_bucket_size | Reduce操作的桶大小 | 作用：影响梯度同步效率和聚合性能。通常建议与allgather_bucket_size保持一致。 | |
| 混合精度参数 | enabled (fp16) | 启用FP16混合精度训练 | 效果：减少50%显存，提高训练速度。注意：可能出现数值不稳定，现代GPU建议开启。 |
| initial_scale_power | 初始损失缩放因子 (2^power) | 作用：防止梯度下溢。建议：12-16通常是好的选择。 | |
| 激活检查点参数 | partition_activations | 在多GPU间分割激活值 | 效果：减少每个GPU的激活值内存，但会增加通信开销。适合显存紧张时。 |
| cpu_checkpointing | 将激活值检查点存储在CPU | 权衡：大幅减少GPU显存（增加CPU内存使用），但增加CPU-GPU传输开销（牺牲速度）。 |
显存不足时的调优顺序，按效果递增的顺序：
速度优化的调优顺序，按效果递增的顺序：
四 总结
- 总而言之，DeepSpeed 就像一个深度学习训练的超级加速引擎和资源管家。它通过ZeRO等显存优化技术，打破了单卡显存的桎梏；通过3D并行策略，让训练超大规模模型成为可能；通过混合精度和通信优化，榨干了硬件的每一分性能；并通过DeepSpeed-Chat等模块，简化了复杂训练流程。
- 无论是想训练万亿参数模型的顶尖研究机构，还是想在个人显卡上体验大模型微调的爱好者，DeepSpeed 都提供了强大且必要的工具支持，是现代深度学习，特别是大模型领域不可或缺的基石之一。
