Critic FP32 masters: one gather buffer per leaf; single-rank DP gathers on the CPU (CompactionRL G3 OOM)

[michaellchung]

背景

s19-compaction-g3-20261010a（CompactionRL G3，两岛 strict-avg，Qwen3.5-0.8B，16k 序列）失败：

• 第 0 轮两岛都训完了。
• 岛 0 的 critic 在 export_critic_tensors → _critic_masters → full_masters 中 CUDA OOM：要分配 3.25 GiB，整卡只剩 1.46 GiB。
• 随后报 CrossChannelCommitError，岛 0 退出，head job FAILED。

根因

1. 导出时 actor 刚训完长序列，仍占着 GPU：critic sleep 下整卡已用 58.9 GB；critic 唤醒后只剩 4.2 GB 空闲。
2. full_masters 先为每个参数在 GPU 上建一份 fp32 全量张量，再用 torch.cat 拼出第二份，峰值约 2×3.25 GiB。

为什么之前没出问题：

• G1 单岛不同步，不走这条导出路径。
• SAO/PPO G3 序列短（1024），actor 训完残留的显存小。

改动

• full_masters：每个优化器叶子只分配一个扁平 fp32 缓冲区，各主权重是它的视图，不再有拼接副本；新增可选参数 device。
• _critic_masters：torch.distributed 已初始化且每个 DP 组只有 1 个 rank 时，在 CPU 上拼，不做集合通信。多 rank 组保持 GPU 缓冲区 + all-reduce；CPU 缓冲区遇到多 rank 组时拒绝执行。
• _write_critic_masters：targets 不再整体搬上 GPU，由 copy_ 逐个搬。stale 检查改为先把主权重转到参数所在设备再比较。

子 agent 代拍板 + 理由：只改 yeto 插件，不改 Miles 的 sleep 顺序。理由是最少改动，且对 DP>1 的行为不变。

测试

• 新增 test_single_rank_dp_gathers_on_cpu_in_one_buffer，检查三点：一个缓冲区、不调用集合通信、CPU 上 export/import 往返哈希一致。
• 本机 state_plugin/critic/compaction 相关 646 个通过。
• golden 相关 227 个通过。test_delta_protocol 的 12 个报错在 main 上同样存在，与本改动无关。
• scripts/precommit_check.py rc=0。
• 上卡：未验证。

🤖 Generated with Claude Code