Speeding Up MACE: Low-Precision Tricks for Equivarient Force Fields
Created by
Haebom
저자
Alexandre Benoit
개요
MACE와 같은 SO(3) 등변 모델의 계산 비용 절감을 위해, 저정밀 산술과 GPU 최적화 커널을 사용하여 정확도를 유지하면서 모델을 더 저렴하고 빠르게 만드는 방법을 연구합니다. MACE의 성능 병목 현상을 파악하고, 다양한 정밀도 설정을 평가하여 추론, 시뮬레이션 및 훈련 과정을 분석합니다.
시사점, 한계점
•
cuEquivariance 백엔드는 추론 지연 시간을 약 3배 감소시킵니다.
•
선형 레이어만 BF16/FP16으로 캐스팅하면 모델 속도가 약 4배 더 빨라집니다.
•
NVT/NPT MD에서 에너지 및 열역학적 관측값은 run-to-run 변동 내에서 유지됩니다.
•
훈련 중 반정밀도 가중치는 힘 RMSE를 저하시킵니다.
•
e3nn과 cuEq 모듈을 명시적인 어댑터 없이 혼합하면 표현 불일치가 발생합니다.
•
융합된 등변 커널과 혼합 정밀도 추론은 다운스트림 MD에 거의 영향을 미치지 않으면서 최첨단 힘장(force field)을 크게 가속화할 수 있습니다.
•
실용적인 정책은 기본적으로 cuEquivariance와 FP32를 사용하고, 최대 처리량을 위해 선형 레이어에 BF16/FP16을 활성화하는 것입니다 (FP32 누적 유지). 훈련은 FP32로 유지됩니다.
