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

Abierto Apto para principiantes
  • ¿Algo no está bien?
#627 5 comentarios 1 reacción 0 asignados Ver en GitHub

Los mantenedores suelen responder en 2 días

@lil-lon ya está trabajando en esto.

Desde el 31/10/2026.

  • #631 de @lil-lon — abierto

Evaluación

Dificultad
2/5
Tiempo estimado
1-3 horas
Aptitud para principiantes
70/100
Tipo de issue
Error
Claridad
Bien especificado
Estado de actividad
Activo
Stack tecnológico
python, pytorch
Área
backend, performance

Línea de trabajo

Localiza determine_max_batch_size y DEFAULT_OOM_ERROR_MESSAGES en el código de autobatching. Lee cómo se reconocen actualmente las excepciones de sondeo y, después, revisa las pruebas existentes para el manejo de falta de memoria. Añade cobertura para torch.OutOfMemoryError y un error de ejecución simple que contenga “out of memory”; el trabajo estará terminado cuando ambos detengan el sondeo correctamente, mientras que los errores no relacionados sigan propagándose.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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)

Lenguaje dominante
Python
Estrellas
502
Forks
108
Merge medio
1 d 1 h
PR fusionados (30 d)
12

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de TorchSim/torch-sim

Todos los issues de TorchSim/torch-sim

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.