perf: enable multi-GPU GPU-resident replay for SAC - #817
Merged
Conversation
- 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%).
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.
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
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
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
MultiGPUGPUResidentReplayPipelinewraps the single-GPUGPUResidentReplayPipeline(feat(replay): opt-in GPU-resident replay pipeline for single-GPU off-policy #806) per rank, providing the same interface asMultiGPUCPUPinnedReplayPipelinewithout collector pack IPC.MultiGPUOffPolicyRunnerselects pipeline viatraining.replay_pipeline; GPU-resident path skips collector pack queues/slots.scripts/train_offpolicy.pyfortraining.replay_pipeline=gpu_resident.make test-allpasses withpyproject.rocm.tomlactive.benchmark/rl/extract_offpolicy_metrics.pyfor TensorBoard metric extraction.Benchmarks (8x AMD Instinct MI300X, ROCm 7.2)
Task:
g1_motion_tracking/motrix, 30 iterations, last-20 average.Key results:
replay_msstays 1.7–1.8 ms across 2/4/8 GPU (no rank degradation, DoD met).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-allpassed: 1634 passed, 26 skipped, 267 deselected, 1 xfailed.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.pyall pass.g1_motion_tracking/motrixtraining completes successfully.Notes
pyproject.toml/uv.lockwere swapped bymake sync-rocmfor AMD development and are intentionally not part of this PR.