diff --git a/tests/test_autobatching.py b/tests/test_autobatching.py index a05e0608..5169a53a 100644 --- a/tests/test_autobatching.py +++ b/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, diff --git a/torch_sim/autobatching.py b/torch_sim/autobatching.py index d4324166..07ee9525 100644 --- a/torch_sim/autobatching.py +++ b/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 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