Hybrid 混合训练模式
概述
Hybrid 模式 是 Relax 在 Colocate(同步) 与 Fully Async(全异步) 之间的第三种执行模式。它将以下两者结合:
- Fully Async 的 流式数据流水线(TransferQueue +
max-staleness控制 off-policy 容忍度),以及 - Colocate 的 进程内权重共享(TensorBackuper +
_switch_model,使 ref / actor_fwd / advantages 全部在 actor 自身 GPU 上完成)。
具体来说,Actor 与 Rollout 仍然部署在 独立的 GPU placement group 上(与 Fully Async 一致),但 actor 不再把权重广播到独立的 ActorFwd / Reference / Advantages 服务,而是通过 CPU/GPU TensorBackuper 在 actor、ref、old_actor、teacher 等 tag 之间切换同一套权重,所有 forward 在本地完成,最后通过同步的 UpdateWeightFromTensor 路径把权重推给 rollout。
模式对比
| 维度 | Colocate(同步) | Fully Async(全异步) | Hybrid |
|---|---|---|---|
| GPU 布局 | Actor 与 Rollout 分时复用同一组 GPU | Actor / Rollout / ActorFwd / Reference 各自独立 GPU | Actor 与 Rollout 独立 GPU;ref / actor_fwd / adv 复用 actor 的 GPU |
| 数据流水线 | TransferQueue,批同步 | TransferQueue + StreamingDataLoader,完全流式 | TransferQueue + 子批次流式(num-iters-per-train-update) |
| 权重同步 | 进程内 tensor 拷贝 | 通过 DCS(Checkpoint Engine)做 NCCL broadcast | 同步的 UpdateWeightFromTensor 推给 rollout;ref/actor_fwd 走 TensorBackuper |
| Staleness | max_staleness = 0(严格 on-policy) | 可配置 max_staleness | 可配置 max_staleness |
| 部署的角色 | actor, critic, rollout | actor, critic, rollout, advantages, reference, actor_fwd | actor, critic, rollout(与 Colocate 相同;ref/actor_fwd 在 actor 内部) |
--balance-data | 支持 | 不支持 | 支持(Hybrid 存在的核心原因之一) |
何时选择 Hybrid
选择 Hybrid 的场景:
- 希望获得独立 rollout GPU 与流水线数据带来的吞吐收益,但
- 模型较大,单独部署 ref / actor_fwd 服务会浪费 GPU,或者
- 需要
--balance-data(DP 间均衡 micro-batch 切分),而纯 Fully Async 不支持此功能。
选择 Fully Async:拥有充足的 GPU 单独运行 ref / actor_fwd / advantages 服务,并希望在 step 之间做真正的并行流水线。
选择 Colocate:GPU 紧张,可以接受 rollout → train 的串行执行。
架构
角色布局
Hybrid 与 Colocate 使用相同的角色集合 —— 只部署 actor、critic(可选)、rollout 三个 Ray Serve 服务。判断逻辑位于 relax/core/registry.py:
def process_role(config):
if config.hybrid:
# hybrid mode: actor handles ref/actor_fwd internally
# via _switch_model, only need actor + rollout services
return ROLES_COLOCATE
if config.fully_async:
...但与 Colocate 不同,actor 与 rollout 的 placement group 是 互相独立 的,这与 Fully Async 一致。参见 relax/core/controller.py:
if colocate and not self.config.hybrid:
# Sync colocate: actor and rollout share GPUs via time-sharing (offload/onload)
actor_rollout_pgs = create_placement_group(num_gpus=num_gpus)
else:
# fully_async (pure or hybrid): actor and rollout use separate GPUs
actor_rollout_pgs = None参数解析
--hybrid 是唯一对外暴露的开关。relax/utils/arguments.py 将其展开为下游已识别的两个底层参数:
if args.hybrid:
args.fully_async = True
args.colocate = True直接传入 --fully-async --colocate 会被拒绝,必须使用 --hybrid。单一开关的设计使 args.hybrid 成为 registry、controller 分发以及 train_hybrid 调用点中识别 hybrid 模式的唯一权威标志。
架构图
ASCII 图中保留英文以避免框线对齐错乱:
┌────────────────────────────────────────────────────────────────────────────┐
│ Controller (Orchestrator) │
│ relax/core/controller.py │
│ │
│ ┌───────────────────────────────────┐ ┌────────────────────┐ │
│ │ Actor Service │ │ Rollout Service │ │
│ │ (own placement group, N GPUs) │ │ (own PG, M GPUs) │ │
│ │ │ │ SGLang engines │ │
│ │ ┌────────────────────────────┐ │ └─────────┬──────────┘ │
│ │ │ TensorBackuper │ │ ▲ │
│ │ │ tags: actor / ref / │ │ │ │
│ │ │ old_actor / teacher │ │ │ │
│ │ │ _switch_model(tag) swaps │ │ │ │
│ │ │ weights on the same GPUs │ │ │ │
│ │ └────────────────────────────┘ │ │ │
│ │ train_hybrid(): │ │ │
│ │ ├─ ref forward (switch:ref) │ │ │
│ │ ├─ actor forward (switch:actor)│ │
│ │ ├─ advantages (in-process) │ │ │
│ │ └─ train (switch:actor)│ │
│ └──────────────┬────────────────────┘ │ │
│ │ UpdateWeightFromTensor (sync) ───┘ │
└──────────────────────┼─────────────────────────────────────────────────────┘
▼
┌───────────────────────────────────────────────────────────────────────────┐
│ TransferQueue (Data Plane) │
│ Rollout writes train_N partition incrementally ──► Actor consumes │
│ in sub-batches via get_meta(batch_size, batch_index) with max-staleness │
└───────────────────────────────────────────────────────────────────────────┘train_hybrid 训练流程
relax/backends/megatron/actor.py:708 实现的 Hybrid 训练步骤分为三个阶段:
采集子批次并完成小批次 forward 计算(峰值显存小)
全局 batch 被切分为
num_iters_per_train_update份子批次。对每个子批次,actor 会:- 从 TransferQueue 拉取数据(
_get_data_from_transfer_queue("train", rollout_id, fields, batch_size, batch_index)) - 若已备份 ref 权重,执行
_switch_model("ref")并计算 ref log-probs - 若已备份 teacher 权重(OPD 场景),执行
_switch_model("teacher")并计算 teacher log-probs - 执行
_switch_model("old_actor" 或 "actor")并计算当前 actor 的 log-probs - 把扩充后的子批次追加到内存列表
- 从 TransferQueue 拉取数据(
合并子批次并做全局 Advantages 归一化
所有子批次 dict 被拼接成一个
rollout_data,随后compute_advantages_and_returns(self.args, rollout_data)在合并后的整批数据上执行一次。这是两阶段设计的 核心正确性要求 —— Advantages 归一化必须看到完整的 DP-group 批次,而不是各个子批次切片。在合并批次上训练并推送权重
一次
train(...)调用基于合并后的 batch 完成优化器步进。随后 actor 把新权重备份到actortag(如果到达 ref 更新间隔,也刷新reftag),然后调用self.update_weights()通过UpdateWeightFromTensor把最新权重同步给 rollout。
子批次 forward 控制了激活峰值显存(与 Fully Async 行为一致),而合并后的训练步则保留了 Colocate 风格的全局统计量。
配置
必需参数
| 参数 | 用途 |
|---|---|
--hybrid | 启用 Hybrid 模式(内部展开为 fully_async=True, colocate=True) |
--resource '{...}' | 分别声明 actor 与 rollout 的 placement group,例如 {"actor":[1,4],"rollout":[1,4]} |
--num-iters-per-train-update | 每个全局 batch 切分的子批次数量(越大 → 峰值显存越小,TransferQueue 轮询次数越多) |
--max-staleness | Off-policy 容忍度(0 = 严格 on-policy,>0 允许一定程度滞后) |
常用可选参数
| 参数 | 说明 |
|---|---|
--balance-data | Hybrid 模式下支持(纯 fully-async 下被拒绝)。启用后做 DP 间负载均衡。 |
--num-data-storage-units | TransferQueue 存储 actor 的数量。 |
--use-streaming-dataset | 从磁盘流式读取 prompts,而不是全量载入内存。 |
--ref-update-interval | 周期性地用最新 actor 权重刷新缓存的 ref 权重。 |
默认值覆盖
启用 --hybrid 后,relax/utils/arguments.py 会按以下方式设置默认值(除非用户显式传入):
offload_train = False且offload_rollout = False—— actor 与 rollout GPU 独立,不需要 offloadcompute_advantages_and_returns = True—— actor 必须在本地计算 advantagesfully_async = True、colocate = True—— 由--hybrid推导得出
WARNING
如果你既想做流式数据流水线,又需要 --balance-data,必须使用 --hybrid。--fully-async --balance-data(不带 --hybrid)会在参数解析阶段被拒绝。
快速开始
8 GPU 多模态 Hybrid 训练的参考启动脚本位于 scripts/training/multimodal/run-qwen35-9B-8xgpu-openr1mm-hybrid-async.sh。
它构建的 Hybrid 调用命令为:
ray job submit --address="http://127.0.0.1:8265" \
--runtime-env-json="${RUNTIME_ENV_JSON}" \
-- python3 -m relax.entrypoints.train \
--resource '{"actor": [1, 4], "rollout": [1, 4]}' \
--max-staleness 2 \
--num-data-storage-units 1 \
--num-iters-per-train-update 8 \
--balance-data \
--hybrid \
"${MODEL_ARGS[@]}" \
"${CKPT_ARGS[@]}" \
"${ROLLOUT_ARGS[@]}" \
"${OPTIMIZER_ARGS[@]}" \
"${GRPO_ARGS[@]}" \
"${PERF_ARGS[@]}" \
"${SGLANG_ARGS[@]}" \
"${MISC_ARGS[@]}"该配置的关键点:
- 8 GPU 总量,actor 与 rollout 各占 4 张
max-staleness 2—— actor 可以消费比最新权重落后最多 2 个 step 的 rollout 输出num-iters-per-train-update 8—— 每个全局 batch 在 forward 阶段被切分为 8 个子批次balance-data—— 启用 DP 间负载均衡- 算法采用 GRPO,附带
--use-kl-loss与--use-tis(这些是算法参数,与 Hybrid 正交)
故障排除
train_hybrid(rollout_id=N) batch_index=K stalled for ... seconds
该警告由 relax/backends/megatron/actor.py 抛出,发生在 actor 不断尝试拉取下一个子批次、partition 却始终未被标记 all_consumed 时。常见原因:
- Rollout 漏填了当前 partition(丢弃样本后未补齐)。
- Rollout 因健康检查失败或重启而处于暂停状态。
- Staleness 预算耗尽:rollout 必须等待新权重,因此无法继续产出数据。
在认定为代码 bug 前,请先排查 rollout 侧日志及 partition 状态。
--balance-data is not supported in pure fully-async mode
你同时传入了 --fully-async --balance-data,但缺少 --hybrid。请去掉 --balance-data,或者改用 --hybrid(Hybrid 模式原生支持 DP 数据均衡)。
Rollout 长时间看到旧权重
Hybrid 在每次 train_hybrid 结束时通过同步的 UpdateWeightFromTensor 路径推送权重。如果观察到权重更新间隔过大,请检查:
- Actor 日志中的
update_weights()耗时 - Rollout 健康检查是否触发了 actor 等待(权重同步前会调用
_check_services_health())
下一步
Hybrid 模式计划推进的工作:
- 接入 DCS 做权重同步 —— 用 Distributed Checkpoint Service 替换当前同步的
UpdateWeightFromTensor路径,使权重向 rollout 广播能够与下一轮训练迭代重叠,消除每次train_hybrid结束时残留的同步阻塞。 - 将
train_actor拆分为num_iters_per_train_update次训练迭代 —— 目前num_iters_per_train_update只对 forward 阶段做了切分,合并后的训练步仍然在整个全局 batch 上跑一次。下一步将训练步同样切成num_iters_per_train_update次,让优化器更新与 TransferQueue 数据消费形成流水线,并进一步压低训练侧峰值显存。
相关文档:
