[논문 리뷰] High-performance training and inference for deep equivariant interatomic potentials

Chuin Wei Tan, Marc Descoteaux외 12명|ArXiv.org|2025. 04. 22.

Topic Modeling인용 수 10

한 줄 요약

이 논문은 NequIP를 TorchInductor로 재구성하여 분산 학습을 가능하게 하고 Ahead-of-Time Inductor (AOTI)와 맞춤형 Triton 텐서-곱 커널을 도입하여 MD 추론에서 최대 18배의 속도 향상을 달성하고 SPICE 2 데이터에 대한 대규모 학습을 가능하게 한다.

ABSTRACT

Machine learning interatomic potentials, particularly those based on deep equivariant neural networks, have demonstrated state-of-the-art accuracy and computational efficiency in atomistic modeling tasks like molecular dynamics and high-throughput screening. The size of datasets and demands of downstream workflows are growing rapidly, making robust and scalable software essential. This work presents a major overhaul of the NequIP framework focusing on multi-node parallelism, computational performance, and extensibility. The redesigned framework supports distributed training on large datasets and removes barriers preventing full utilization of the PyTorch 2.0 compiler at train time. We demonstrate this acceleration in a case study by training Allegro models on the SPICE 2 dataset of organic molecular systems. For inference, we introduce the first end-to-end infrastructure that uses the PyTorch Ahead-of-Time Inductor compiler for machine learning interatomic potentials. Additionally, we implement a custom kernel for the Allegro model's most expensive operation, the tensor product. Together, these advancements speed up molecular dynamics calculations on system sizes of practical relevance by up to a factor of 18.

연구 동기 및 목표

• 대규모 데이터셋에서 MLIP의 학습 확장성과 성능 향상 목표를 달성한다.
• 전통적인 MD 코드와의 엔드-투-엔드 고성능 추론 통합을 가능하게 한다.
• 깊은 대칭 불변 interatomic 포텐셜에 대한 컴파일러 기반 가속화를 엔드-투-엔드에서 제공한다.
• 대규모 다채로운 데이터셋(SPICE 2)에서 실용적 이점 시연.
• 향후 MLIP 아키텍처(예: Allegro)를 위한 프레임워크 확장성 강화.

제안 방법

• 학습 시간 컴파일 및 추론 호환성을 위해 파생을 포함한 전체 모델을 단일 PyTorch FX 그래프로 추적한다.
• torch.compile를 통해 순전파 및 역전파에 대해 최적화 커널을 생성하기 위해 TorchInductor를 적용한다.
• 그래디언트 버킷화를 피하는 맞춤형 DDP 접근법으로 분산 데이터-병렬 학습을 구현하여 융합 기회를 극대화한다.
• 비Python 환경(예: LAMMPS)에서 사용하기 위해 컴파일된 PyTorch 모델을 네이티브 코드로 내보내는 Ahead-of-Time Inductor(AOTI)을 도입한다.
• Allegro 텐서 곱에 대해 PyTorch 구현의 비융합 부분을 대체하기 위해 맞춤형 융합 Triton 커널을 개발한다.
• 메모리 사용을 줄이기 위해 Wigner 3-j 수축 계수를 압축 희소 형식으로 표현하고 중간물을 구체화하지 않는다.

실험 결과

연구 질문

• RQ1 TorchInductor를 어떻게 활용하여 깊은 대칭 불변 interatomic 포텐셜의 학습 및 추론을 가속할 수 있는가?
• RQ2 AOTI를 추론에 통합하여 LAMMPS 같은 고성능 MD 코드와의 성능 향상을 달성하는 방법은 무엇인가?
• RQ3 Allegro에서 텐서 곱에 대한 맞춤형 융합 커널이 메모리 및 계산 병목을 얼마나 더 줄일 수 있는가?
• RQ4 대규모 MLIP 데이터셋(SPICE 2)에 대한 다중-GPU 분산 학습은 어떻게 확장되는가?
• RQ5 SPICE 2에서 Allegro 모델을 서로 다른 전하 상태 하위집합으로 학습할 때 정확도와 일반화에 어떤 영향이 나타나는가?

주요 결과

지표	소형 모델	중형 모델	대형 모델
장벽 MAE [meV]	22.75	15.42	11.37
장벽 RMSE [meV]	32.36	21.77	15.38

• torch.compile으로 학습 시 TorchScript 대비 2.4–5.0x 속도 증가를 보인다(테스트 GPU에서).
• 분산 학습은 128 랭크까지 양호하게 확장되며 256 랭크에서 합리적인 효율성(AMD 기준 40%, NVIDIA 기준 24%의 Baseline 대비)을 보인다.
• AOTI와 최적화된 텐서 곱 커널을 이용한 추론은 소분자 및 물 박스에서 TorchScript 대비 최대 18x 속도 증가를 달성한다.
• AOTI는 대규모 시스템에 대해 메모리 효율적이고 확장 가능한 추론을 가능하게 하여 GPU당 실 feasibile 시스템 크기를 확장한다(하드웨어에 따라 4320–5184 원자까지의 대형 모델 포함).
• 생체분자 시스템(DHFR, 셀룰로오스)에 대한 프런티어에서의 강한 스케일링은 프런티어와 펄름머터에서 최대 256 노드까지 우수한 성능을 보이며, AOTI+TP의 상당한 처리량 이점을 보여준다.
• SPICE 2 벤치마크와 비교할 때 세 Allegro 모델(소형/중형/대형) 모두 모델 크기가 커질수록 MAE/RMSE가 감소하는 방향으로 바람직한 경향의 토션 벤더 예측을 보인다.