Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Skip trial instead of aborting when CUDA graph capture fails#605

Open
noahfarr wants to merge 1 commit into
PufferAI:4.0from
noahfarr:fix/cudagraph-oom-graceful-skip
Open

Skip trial instead of aborting when CUDA graph capture fails#605
noahfarr wants to merge 1 commit into
PufferAI:4.0from
noahfarr:fix/cudagraph-oom-graceful-skip

Conversation

@noahfarr

@noahfarr noahfarr commented Jul 2, 2026

Copy link
Copy Markdown

cudaStreamEndCapture / cudaGraphInstantiate / cudaGraphDestroy were guarded by assert(), which aborts the whole process on failure. When running PROTEIN my runs sometimes OOM when a trial's config sits near the VRAM ceiling. Because graph capture happens during create_pufferl_impl's warmup, a single sweep trial can take down the entire sweep process rather than being skipped. Throw std::runtime_error instead.

cudaStreamEndCapture / cudaGraphInstantiate / cudaGraphDestroy were
guarded by assert(), which abort()s the whole process on failure. These
calls allocate GPU memory and fail with OOM when a trial's config sits
near the VRAM ceiling. Because graph capture happens during
create_pufferl_impl's warmup, a single memory-hungry sweep trial would
take down the entire sweep process rather than being skipped.

Throw std::runtime_error instead. The exception propagates up through
create_pufferl to the existing RuntimeError handler in pufferl.py, which
already skips a trial and continues the sweep (same path as the
"CUDA OOM: failed to allocate training buffers" case). Applied to both
the rollout (net_callback_wrapper) and train (train_impl) graph captures;
the graph is destroyed before throwing on an instantiate failure to avoid
leaking it.
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.

1 participant