Autobatcher memory probe does not recognize cuequivariance out-of-memory errors

Open Beginner friendly
  • Is something wrong?
#627 5 comments 1 reaction 0 assignees View on GitHub

Maintainers usually reply within 2 days

@lil-lon is already working on this.

Since Oct 31, 2026.

  • #631 by @lil-lon — open

Assessment

Difficulty
2/5
Estimated time
1-3 hours
Newbie friendliness
70/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Active
Tech stack
python, pytorch
Domain
backend, performance

Research direction

Locate determine_max_batch_size and DEFAULT_OOM_ERROR_MESSAGES in the autobatching code. Read how probe exceptions are currently recognized, then check the existing tests for out-of-memory handling. Add coverage for torch.OutOfMemoryError and a plain runtime error containing “out of memory”; done means both stop the probe cleanly while unrelated errors still propagate.

Written by the indexing model from the issue text.

Description

Summary

determine_max_batch_size grows a batch until the model runs out of GPU memory, and recognizes that point by matching the exception message against DEFAULT_OOM_ERROR_MESSAGES = ("CUDA out of memory", "Failed to allocate"). With MACE using cuequivariance kernels (enable_cueq=True), the out-of-memory error at the last probe step can come from cuequivariance instead of PyTorch:

RuntimeError: CUDA error encountered at: file=.../cuequivariance_ops/cuda/equivariance/uniform1d/kernel_generator.cu line=202: call='cudaMallocAsync(&sort_permutation, s.a * s.b * s.c * index_elem_size, stream)', Reason=cudaErrorMemoryAllocation:out of memory

This message matches neither entry, so the probe re-raises it, and ts.optimize / ts.static with autobatcher=True fails instead of using the largest batch that fit.

Why it happens only sometimes

Which error ends the probe depends on the allocator state at the step that crosses the limit:

  • In a fresh process, PyTorch's caching allocator fails once, frees its cache, and retries successfully. That leaves too little memory for cuequivariance's own cudaMallocAsync buffer, which is allocated outside PyTorch's allocator. cuequivariance then fails with the RuntimeError above, which isn't recognized.
  • Later in the same process, or for other structures, PyTorch's retry also fails, so torch.OutOfMemoryError is raised first. It is recognized and the probe stops cleanly.

So a given structure fails deterministically when it runs in a fresh process. That is the case for workflow managers that run each job in its own process (jobflow-remote batch mode, in our case).

Reproduction

A100-SXM4-40GB, MACE-MH-1, MACE-MH1, 12-atom NaO₂ cell:

import torch
import torch_sim as ts
from mace.calculators.foundation import download_mace_mp_checkpoint
from pymatgen.core import Lattice, Structure
from torch_sim.autobatching import determine_max_batch_size
from torch_sim.models.mace import MaceModel

structure = Structure(
    Lattice([[5.388803, -0.010233, 0.227281],
             [-0.00986, 5.374508, 0.013876],
             [0.225145, 0.014051, 5.410285]]),
    ["Na"] * 4 + ["O"] * 8,
    [[0.0, 0.5, 0.5], [0.5, 0.5, 0.0], [0.5, 0.0, 0.5], [0.0, 0.0, 0.0],
     [0.43045, 0.076054, 0.935304], [0.075358, 0.937325, 0.417139],
     [0.924244, 0.43641, 0.082993], [0.570596, 0.575467, 0.564034],
     [0.429404, 0.424533, 0.435966], [0.075756, 0.56359, 0.917007],
     [0.924642, 0.062675, 0.582861], [0.56955, 0.923946, 0.064696]],
)

model_path = download_mace_mp_checkpoint("mh-1")
model = MaceModel(model=model_path, enable_cueq=True, compute_stress=True)
state = ts.initialize_state([structure], model.device, model.dtype)
determine_max_batch_size(state, model)

Output (truncated):

Model Memory Estimation: Running forward pass on state with 6792 atoms and 566 systems.
Model Memory Estimation: Running forward pass on state with 10872 atoms and 906 systems.
[W CUDACachingAllocator.cpp:3933] memory allocation failed with OOM on device 0 while trying to allocate 7348420608 bytes (free: 6719864832, total: 42436526080).
RuntimeError: CUDA error encountered at: ... cudaMallocAsync(&sort_permutation, ...) ... Reason=cudaErrorMemoryAllocation:out of memory

Run again in the same process, the same step raises torch.OutOfMemoryError (two allocator warnings: 6.66 GB, then 7.29 GB free), and the probe returns 354 systems. Two other structures we tried (Al₂O₃ and Mg₂FeH₆, 10 and 9 atoms) also ended with torch.OutOfMemoryError. With enable_cueq=False, the probe ends cleanly at 86 systems.

Suggested fix

In determine_max_batch_size:

  1. Treat torch.OutOfMemoryError as out of memory by type, whatever its message.
  2. Add a generic "out of memory" to DEFAULT_OOM_ERROR_MESSAGES, to cover kernels that raise a plain RuntimeError with a CUDA allocation error, like cuequivariance's.
except Exception as exc:
    if isinstance(exc, torch.OutOfMemoryError) or any(
        msg in str(exc) for msg in oom_error_message
    ):
        ...

Another thing that might be interesting to patch is to skip the probe when the state has a single system. Nothing can be batched, and the probe fills the GPU for no benefit, which is the most common case when the autobatcher is turned on by default.

Environment
  • torch-sim-atomistic 0.6.2
  • mace-torch 0.3.16
  • cuequivariance / cuequivariance-torch / cuequivariance-ops-torch-cu12 0.12.0
  • torch 2.13.0+cu126

(issue message was generated by Claude Code)

Dominant language
Python
Stars
502
Forks
108
Avg merge
1d 1h
Merged PRs (30d)
12

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from TorchSim/torch-sim

All issues in TorchSim/torch-sim

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.