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

Đang mở Phù hợp với người mới
  • Có gì không ổn?
#627 4 bình luận 1 reaction 0 người được giao Xem trên GitHub

Maintainer thường phản hồi trong vòng 2 ngày

@lil-lon đang làm issue này rồi.

Từ ngày 31/10/2026.

  • #631 của @lil-lon — đang mở

Đánh giá

Độ khó
2/5
Thời gian dự kiến
1-3 giờ
Mức phù hợp với người mới
70/100
Loại issue
Lỗi
Độ rõ ràng
Đặc tả rõ ràng
Mức độ hoạt động
Sôi nổi
Công nghệ
python, pytorch
Lĩnh vực
backend, performance

Hướng nghiên cứu

Tìm determine_max_batch_size và DEFAULT_OOM_ERROR_MESSAGES trong mã autobatching. Đọc cách các ngoại lệ probe hiện được nhận diện, sau đó kiểm tra các bài kiểm tra hiện có về xử lý hết bộ nhớ. Bổ sung kiểm thử cho torch.OutOfMemoryError và một lỗi runtime thông thường có chứa “out of memory”; coi như hoàn tất khi cả hai đều dừng probe một cách bình thường, còn các lỗi không liên quan vẫn được truyền lên.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

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)

Ngôn ngữ chính
Python
Star
502
Fork
108
Merge trung bình
1 ngày 1 giờ
Pull request đã merge (30 ngày)
12

Chuẩn bị môi trường

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của TorchSim/torch-sim

Tất cả issue của TorchSim/torch-sim

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.