Skip to content

perf: enable multi-GPU GPU-resident replay for SAC - #817

Merged
TATP-233 merged 2 commits into
mainfrom
perf/issue-813-multi-gpu-gpu-resident-replay
Jul 30, 2026
Merged

perf: enable multi-GPU GPU-resident replay for SAC#817
TATP-233 merged 2 commits into
mainfrom
perf/issue-813-multi-gpu-gpu-resident-replay

Conversation

@TATP-233

Copy link
Copy Markdown
Collaborator

Summary

Enable rank-local GPU-resident replay for multi-GPU SAC off-policy training on AMD ROCm. Each learner rank owns an independent GPU-resident mirror of the shared CPU replay buffer, eliminating per-tick collector CPU pack and per-rank batch H2D.

Fixes #813

Changes

  • New pipeline: MultiGPUGPUResidentReplayPipeline wraps the single-GPU GPUResidentReplayPipeline (feat(replay): opt-in GPU-resident replay pipeline for single-GPU off-policy #806) per rank, providing the same interface as MultiGPUCPUPinnedReplayPipeline without collector pack IPC.
  • Runner wiring: MultiGPUOffPolicyRunner selects pipeline via training.replay_pipeline; GPU-resident path skips collector pack queues/slots.
  • Config unlock: remove the single-GPU-only restriction in scripts/train_offpolicy.py for training.replay_pipeline=gpu_resident.
  • ROCm test compatibility: update repo hygiene / torch-source tests to accept both CUDA and ROCm configs so make test-all passes with pyproject.rocm.toml active.
  • Benchmark helper: add benchmark/rl/extract_offpolicy_metrics.py for TensorBoard metric extraction.

Benchmarks (8x AMD Instinct MI300X, ROCm 7.2)

Task: g1_motion_tracking / motrix, 30 iterations, last-20 average.

GPUs Pipeline iter_ms steps/s Collector/s collector replay_ms collector wait_ms rank barrier_ms
2 cpu_pinned 77.0 27,420 28,982 2.00 2.75 2.69
2 gpu_resident 81.5 25,482 24,751 1.67 4.96 2.38
4 cpu_pinned 86.4 24,815 22,556 1.95 11.19 2.19
4 gpu_resident 75.8 27,544 27,777 1.79 2.36 2.23
8 cpu_pinned 118.4 19,536 18,477 1.83 44.66 1.49
8 gpu_resident 86.2 25,146 24,802 1.75 12.71 2.31

Key results:

  • collector replay_ms stays 1.7–1.8 ms across 2/4/8 GPU (no rank degradation, DoD met).
  • 8 GPU vs CPU-pinned: iter_ms -27% (118.4 → 86.2), steps/s +29% (19,536 → 25,146), learner collector wait -72% (44.7 → 12.7 ms).

Microbenchmark baseline (host replay harness, capacity 1M rows): CPU sample+H2D 7.0 → 10.8 → 20.6 ms (1→4→8 GPU); GPU-resident sample 0.14 → 0.50 → 0.66 ms.

Validation

  • make test-all passed: 1634 passed, 26 skipped, 267 deselected, 1 xfailed.
  • Benchmark smoke test passed (34/35 module-mode, 35/36 script-mode; only mlx optional skipped).
  • Relevant unit tests: tests/ipc/test_replay_pipeline_gpu_resident.py, tests/ipc/test_multi_gpu_replay_pack.py, tests/algos/test_offpolicy_double_buffer_runner.py, tests/algos/test_offpolicy_runner_unit.py all pass.
  • End-to-end smoke: 2/4/8 GPU g1_motion_tracking / motrix training completes successfully.

Notes

- Add MultiGPUGPUResidentReplayPipeline: per-rank GPU-resident replay mirror
  wrapping the single-GPU GPUResidentReplayPipeline (#806) for multi-GPU
  off-policy training.
- Wire into MultiGPUOffPolicyRunner: select pipeline via
  training.replay_pipeline; GPU-resident path skips collector pack IPC.
- Remove the single-GPU-only restriction in scripts/train_offpolicy.py.
- Add ROCm compatibility to repo hygiene / torch-source tests so make
  test-all passes with pyproject.rocm.toml active.
- Add benchmark/rl/extract_offpolicy_metrics.py for TensorBoard metric
  extraction.

Benchmarks on 8x AMD Instinct MI300X (g1_motion_tracking/motrix, 30 iters):
- collector replay_ms stays ~1.7-1.8 ms across 2/4/8 GPU (DoD: no rank
  degradation).
- 8 GPU: iter_ms 118.4 -> 86.2 (-27%), steps/s 19,536 -> 25,146 (+29%),
  learner collector wait 44.7 -> 12.7 ms (-72%).
@TATP-233
TATP-233 requested a review from caozx1110 as a code owner July 30, 2026 10:07
CI (ubuntu-slim) has no CUDA device, and MultiGPUOffPolicyRunner.build_runner
requires a CUDA device for num_gpus > 1. Add a skipif so the test is skipped
there instead of failing.
@TATP-233
TATP-233 merged commit b2a00c4 into main Jul 30, 2026
6 checks passed
@TATP-233
TATP-233 deleted the perf/issue-813-multi-gpu-gpu-resident-replay branch July 30, 2026 10:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

perf: SAC replay 多卡(2/4/8 GPU)traffic scaling 与 rank-local GPU cache 验证 — 自 #694 拆出

2 participants