【Task.23】Colocate 纯文本性能：降低权重同步窗口的进程组拆建开销

[ldemon2333]

Task 23 — Colocate 纯文本性能：降低权重同步窗口的进程组拆建开销

Author: @ldemon2333, @Religious-J
Task tracker: Task 23 — Colocate 纯文本性能（#86）
Report issue: redai-infra/Relax#230
Upstream PR: redai-infra/Relax#229

本报告记录一次针对 colocate 模式下 actor/rollout 切换开销的优化，包括问题定位、实现机制、实测结果与适用边界。

实验数据包括本地 4×A6000 上 Qwen3-4B 单次运行内的 4 个稳态 step，以及 8×H20 上 Qwen3-30B-A3B 的单个稳态 step 对照。两组实验的硬件、模型、并行拓扑和实现版本不同，因此分别报告，不做横向性能比较，也不据此推断其他配置下的收益。

---

0. 结论速览

在 4×A6000 的 Qwen3-4B 实验中，稳态步（step 2–5，n=4）的切换阶段总开销如下：

方案	sleep	update_weights	wake_up	切换合计
基线 1f33585	5.34	3.99	2.24	11.56 s
优化后 053a678	2.46	2.17	2.15	6.79 s
差值	−2.88	−1.82	−0.09	−4.77 s (−41.3 %)

基线与优化方案在这 4 个 step 内的标准差分别为 0.12 s 和 0.11 s，观测均值相差 4.77 s。由于这些 step 来自各自的一次连续运行，而非多次独立重复实验，这里的标准差只描述运行内波动，不能作为跨运行不确定性的估计。

在 8×H20 的 Qwen3-30B-A3B 补充实验中，单个稳态 step 的切换合计从 23.507 s 降至 18.094 s，观测差值为 −5.413 s（−23.03%），相当于基线整步时间的 1.31%；观测吞吐差值约为 +0.8%。由于每个方案只有一个稳态 step，且优化方案的 actor train tokens 少 1.48%，这些结果不足以估计稳定的端到端收益。

在 4B 实验中，4.77 s 的切换阶段差值约占基线稳态平均 step_time（991.94 s）的 0.48%。理论上，固定生命周期开销在整步时间较短时占比更高；但其绝对值也可能随硬件、进程组数量和并行拓扑变化，不能用本实验直接外推。

---

1. 环境信息

项	值
硬件	4 × NVIDIA RTX A6000（48 GB，Ampere），NVLink bridge on 0-1 / 2-3
驱动	580.173.02
PyTorch / CUDA	2.11.0+cu129 / 12.9
镜像	ghcr.io/redai-infra/relaxrl:dev-20260715-8325919e
基线 commit	1f33585 feat(megatron): add selective CPU offload… (#85)
优化 commit	053a678（含 ad4af04）
模型	Qwen3-4B
数据集	dapo-math-17k

并行与工作负载：

项	值
训练侧	TP=2, PP=1, CP=1, EP=1, DP=2
推理侧	单 engine 跨 4 卡，--rollout-num-gpus-per-engine 4
rollout-batch-size	16
n-samples-per-prompt	8
global-batch-size	64
max-response-len	8192
max-tokens-per-gpu	4096
激活重算	full / uniform / 1
sglang-mem-fraction-static	0.65
NUM_ROLLOUT	6

---

2. 待优化路径：切换窗口里的进程组生命周期操作

colocate 模式下，actor 与 rollout 时分复用同一组 GPU。每步训练结束后的切换序列为：

train → save_model → sleep() → update_weights() → rollout generate → wake_up() → train

sleep() 调用 torch_memory_saver.pause()，在 VMM 层面暂停训练侧显存页，把显存让给 SGLang。NCCL communicator 使用的 native resources 不由 PyTorch allocator 统一管理，TMS pause 不会处理这些资源。基线路径因此选择在 pause 前调用 destroy_process_groups()，销毁训练侧 NCCL 组。

该路径包含以下两类可测量的额外操作。

2.1 拆组后的 2 秒固定等待

destroy_process_groups() 末尾有一段无条件的 time.sleep(2.0)：

# Wait for OS to release NCCL socket ports (TCP TIME_WAIT),
# preventing "Address already in use" on subsequent reload.
time.sleep(post_destroy_delay)

这段等待用于规避 NCCL socket 端口尚未释放时，下次建组可能出现的 Address already in use。从该进程的执行路径看，这 2 秒没有执行后续工作，因此会直接增加切换阶段时长；本文没有单独测量等待期间整个系统的 CPU/GPU 利用率。

2.2 一步之内拆建两次

此外，update_weights() 内部需要做 TP all-gather，把分片权重拼成完整张量并导出给 SGLang（weight_update/common.py），因此需要 reload_process_groups() 重新创建刚销毁的组；导出完成后再次销毁。

于是基线的单步路径包含 destroy → reload → destroy 循环，每次 destroy 各带一个 2 s 等待。基线日志中每步可见两条：

Destroyed 14 process groups, waiting 2.0s for NCCL socket port release

这些操作不随该 step 的训练 token 数直接缩放，但实际耗时可能受进程组数量、拓扑、硬件和系统状态影响。

---

3. 改动内容

3.1 机制一：延迟端口冷却（_defer_port_release_wait）

文件：relax/utils/reloadable_process_group.py

不再在 destroy 后原地 sleep(2.0)，而是记录一个截止时刻，把等待推迟到下一次 reload_process_groups() 真正需要端口时才补足：

def _defer_port_release_wait(pid: int, delay: float) -> None:
    """Overlap NCCL socket cooldown with useful work before the next reload."""
    _port_release_deadline_by_pid[pid] = max(
        _port_release_deadline_by_pid.get(pid, 0.0),
        time.monotonic() + delay,
    )

def _wait_for_deferred_port_release(pid: int) -> None:
    deadline = _port_release_deadline_by_pid.pop(pid, None)
    if deadline is None:
        return
    remaining = deadline - time.monotonic()
    if remaining > 0:
        time.sleep(remaining)

TIME_WAIT 按墙钟时间推进，与进程是否阻塞无关。在本次 4B 运行中，destroy 与下一次 reload 之间的 rollout 生成持续数百秒，因此预期能够覆盖 2 秒冷却窗口。是否仍需补足等待由运行时的 remaining 判断；报告没有单独列出每次 remaining 的日志。

当运行到 reload 时冷却窗口尚未结束，代码仍会等待剩余时间。该机制保留了原有等待条件，但仍需通过不同网络和 NCCL 配置验证其兼容性。

3.2 机制二：跨 TMS pause 保活进程组（PG_PAUSED_LIVE）

文件：relax/backends/megatron/actor.py

引入三态机 PG_ACTIVE / PG_PAUSED_LIVE / PG_DESTROYED。当条件满足时，sleep() 不拆组，直接进入 TMS pause，让 communicator 存活穿过整个 pause 窗口；update_weights() 随后直接复用，省掉 reload。

这依赖一个前提：TMS 的 VMM pause 不触碰 NCCL 自行分配的显存，communicator 因而可能在 pause 期间继续可用。实现通过运行期 canary 检查当前运行中的通信组；canary 通过只能说明本次检查成功，不能证明所有硬件、拓扑和故障场景均安全。

11 条静态门控（_local_pg_preservation_eligibility）：

self.role == "actor"
and self.args.colocate
and self.args.offload_train
and self._per_step_rollout
and self._train_state_offloader.uses_tms      # 仅 TMS 路径
and not self.args.use_critic
and not self.args.fully_async
and getattr(self.args, "enable_weights_backuper", False)
and isinstance(getattr(self, "weight_updater", None), UpdateWeightFromTensor)
and mpu.get_pipeline_model_parallel_world_size() == 1
and mpu.get_context_parallel_world_size() == 1
and mpu.get_expert_model_parallel_world_size() == 1
and not device_utils.is_klx()

运行期 canary（_run_preserved_pg_canary）：pause 之后对 TP 组和 DP 组各跑一次 all-reduce，校验结果等于 world_size。检测到异常时，实现会退回旧路径，并在当前进程生命周期内关闭保活。

三项回退与一致性措施：

1. 所有 rank 用 Gloo 组 all_reduce(MIN) 汇总决策（_all_ranks_agree），降低部分 rank 保活、部分销毁所导致的死锁风险。
2. 两段式启用：首次 sleep() 仅置 eligible=True，仍走原有路径；一次本地 IPC 权重同步完成后才置 ready=True。
3. update_weights() 整体包在 try/except 中，失败时调用 destroy_process_groups(post_destroy_delay=0)，用于回收仍存活的训练侧 communicator。

---

4. 实测数据

4.1 实验方法

基线与优化方案各运行 6 步，使用 step 0–1 预热，并统计 step 2–5。每个方案只有一次运行，未进行随机顺序、交替运行或多次独立重复，因此结果可能包含运行顺序、系统状态和 rollout 随机性的影响。

4.2 逐步原始数据

基线 1f33585（单位 s）：

step	sleep	upd_w	wake	切换合计	train_wait	step_time
0	16.11	3.56	1.62	21.30	653.54	1126.12
1	12.71	6.13	2.22	21.06	509.62	944.03
2	5.42	3.80	2.28	11.50	512.90	956.26
3	5.11	4.10	2.19	11.40	515.68	966.42
4	5.36	4.20	2.17	11.73	611.36	1158.27
5	5.47	3.85	2.31	11.62	469.87	886.79

优化后 053a678（单位 s）：

step	sleep	upd_w	wake	切换合计	train_wait	step_time
0	16.60	3.53	1.48	21.61	651.29	1115.52
1	10.70	2.29	2.10	15.10	498.40	931.57
2	2.46	2.42	2.00	6.88	500.14	945.39
3	2.47	2.24	1.99	6.71	504.39	948.55
4	2.43	1.87	2.36	6.66	608.64	1140.84
5	2.48	2.15	2.26	6.90	478.36	907.31

4.3 稳态统计（step 2–5，n=4）

指标	基线 mean±sd	优化 mean±sd	差值	相对
sleep	5.34 ± 0.14	2.46 ± 0.02	−2.88	−53.9 %
update_weights	3.99 ± 0.17	2.17 ± 0.20	−1.82	−45.6 %
wake_up	2.24 ± 0.06	2.15 ± 0.16	−0.09	−4.0 %
切换合计	11.56 ± 0.12	6.79 ± 0.11	−4.77	−41.3 %

在这次运行中，sleep 与 update_weights 的均值分别下降 2.88 s 和 1.82 s，变化方向与跳过等待及减少 reload 的设计一致。

4.4 Qwen3-30B-A3B 补充实验

作为更大模型和不同并行拓扑下的补充观察，本文引用 pg-reuse-issue.md 中的 8×H20 实验。30B 优化实现以 Relax baseline 0694cd536aa9bddb9cd78452585e7f56feebd6a2 为基准开发，来源报告记录的 prototype commit 为 f98aa2bab591674b896183fff6b9f037bc0cbb2c。实际 A/B 在同一 PG-reuse worktree 上运行，通过开关控制是否启用复用。根据来源报告，两次运行使用相同模型、数据、训练参数和脚本，脚本 SHA256 相同；但 rollout 为独立采样。

实验配置：

项	值
硬件	8 × NVIDIA H20（97,871 MiB/GPU）
PyTorch / CUDA / NCCL	2.11.0+cu129 / 12.9 / 2.28.9
Relax baseline	0694cd536aa9bddb9cd78452585e7f56feebd6a2
PG-reuse prototype	f98aa2bab591674b896183fff6b9f037bc0cbb2c
A/B 控制方式	同一 PG-reuse worktree，通过运行时开关关闭或启用复用
模型	Qwen3-30B-A3B
数据集	dapo-math-17k
部署方式	actor/rollout colocate
训练并行	TP=4, CP=2, EP=8，DeepEP flex dispatcher
推理并行	SGLang TP=8
batch 配置	global batch size=256，rollout batch size=32，每个 prompt 采样 8 次
优化器	CPU optimizer offload

首个稳态 step 的结果：

指标	基线	优化后	变化
sleep_time	12.046 s	10.042 s	−2.003 s（−16.63%）
update_weights_time	9.090 s	6.156 s	−2.934 s（−32.27%）
wake_up_time	2.371 s	1.895 s	−0.476 s（−20.10%）
切换合计	23.507 s	18.094 s	−5.413 s（−23.03%）
step_time	412.579 s	403.170 s	−2.28%
step_token_per_s	4,150.01	4,183.82	+0.81%
step_resp_token_per_s	4,056.16	4,087.78	+0.78%
actor train tokens	1,712,206	1,686,793	−1.48%
峰值 GPU 显存	75,473 MiB	75,477 MiB	+4 MiB

在本次正确性检查范围内，pause 后 TP/DP/CP/EP communicator canary 通过，训练与权重同步完成，并进入下一轮 rollout；日志中未观察到进程组 reload 失败、weight-version mismatch、traceback 或 CUDA OOM。相关 focused tests 结果为 9 passed。这些检查降低了已覆盖路径的风险，但不能替代长时间运行、故障注入和更多拓扑测试。

进程组复用会在 pause 后的短窗口内保留 NCCL native resources。该次实验中，paused-window 的物理显存由约 8.86–8.87 GiB 上升到 12.45–14.24 GiB；同时，记录到的峰值 GPU 显存相差 +4 MiB。前者意味着该窗口可供 SGLang 使用的瞬时显存余量可能减少，但两项数据均来自单次对照，需要重复测量。

[xiaoliang0601]

❌ 验收不通过。

整体性能优化效果偏小，并且显存有提升。