fix(autobatching): recognise Warp OOM in memory estimation and free cache before the real run

tests/test_autobatching.py

@@ -488,6 +488,63 @@ def mock_measure(*_args: Any, **_kwargs: Any) -> float:
     assert max_size == 2
 
 
+@pytest.mark.parametrize(
+    "oom_message",
+    [
+        "CUDA out of memory. Tried to allocate 20.00 MiB",
+        # Warp / nvalchemiops allocator (used by ORB v3 neighbor lists) phrases
+        # OOM differently and must still be recognised by the default matcher.
+        "Failed to allocate 2556 bytes on device 'cuda:0'",
+    ],
+)
+def test_determine_max_batch_size_recognises_oom_variants(
+    si_sim_state: ts.SimState,
+    lj_model: LennardJonesModel,
+    monkeypatch: pytest.MonkeyPatch,
+    oom_message: str,
+) -> None:
+    """OOM is detected for both PyTorch and Warp-style allocator messages.
+
+    Regression test: the default ``oom_error_message`` must cover the Warp
+    allocator wording, and a non-matching first entry in the message list must
+    not short-circuit the check before later entries are compared.
+    """
+    call_count = {"n": 0}
+
+    def mock_measure(*_args: Any, **_kwargs: Any) -> float:
+        call_count["n"] += 1
+        if call_count["n"] >= 3:  # OOM once the batch grows past a couple probes
+            raise RuntimeError(oom_message)
+        return 0.1
+
+    monkeypatch.setattr(
+        "torch_sim.autobatching.measure_model_memory_forward", mock_measure
+    )
+
+    # Uses the (broadened) default oom_error_message. Should degrade to a safe
+    # batch size instead of propagating the OOM RuntimeError.
+    max_size = determine_max_batch_size(si_sim_state, lj_model, max_atoms=10_000)
+    assert max_size >= 1
+
+
+def test_determine_max_batch_size_reraises_non_oom_error(
+    si_sim_state: ts.SimState,
+    lj_model: LennardJonesModel,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A genuine (non-OOM) error is still propagated, not swallowed."""
+
+    def mock_measure(*_args: Any, **_kwargs: Any) -> float:
+        raise RuntimeError("shape mismatch in einsum")
+
+    monkeypatch.setattr(
+        "torch_sim.autobatching.measure_model_memory_forward", mock_measure
+    )
+
+    with pytest.raises(RuntimeError, match="shape mismatch"):
+        determine_max_batch_size(si_sim_state, lj_model, max_atoms=10_000)
+
+
 @pytest.mark.parametrize("scale_factor", [1.1, 1.4])
 def test_determine_max_batch_size_small_scale_factor_no_infinite_loop(
     si_sim_state: ts.SimState,

torch_sim/autobatching.py

@@ -36,6 +36,12 @@
 
 logger = logging.getLogger(__name__)
 
+# Substrings used to recognise out-of-memory errors raised during the memory
+# estimation forward passes. Different backends word this differently: PyTorch
+# raises "CUDA out of memory", while Warp-based neighbor lists (e.g. the
+# nvalchemiops kernels used by ORB v3) raise "Failed to allocate <n> bytes".
+DEFAULT_OOM_ERROR_MESSAGES = ("CUDA out of memory", "Failed to allocate")
+
 
 def to_constant_volume_bins(  # noqa: C901
     items: dict[int, float] | list[Any],
@@ -192,7 +198,7 @@ def determine_max_batch_size(
     max_atoms: int = 500_000,
     start_size: int = 1,
     scale_factor: float = 1.6,
-    oom_error_message: str | list[str] = "CUDA out of memory",
+    oom_error_message: str | Sequence[str] = DEFAULT_OOM_ERROR_MESSAGES,
 ) -> int:
     """Determine maximum batch size that fits in GPU memory.
 
@@ -210,8 +216,9 @@ def determine_max_batch_size(
         scale_factor (float): Factor to multiply batch size by in each iteration.
             Defaults to 1.6.
         oom_error_message (str | list[str]): String or list of strings to match in
-            RuntimeError messages to identify out-of-memory errors. Defaults to
-            "CUDA out of memory".
+            RuntimeError messages to identify out-of-memory errors. If any
+            listed substring appears in the error, it is treated as OOM.
+            Defaults to ``("CUDA out of memory", "Failed to allocate")``.
 
     Returns:
         int: Maximum number of batches that fit in GPU memory.
@@ -251,19 +258,18 @@ def determine_max_batch_size(
         except Exception as exc:
             exc_str = str(exc)
             # Check if any of the OOM error messages match
-            for msg in oom_error_message:
-                if msg in exc_str:
-                    safe_size = sizes[max(0, sys_idx - 2)]
-                    logger.debug(
-                        "OOM at %d systems (%d atoms), returning safe batch size %d",
-                        n_systems,
-                        concat_state.n_atoms,
-                        safe_size,
-                    )
-                    return safe_size
+            if any(msg in exc_str for msg in oom_error_message):
+                safe_size = sizes[max(0, sys_idx - 2)]
+                logger.debug(
+                    "OOM at %d systems (%d atoms), returning safe batch size %d",
+                    n_systems,
+                    concat_state.n_atoms,
+                    safe_size,
+                )
+                return safe_size
 
-                # No OOM message matched - re-raise the error
-                raise
+            # Not an OOM error - re-raise
+            raise
 
     return sizes[-1]
 
@@ -442,6 +448,15 @@ def estimate_max_memory_scaler(
     min_state_max_batches = determine_max_batch_size(min_state, model, **kwargs)
     max_state_max_batches = determine_max_batch_size(max_state, model, **kwargs)
 
+    # The probing above deliberately grows batches until an OOM, which leaves
+    # PyTorch's caching allocator holding most of the device memory. Release it
+    # so the real optimization run - and in particular any separate allocator
+    # such as the Warp/cudaMallocAsync pool used by ORB's neighbor lists - is
+    # not starved on the first real forward pass.
+    if torch.cuda.is_available():  # pragma: no cover
+        torch.cuda.synchronize()
+        torch.cuda.empty_cache()
+
     return min(
         min_state_max_batches * min_metric.item(),
         max_state_max_batches * max_metric.item(),
@@ -500,7 +515,7 @@ def __init__(
         max_atoms_to_try: int = 500_000,
         memory_scaling_factor: float = 1.6,
         max_memory_padding: float = 1.0,
-        oom_error_message: str | list[str] = "CUDA out of memory",
+        oom_error_message: str | Sequence[str] = DEFAULT_OOM_ERROR_MESSAGES,
     ) -> None:
         """Initialize the binning auto-batcher.
 
@@ -528,8 +543,9 @@ def __init__(
             max_memory_padding (float): Multiply the auto-determined max_memory_scaler
                 by this value to account for fluctuations in max memory. Defaults to 1.0.
             oom_error_message (str | list[str]): String or list of strings to match in
-                RuntimeError messages to identify out-of-memory errors. Defaults to
-                "CUDA out of memory".
+                RuntimeError messages to identify out-of-memory errors. If any
+                listed substring appears in the error, it is treated as OOM.
+                Defaults to ``("CUDA out of memory", "Failed to allocate")``.
         """
         self.max_memory_scaler = max_memory_scaler
         self.max_atoms_to_try = max_atoms_to_try
@@ -804,7 +820,7 @@ def __init__(
         memory_scaling_factor: float = 1.6,
         max_iterations: int | None = None,
         max_memory_padding: float = 1.0,
-        oom_error_message: str | list[str] = "CUDA out of memory",
+        oom_error_message: str | Sequence[str] = DEFAULT_OOM_ERROR_MESSAGES,
     ) -> None:
         """Initialize the hot-swapping auto-batcher.
 
@@ -835,8 +851,9 @@ def __init__(
             max_memory_padding (float): Multiply the auto-determined max_memory_scaler
                 by this value to account for fluctuations in max memory. Defaults to 1.0.
             oom_error_message (str | list[str]): String or list of strings to match in
-                RuntimeError messages to identify out-of-memory errors. Defaults to
-                "CUDA out of memory".
+                RuntimeError messages to identify out-of-memory errors. If any
+                listed substring appears in the error, it is treated as OOM.
+                Defaults to ``("CUDA out of memory", "Failed to allocate")``.
         """
         self.model = model
         self.memory_scales_with = memory_scales_with