Skip to content

[Performance] DreamerV3: freeze the setup heap before the training loop - #4308

Open
vmoens wants to merge 2 commits into
dreamerv3-learner-cpu-offloadfrom
collector-filelock-audit-hook
Open

[Performance] DreamerV3: freeze the setup heap before the training loop#4308
vmoens wants to merge 2 commits into
dreamerv3-learner-cpu-offloadfrom
collector-filelock-audit-hook

Conversation

@vmoens

@vmoens vmoens commented Sep 8, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

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

Freeze the DreamerV3 setup heap so cyclic GC stops rescanning it during training.

# After learner, collector, and replay setup:
gc.freeze()
for _ in collector:
    ...

Reported validation: CPU SOTA smoke and lint pass; local reproduction reduces GC time ~8×. Profiled GPU configuration still needs validation.

@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/4308

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

❌ 8 New Failures, 1 Pending

As of commit 7ca98b5 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 sota-implementations/ Performance Performance issue or suggestion for improvement labels Sep 8, 2026
@vmoens
vmoens force-pushed the collector-filelock-audit-hook branch from 0b0e1f0 to 9be4cce Compare September 8, 2026 15:55
@vmoens
vmoens changed the base branch from main to dreamerv3-learner-cpu-offload 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 
Measured local GC time:      0.208 -> 0.025 s (~8x less)
Measured local receive loop: 165 -> ~108 us/transition

Projected end-to-end gain is at most a few percent and remains unmeasured on the GPU node. The benefit should shrink after #4305 reduces unpickling.

vmoens and others added 2 commits September 8, 2026 17:34
@vmoens @claude
[Performance] DreamerV3: freeze the setup heap before the training loop
5215ed1 
Profiling a 64-environment DreamerV3 Minecraft run on a 4-GPU node with
process environment workers attributed 7% of the driver main thread
(218 of 3000 py-spy samples at 200 Hz over 15 s) to filelock's
sys.addaudithook callback, reached from AsyncBatchedCollector._rollout_frames
through multiprocessing.Queue.get and _ForkingPickler.loads, once per
transition received from the workers.

The hook itself is not the cost: it fires on pickle.find_class, 8 times per
consolidated transition regardless of its size, for 1-2 us per unpickle.
Since Python 3.12 the cyclic GC is scheduled on the eval breaker and runs at
the next check, which during a C-level unpickle is the RESUME of the first
Python callback, i.e. the audit hook. The samples are full collections
re-traversing the large setup heap, triggered by the unpickler's
allocations; without filelock they land on torch's rebuild_storage_fd.

Freeze the setup objects once everything is built so later collections only
visit what the loop allocates. Locally (3000 unpickles over a 1M-dict live
heap) this cuts GC time from 0.21 s to 0.025 s and the per-item loop from
165 us to about 108 us.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
@vmoens
[BugFix] Unfreeze Dreamer setup objects during cleanup
7ca98b5
@vmoens
vmoens force-pushed the collector-filelock-audit-hook branch from 9be4cce to 7ca98b5 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 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 sota-implementations/

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@vmoens