Skip trial instead of aborting when CUDA graph capture fails#605
Open
noahfarr wants to merge 1 commit into
Open
Skip trial instead of aborting when CUDA graph capture fails#605noahfarr wants to merge 1 commit into
noahfarr wants to merge 1 commit into
Conversation
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.
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.
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.