DistMLIP 머신러닝 원자간 포텐셜 분산 추론 플랫폼

초록

DistMLIP은 그래프 수준의 제로 중복 파티셔닝을 이용해 머신러닝 기반 원자간 포텐셜(MLIP)의 멀티‑GPU 추론을 효율적으로 수행하는 플랫폼이다. 기존의 공간 분할 방식과 달리 그래프 파티셔닝으로 노드·엣지 정보를 레이어마다 교환함으로써 장거리 상호작용을 갖는 GNN 기반 MLIP을 손쉽게 확장한다. CHGNet, MACE, TensorNet, eSEN 네 가지 최신 MLIP을 테스트했으며, 8 GPU 환경에서 최대 8배 가속 및 3.4배 더 큰 시스템을 시뮬레이션할 수 있음을 보였다.

상세 분석

DistMLIP은 MLIP 추론을 위한 분산 프레임워크로, 기존 LAMMPS 기반의 공간 파티셔닝이 갖는 ‘ghost atom’ 중복 계산을 완전히 제거한다. 핵심 아이디어는 원자들을 노드, 원자 간 거리와 화학종을 엣지로 하는 그래프를 구축한 뒤, 전체 그래프를 수직 벽을 기준으로 파티션하고 각 파티션에 필요한 1‑hop 및 2‑hop 이웃을 포함하는 확장 서브그래프(G′i)를 생성하는 것이다. 이렇게 하면 각 GPU는 자신의 파티션에 속한 노드와, 인접 파티션으로부터 전달받은 ‘border node’·‘border edge’ 정보를 이용해 완전한 메시지 패싱을 수행할 수 있다.

특히 DistMLIP은 두 종류의 그래프를 동시에 분산한다. 첫 번째는 원자 그래프이며, 두 번째는 세‑원자 상호작용을 표현하는 라인 그래프(또는 bond graph)이다. 라인 그래프는 원자 그래프의 1‑hop·2‑hop 이웃을 기반으로 재구성되며, 엣지 테이블을 이용해 효율적인 매핑과 재귀 탐색을 수행한다. 이 과정은 모델‑아그노스틱하게 설계돼, CHGNet와 같은 다중 레이어 GNN, MACE와 같은 고차원 메시지 패싱, TensorNet과 같은 경량 모델까지 동일한 파이프라인으로 분산할 수 있다.

데이터 전송은 ‘TO’, ‘FROM’, ‘PURE’ 버킷을 이용해 메모리 레이아웃을 사전 정의하고, 마커 배열을 통해 각 버킷의 시작·끝 인덱스를 빠르게 조회한다. 이렇게 하면 레이어마다 필요한 border node/edge 피처만 최소한의 통신으로 교환할 수 있어, 통신 오버헤드가 크게 감소한다. 실험에서는 기존 멀티‑GPU 구현 대비 2‑3배 정도의 통신 비용 절감과, 전체 실행 시간에서 최대 8배 가속을 달성했다.

또한 DistMLIP은 LAMMPS와 같은 외부 시뮬레이터에 종속되지 않는다. 플러그인 형태의 인터페이스를 제공해, 사용자는 기존 PyTorch 기반 MLIP 모델을 그대로 불러와 distmlip.run() 호출만으로 분산 시뮬레이션을 시작할 수 있다. 이는 연구자들이 새로운 모델을 별도 포팅 없이 바로 대규모 시뮬레이션에 적용할 수 있게 해준다.

성능 평가에서는 8 GPU(A100‑80GB) 환경에서 1백만 원자 규모의 시스템을 수 초 내에 한 스텝 계산했으며, 3.4배 더 큰 시스템(≈3.4M 원자)에서도 메모리 초과 없이 실행 가능했다. 특히 장거리 상호작용을 필요로 하는 MACE와 CHGNet에서 기존 공간 파티셔닝 방식이 겪는 메모리 폭증 문제를 완전히 회피했다.

요약하면, DistMLIP은 그래프‑레벨 파티셔닝, 제로 중복 데이터 전송, 고차원 라인 그래프 지원, 그리고 프레임워크 독립성을 결합해, 현재와 미래의 대규모 MLIP 기반 원자 시뮬레이션에 필요한 핵심 인프라를 제공한다.