feat(multi-gpu): integrate single-node off-policy and PPO data parallelism - #982
Merged
Conversation
* feat(offpolicy): add multi-GPU data-parallel rank topology Add training.devices config surface and rank topology for off-policy data-parallel training (issue #965, child of roadmap #964). Rank 0 owns the ExperimentTracker and canonical checkpoints and supervises ranks 1..N-1 via DpRankSupervisor; each spawned rank re-runs scripts/train_offpolicy.py with UNILAB_DP_* env, its own cuda device (cuda:devices[rank]) and seed+rank. Parameter sync between ranks is out of scope for this stage. null/[]/[0] degenerate to the exact current single-device behavior. * fix(offpolicy): wait for sibling ranks on normal rank-0 exit DpRankSupervisor previously terminated still-running rank subprocesses as soon as rank 0 finished, so sibling ranks never completed their own runs (exit 0 masked the teardown). Give ranks a grace window to finish; a rank that exceeds the grace or exits non-zero still fails the whole run. * style(ipc): annotate supervisor __exit__ as Literal[False] * style(tests): apply ruff format to dp_launcher tests
* feat(offpolicy): partition collector CPU cores per data-parallel rank Each rank's collector owns one contiguous CPU block (cpu_count // world_size, rank i -> [i*size, (i+1)*size); remainder CPUs keep default OS scheduling), routed through EnvCfg.cpu_ids into the MuJoCo BatchEnvPool only for the collector env — the learner-side num_envs=1 probe envs keep the base override untouched. training.dp_collector_cpu_ids allows an explicit per-rank override. Single-rank runs keep bit-identical behavior (no cpu_ids, host-sized thread budget); multi-rank runs resolve the torch thread budget against the rank's CPU share. Closes #966 * fix(offpolicy): derive DP world size from training.devices for rank 0 Rank 0 carries no UNILAB_DP_* environment (the supervisor only sets it for spawned ranks), so its collector CPU partition and thread budget silently fell back to the single-rank path and rank 0's BatchEnvPool workers stayed unpinned. Resolve world_size from training.devices instead; rank still comes from the env (0 for rank 0). Verified on a 2-GPU host: rank0 pool workers pinned to CPUs 0-63, rank1 to 64-127. * style: apply ruff format from make check
* feat(offpolicy): 多 learner 周期性参数平均同步（torch.distributed） - 新增 unilab.ipc.dp_sync.DpParameterSync：FileStore rendezvous 的 process group 封装，broadcast_from_rank0 / allreduce_mean 均为原地 集合通信，key 序首次固定，close 幂等。 - FastSACLearner.dp_sync_tensors() 返回 actor/qnet/qnet_target state_dict 的活引用加 log_alpha（不含 optimizer state）。 - DoubleBufferOffPolicyRunner：collector 启动前做 init broadcast； 每个 learner iteration 的 update 边界（finish_update 之后、 log_step 之前）按 dp_sync_interval 做参数 all-reduce 平均， 耗时以 dp_sync_time 计入 iter_metrics；manifest 记录 dp_sync。 - build_runner 冷路径装配：training.dp_sync_interval（默认 8，<1 即 ValueError），world_size>1 时构造 DpParameterSync，rendezvous 锚定 每个 run 唯一的 log 根目录（rank>0 走 UNILAB_DP_LOG_DIR）。 - world_size=1 路径不建 process group，逐位等价。 - 测试：gloo 双进程 broadcast/allreduce/key 序/原地语义；runner 集成 （fake learner + fake dp_sync）；build_runner 装配与非法值校验； NCCL 双卡冒烟（slow + cuda>=2 skipif）。 Validation: pytest tests/ipc tests/algos 全绿；全量非 slow 单测 1582 passed；make type 0 error。 Fixes #967 * fix(ipc): pin rank CUDA device and force TCP loopback for NCCL dp sync Two host-level NCCL failures surfaced in the 2-GPU acceptance run: - ranks other than 0 kept the default current device (cuda:0) while their learner tensors lived on cuda:1, hanging the first collective (NCCL watchdog timeout after 120s); - with P2P disabled the SHM transport crashed with CUDA illegal memory access during multi-tensor all_reduce (reproduced with repo-free code). DpParameterSync.start() now pins torch.cuda.set_device to the rank device and defaults NCCL_P2P_DISABLE/NCCL_SHM_DISABLE=1 (env override wins), so collectives use the TCP loopback transport. Verified: NCCL unit smoke passes in 3s; a 30-iteration K=1 two-rank run leaves rank0/rank1 checkpoints bitwise identical across all 77 synced tensors. * style: apply ruff format from make check
- benchmark/rl/benchmark_offpolicy_dp_scaling.py: subprocess 跑真实 train_offpolicy.py，对比 N=1 与 N-way DP；稳态吞吐取 tfevents perf/steps_per_sec 后 50% 采样点均值，聚合为各 rank 求和；输出 scaling ratio 与 1.7x 阈值 verdict；run 失败退出码非零 - tests/benchmark/test_offpolicy_dp_scaling_benchmark.py: 假 run 目录 (run_summary.json + 假 tfevents) 覆盖稳态均值、N=2 聚合、rank1 缺失报错、verdict 判定，不起真实训练 - docs: SAC 页补多卡数据并行用法与限制（单节点 / sac only / NCCL TCP loopback 原因）；支持矩阵结构无法表达多卡，未改生成器
FlashSACLearner 实现 dp_sync_tensors()（actor/critic/target_critic/ temperature state_dict 活引用），build_runner 的 flashsac 分支接通 dp_sync / dp_sync_interval / collector_cpu_ids，与 sac 共用同一套 runner 机制；单卡路径 dp_sync=None 保持不变。 Refs #969
* fix(offpolicy): stabilize multi-GPU graphs and logging Warm up and capture NCCL gradient collectives safely for FastSAC and FlashSAC CUDA graphs. Consolidate rank-zero terminal ownership, fixed-rate rolling statistics, aggregate throughput reporting, and synchronous collector telemetry. * fix(offpolicy): clean up multi-GPU interrupts * feat(ppo): add single-node RSL-RL multi-GPU training (#980)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This is the final umbrella integration PR for #964. The implementation work was split across reviewed child PRs and accumulated on
feat/issue-964-offpolicy-multi-gpu-dp; this PR does not introduce a new implementation scope.Final off-policy contract
The original roadmap proposed periodic parameter averaging through
training.dp_sync_interval. The final implementation and #964 have been explicitly revised:training.dp_sync_intervaltraining.devicespreserves the existing single-GPU pathThis uses runner-owned collectives rather than a PyTorch DDP wrapper and does not introduce a second collector/learner lifecycle.
Included work
Validation
The final PR head tree is identical to the tree validated by:
make test-all— 1657 passed, 24 skipped, 1 xfailed; coverage 72%; benchmark import smoke passed--cfg job, resolving the documented 2/4/8-device lists; the hardware validation boundary remains the 2-GPU platform stated belowtaskset -c 0) IPC regression — 10 passed, 1 slow test deselected; sparse-statistics worker stress — 10/10 passedDone., without a second NCCL process groupThe branch is 13 commits ahead of
main; a merge-tree check reports no conflicts.Throughput
Off-policy validation used 2 × RTX 6000D and Threadripper 9980X with the production FastSAC configuration:
PPO results are weak-scaling measurements with 1024 environments per rank, rollout length 24, and 10 steady-state iterations after discarding 2 warmup iterations:
No PPO model-quality claim is made; acceptance is semantic correctness, lifecycle completion, and throughput reporting.
Support boundary
NCCL_P2P_DISABLE=1andNCCL_SHM_DISABLE=1are defaults for the validated platform; explicit user environment values take precedenceCloses #964
Related: #978