Skip to content

[Performance] Keep the DreamerV3 replay write-back off the learner stream - #4307

Open
vmoens wants to merge 2 commits into
inference-server-pinned-passfrom
dreamerv3-learner-cpu-offload
Open

[Performance] Keep the DreamerV3 replay write-back off the learner stream#4307
vmoens wants to merge 2 commits into
inference-server-pinned-passfrom
dreamerv3-learner-cpu-offload

Conversation

@vmoens

@vmoens vmoens commented Sep 8, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Stack: #4305 → #4306 → #4307 → #4308; depends on #4306.

Enqueue CUDA→CPU replay copies at submission; the update thread waits on CUDA events, avoiding copies queued behind the next learner step. Also deduplicate coordinates and group ensemble writes with stable sorts.

index, generation, patch = replay_context_update(sample, state, belief)
future = replay_buffer.submit_update_if_present(
    index=index, generation=generation, patch=patch
)

Reported validation: CPU DreamerV3/replay tests pass except eight missing-h5py failures. CUDA validation pending.

@pytorch-bot

pytorch-bot Bot commented Sep 8, 2026 •
edited
Loading

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4307

Note: Links to docs will display an error until the docs builds have been completed.

❌ 11 New Failures

As of commit 1c19214 with merge base 14324ea (image):

NEW FAILURES - The following jobs have failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Sep 8, 2026
@github-actions github-actions Bot added Performance Performance issue or suggestion for improvement Benchmarks rl/benchmark changes Objectives sota-implementations/ ReplayBuffers labels Sep 8, 2026
@vmoens
vmoens force-pushed the dreamerv3-learner-cpu-offload branch from 726f85c to d1ab7fc Compare September 8, 2026 15:55
@vmoens
vmoens changed the base branch from main to inference-server-pinned-pass September 8, 2026 15:55
This was referenced Sep 8, 2026
Open
Open
Open
@vmoens vmoens added the benchmarks/trigger Run the full benchmark suite on main right after this PR merges label Sep 8, 2026
@vmoens

vmoens commented Sep 8, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author 
Projected GPU: ~1.4 -> 0.7–0.9 s/update (~1.5–2x faster)
Measured CPU:  overall neutral; dedup 6% faster, write-back 17% faster

The GPU estimates remain unmeasured on the profiled node. Faster learner updates should raise collection throughput at a fixed train ratio.

vmoens and others added 2 commits September 8, 2026 17:34
@vmoens @claude
[Performance] Keep the DreamerV3 replay write-back off the learner st…
cdacfb3 
…ream

Profiling a 64-environment DreamerV3 Minecraft run on a 4-GPU node showed
one learner update taking 1.37 s with the GPU busy about a quarter of the
time. Three CPU-side items stood out: replay_context_update at 18% of the
main thread (about 0.25 s per update), the update thread spending 17% of
wall time in _validate_conditional_patch, and the prefetch thread spending
13% in sampling.

The per-line profile attributes almost all of the first item to the
host-to-device copy of the deduplicated index, not to the argsort loop
(about 0.5 ms). A pageable copy synchronizes the stream, so the main thread
waited for the CUDA-graph step there. The second item is the blocking
device-to-host copy of the patch, issued from the update thread after the
next graph replay was enqueued: it waited for that replay, and the sample
prefetched after the update inherited the delay, which serialized sampling
with the learner step.

- replay_context_update packs the coordinates into one key and runs a
  single stable sort (per-column fallback if the key would overflow int64),
  and moves the kept index to the learner device through pinned memory with
  non_blocking=True. The main thread no longer synchronizes there.
- ReplayBuffer.submit_update_if_present copies CUDA inputs destined for a
  CPU storage to pinned host memory on the caller's current stream at
  submission time and records an event; the update thread waits for the
  event instead of issuing blocking copies. Later work on the stream does
  not delay the update or the samples that depend on it.
- ReplayBufferEnsemble.update_if_present moves the patch to the members'
  common storage device once and groups the records per member with one
  stable sort and bincount instead of a boolean mask pass per member.
- bench_dreamer_v3_learner.py gains --device cpu and updates_per_second so
  the replay workload can be timed without CUDA.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
@vmoens
[BugFix] Match replay patch indices to the patch device
1c19214
@vmoens
vmoens force-pushed the dreamerv3-learner-cpu-offload branch from d1ab7fc to 1c19214 Compare September 8, 2026 16:35
@vmoens

vmoens commented Sep 8, 2026

Copy link
Copy Markdown
Collaborator Author

Measured on the target hardware (64-environment DreamerV3 Minecraft run, one 4-GPU node, process-slot inference server with a static batch of 64, compiled and CUDA-graph learner step, bf16, train ratio 16), comparing the full stack #4305 + #4306 + #4307 + #4308 (with transition_chunk_size=64) against the same configuration without these four PRs, over the same elapsed window of the two runs (1,000 s to 4,428 s after training start):

environment steps/s learner updates/s (32x256)
without the stack 361 0.70
with the stack 1,249 2.44
speed-up 3.46x 3.46x

The unthrottled collection phase before training reached 1,254 steps/s (298 before), the learner GPU now shows sustained full-utilization bursts instead of about a quarter duty cycle, and the run had no errors over the measured 1.3 hours. The two rates move together because the train ratio couples them, so the stack lifted both the collection ceiling and the learner cadence past the previous limits. This measures the combined effect of the four PRs; per-PR attribution has not been measured.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

benchmarks/trigger Run the full benchmark suite on main right after this PR merges Benchmarks rl/benchmark changes CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. Objectives Performance Performance issue or suggestion for improvement ReplayBuffers sota-implementations/

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@vmoens