Skip to content

Commit

Permalink
fix(pt): improve out-of-memory capture (#3857)
Browse files Browse the repository at this point in the history
I just received another error message that reports out of memory. It's a
bad design of PyTorch that all errors use a general `RuntimeError`.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **Bug Fixes**
  - Improved out-of-memory error detection for CUDA driver issues.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Jinzhe Zeng <[email protected]>
  • Loading branch information
njzjz authored Jun 7, 2024
1 parent 674bad7 commit 057dc11
Showing 1 changed file with 1 addition and 0 deletions.
1 change: 1 addition & 0 deletions deepmd/pt/utils/auto_batch_size.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@ def is_oom_error(self, e: Exception) -> bool:
# (the meaningless error message should be considered as a bug in cusolver)
if isinstance(e, RuntimeError) and (
"CUDA out of memory." in e.args[0]
or "CUDA driver error: out of memory" in e.args[0]
or "cusolver error: CUSOLVER_STATUS_INTERNAL_ERROR" in e.args[0]
):
# Release all unoccupied cached memory
Expand Down

0 comments on commit 057dc11

Please sign in to comment.