总目录 › OpenRLHF › 第 7 章
Chapter 07

DeepSpeed 训练后端配置

📌 commit 6c6056daa522 DeepspeedStrategy + ZeRO 与 RLHF 适配

OpenRLHF 的训练后端选了 DeepSpeed，actor / critic / ref / reward 四个角色每个都跑在独立的 DeepspeedStrategy 上。 本章拆解 --ds.* 系列参数，讲清楚 ZeRO 阶段怎么选、offload 怎么用、packing samples 是什么。 对 DeepSpeed 本身不熟可以先翻 DeepSpeed 那本书。

7.1训练后端的工作划分

角色DeepSpeed 配置关键考虑
Actor ZeRO-3 + Adam + bf16 每步训练，要省显存
Critic 同上 同上
Reference ZeRO-3 + 冻结 + bf16，无 optimizer 只 forward，能 offload
Reward 同上（如果是 RM） 同上

OpenRLHF 把这套配置封装在 openrlhf/utils/deepspeed/ 里，对外暴露 --ds.* 参数。

7.2核心 --ds.* 参数

flag说明默认
--ds.zero_stage ZeRO 阶段 (0/1/2/3) 3
--ds.param_dtype 参数精度 (bf16 / fp16 / fp32) bf16
--ds.adam_offload 把 Adam 优化器状态卸到 CPU 关
--ds.zpg ZeRO++ partition group 大小 —
--ds.tensor_parallel_size 开 DS Tensor Parallel 1
--ds.ring_attn_size Ring Attention 度 1
--ds.ring_attn_head_stride Ring Attention head 步长 —
--ds.packing_samples 把短样本拼成长 sequence（必开） 关
--ds.enable_sleep 训练 actor offload 到 CPU（colocate）关
--ds.gradient_accumulation_steps梯度累积步数 1

7.3ZeRO 阶段怎么选

场景推荐原因
7B PPO + 8 卡 ZeRO-2 显存够，stage 越低通信越少
13B PPO + 8 卡 ZeRO-3 必须切参数
70B PPO + 16 卡 ZeRO-3 + adam_offload 显存极紧
70B PPO + 32 卡 ZeRO-3 16 卡足够时不开 offload，速度更好
GRPO（无 critic） 同 PPO 但能开一档低 stage少一个角色显存
colocate_all ZeRO-3 + enable_sleep必须配 sleep 才装得下

7.4关键开关：packing_samples

--ds.packing_samples 是 OpenRLHF非常关键的开关。它做的事：

为什么要 packing

RLHF 训练里每个 batch 的样本长度差异极大（短的 100 token、长的 4K token）。

不 packing：每个 batch 都按最长样本 padding → 计算大量在 padding 上浪费。

packing：把短样本拼接到同一 sequence 里，用 attention mask 隔开 → 利用率拉到 90%+。

实现细节：

大模型 RLHF 几乎必开。

7.5adam_offload：把优化器换到 CPU

当 ZeRO-3 + bf16 都开了仍然显存不够，可以开 --ds.adam_offload：

位置不开 offload开 offload
fp32 master 参数 GPU CPU
Adam 一阶矩 $m$ GPU CPU
Adam 二阶矩 $v$ GPU CPU
bf16 参数 / 梯度 GPU GPU

显存收益（占总参数倍数）：

$$ \Delta M_\text{adam offload} = -12 \Psi / \mathrm{DP}\ \text{(每张卡)} $$

13B 模型 DP=8 时省下 $12 \times 13 / 8 = 19.5$ GB。代价是每步 optimizer.step 慢 20-40%，因为参数要在 CPU 和 GPU 之间搬运。

建议：除非 OOM，不要轻易开 offload。

7.6Ring Attention：长上下文必备

当 sequence length > 8K，attention 的 $O(L^2)$ 内存开始爆炸。OpenRLHF 集成了 Ring Attention：

--ds.ring_attn_size 4 \             # 沿 seq 维切 4 张卡
--ds.ring_attn_head_stride 2

实现：

使用条件：

seq 长度推荐 ring_attn_size
≤ 4K 1（关）
4K - 16K 2
16K - 64K 4
≥ 64K 8

7.7Tensor Parallel：DeepSpeed 自己的 TP

DeepSpeed 0.15+ 加入了自家 TP（不依赖 Megatron）。OpenRLHF 通过 --ds.tensor_parallel_size 2 启用：

--ds.tensor_parallel_size 2

使用场景：

TP × DP × ZeRO 关系：

$$\text{总卡数} = \text{TP} \times \text{DP},\quad \text{ZeRO 切的是 DP 维度}$$

7.8actor.optim 与 critic.optim 分开调

OpenRLHF 默认 actor 和 critic 用不同学习率（PPO 经验值）：

--actor.adam.lr 5e-7 \      # actor：因为基座大、对扰动敏感，lr 必须小
--critic.adam.lr 9e-6 \      # critic：value head 是新初始化的，lr 大一些

每个角色都能单独配 Adam：

flag含义
--actor.adam.lr actor 学习率
--actor.adam.betas Adam beta1 beta2
--actor.adam.eps Adam epsilon
--actor.adam.weight_decay权重衰减
--critic.adam.lr 同上，critic 的

7.9gradient_checkpointing：用时间换显存

--actor.gradient_checkpointing_enable
--critic.gradient_checkpointing_enable

原理：前向不存中间激活，反向时重算。代价是训练时间 +25-35%，但显存能省 50%+。大模型 RLHF 几乎必开。

7.10常见 --ds.* 配置示例

13B PPO 在 8 × A100 80G 上

--ds.zero_stage 3 \
--ds.param_dtype bf16 \
--ds.packing_samples \
--actor.gradient_checkpointing_enable \
--critic.gradient_checkpointing_enable

70B PPO 在 8 × A100 80G 上（极限）

--ds.zero_stage 3 \
--ds.adam_offload \                  # ★ 优化器卸 CPU
--ds.tensor_parallel_size 2 \         # ★ 加 TP
--ds.param_dtype bf16 \
--ds.packing_samples \
--ds.enable_sleep \                    # ★ colocate 时必开
--actor.gradient_checkpointing_enable \
--critic.gradient_checkpointing_enable

32K 长上下文 GRPO 在 8 × H100 上

--ds.zero_stage 3 \
--ds.ring_attn_size 4 \                # ★ 沿 seq 切 4 张卡
--ds.ring_attn_head_stride 2 \
--ds.param_dtype bf16 \
--ds.packing_samples \
--data.max_len 32768 \
--rollout.max_new_tokens 32000 \
--algo.advantage_estimator group_norm \
--rollout.n_samples_per_prompt 8

7.11checkpoint 保存与恢复

flag说明
--ckpt.output_dir 最终 HF 模型导出目录
--ckpt.path DS shard ckpt 中间目录
--ckpt.save_hf 最后导出 HF 格式
--ckpt.save_steps 多少 step 保存一次
--ckpt.max_num 最多保留几份
--ckpt.load_enable 启动时尝试 resume

DS shard 转 HF 用仓库工具：

bash examples/scripts/ckpt_ds_zero_to_universal.sh

7.12DeepspeedStrategy 内部三件事

OpenRLHF 的训练后端不是直接调 DeepSpeed，而是过一层 DeepspeedStrategy（openrlhf/utils/deepspeed/deepspeed.py:37）。它做这三件事：

职责对应方法
初始化分布式 + device mesh setup_distributed(timeout=...) 把 ring_attn / tp 组建好
动态生成 ds_config get_ds_train_config(is_actor) 按 --ds.* flag 拼出完整 JSON
统一的 backward / step API strategy.backward(loss, model, optimizer) + strategy.optimizer_step(...)

"为什么不直接调 DeepSpeed API？" → 训练过程中要支持 ring attention / sleep mode / muon 优化器，Strategy 是一个统一接入点，把这些扩展挂在上面。未来要换成 FSDP / Megatron 也是改 Strategy 子类，业务层无感。

7.133D device mesh (dp, sp, tp)

OpenRLHF 把每个 group 内部当成一个 3D 网格：

维度含义YAML 字段
dp data parallel 数 自动 = world_size / (sp × tp)
sp sequence parallel (Ring Attention) 数 --ds.ring_attn_size
tp tensor parallel 数（DeepSpeed Tensor Parallel） --ds.tensor_parallel_size

例如 8 卡机开 ring_attn_size=2, tensor_parallel_size=2 → mesh 是 (dp=2, sp=2, tp=2)，每张卡同时是 1 个 DP 组里的、1 个 SP 组里的、1 个 TP 组里的。Tensor Parallel 必须 bf16（fp16 数值精度撑不住）；DS Tensor Parallel 是 0.18+ 才支持的相对新功能。

7.14Muon 优化器（2025-Q4 新增）

DeepSpeed 0.18.9+ 集成了 Muon（Newton-Schulz 风格的二阶优化器）。OpenRLHF 通过 --actor.optim_name muon 启用：

优化器显存收敛什么时候用
AdamW（默认） 2× params（m, v） 稳，最广用 常规
Muon 1× params（无一阶矩） 初期慢，后期更稳 大模型长训练，省显存

Muon 特别适合 2D 权重（Attention QKV / MLP），1D 权重（bias / LayerNorm）自动 fallback 到 AdamW。开启时显存约比 AdamW 省 33%。

7.15Strategy 的 sleep 接口

colocate 模式下，Strategy 提供两个对称的接口让 DS 引擎"睡觉"：

方法含义
strategy.offload_states() 把 optimizer state / gradient / params 都搬到 CPU 内存（保留 ZeRO 切片信息）
strategy.reload_states() 从 CPU 搬回 GPU，继续训练

YAML 字段 --ds.enable_sleep。每个训练 step 末尾 controller 调 offload，rollout 前调 reload。配合 vLLM 的 sleep_level=2，actor 训练显存和 vLLM 推理显存不重叠 —— 这就是 Hybrid Engine 真正省显存的地方。

7.16这章你需要带走的