Catch ROCM/HIP/AMD oom in should_reduce_batch_size (#812)

mcleish7 · clefourrier · web-flow · commit 95ec669127d3 · 2025-06-23T19:16:26.000+02:00
* Catch ROCM/HIP oom in should_reduce_batch_size

* fix formatting

---------

Co-authored-by: Clémentine Fourrier &lt;22726840+clefourrier@users.noreply.github.com&gt;
diff --git a/src/lighteval/utils/parallelism.py b/src/lighteval/utils/parallelism.py
@@ -50,6 +50,7 @@ def should_reduce_batch_size(exception: Exception) -> bool:
         "CUDA out of memory.",  # CUDA OOM
         "cuDNN error: CUDNN_STATUS_NOT_SUPPORTED.",  # CUDNN SNAFU
         "DefaultCPUAllocator: can't allocate memory",  # CPU OOM
+        "HIP out of memory",  # ROCM OOM
     ]
     if isinstance(exception, RuntimeError) and len(exception.args) == 1:
         return any(err in exception.args[0] for err in _statements)

Original file line number	Diff line number	Diff line change
`@@ -50,6 +50,7 @@ def should_reduce_batch_size(exception: Exception) -> bool:`
`50`	`50`	`"CUDA out of memory.", # CUDA OOM`
`51`	`51`	`"cuDNN error: CUDNN_STATUS_NOT_SUPPORTED.", # CUDNN SNAFU`
`52`	`52`	`"DefaultCPUAllocator: can't allocate memory", # CPU OOM`
	`53`	`+ "HIP out of memory", # ROCM OOM`
`53`	`54`	`]`
`54`	`55`	`if isinstance(exception, RuntimeError) and len(exception.args) == 1:`
`55`	`56`	`return any(err in exception.args[0] for err in _statements)`