🔬 physics

jQMC: A JAX-based ab initio quantum Monte Carlo package designed for GPU-accelerated computing

이 논문은 변분 및 격자 정규화 확산 몬테카를로 시뮬레이션 모두에서 기존의 포트란 구현체보다 상당한 성능 향상을 달나기 위해 GPU 가속과 자동 미분을 활용하는, Python 기반의 JAX 지원 ab initio 양자 몬테카를로 패키지인 jQMC를 소개한다.

원저자: Kousuke Nakano, Michele Casula

게시일 2026-07-16

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Kousuke Nakano, Michele Casula

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

우주가 원자핵 주변을 아주 활발하게 움직이는 벌떼처럼, 전자라고 불리는 작고 춤추는 입자들로 이루어져 있다고 상상해 보십시오. 물질이 어떻게 행동하는지, 약이 어떻게 작용하는지, 혹은 왜 어떤 것들은 전기를 통하게 하고 다른 것들은 통하지 않게 하는지를 이해하기 위해, 과학자들은 슈뢰딩거 방정식이라는 거대하고 뒤엉킨 수학 문제를 풀어야 합니다. 이것은 마치 수많은 벌들이 서로 부딪히고 벌집 벽에 부딪히는 와중에, 그 무리 속 모든 벌의 정확한 경로를 예측하려고 노력하는 것과 같습니다. 수십 년 동안 우리가 가졌던 최고의 도구들은 큰 무리를 다루기에는 너무 느리거나, 너무 많은 지름길(가정)을 택해야 했기에 때때로 목표를 놓치는 추측을 내놓기도 했습니다.

여기에 **양자 몬테카를로(Quantum Monte Carlo, QMC)**가 등장합니다. 이것을 단일 계산이 아니라, 수천 명의 '워커(walker)'(전자의 디지털 유령)들이 가상 세계를 돌아다니며 모든 가능한 배치를 탐색하기 위해 무작위로 발걸음을 옮기는 거대한 시뮬레이션 게임이라고 생각해 보십시오. 이 워커들이 도착한 위치를 평균함으로써, 과학자들은 시스템의 에너지와 거동에 대한 매우 정확한 그림을 얻을 수 있습니다. 이는 경기장에 있는 모든 사람을 일일이 측정하는 대신, 수천 명의 사람들이 관중석을 돌아다니며 보고하게 함으로써 경기장 전체 인원의 평균 키를 추측하는 것과 비슷합니다. 문제는 무엇일까요? 이 게임은 계산 비용이 엄청나게 많이 듭니다. 신뢰할 수 있는 답을 얻기 위해 슈퍼컴퓨터를 사용하여 오랫동안 실행해야 하며, 오랫동안 이 게임을 실행하는 소프트웨어들은 더 오래되고 느린 유형의 컴퓨터에 맞춰 제작되어 왔습니다.

이것이 바로 나카노 코스케(Kousuke Nakano)와 미켈레 카술라(Michele Casula)가 소개한 새로운 소프트웨어인 jQMC로 이어집니다. 이 논문은 jQMC를 오늘날 가장 강력한 슈퍼컴퓨터에서 발견되는 최신형 고성능 그래픽 카드(GPU)에서 실행되도록 특별히 설계된 현대적인 파이썬(Python) 기반 엔진으로 제시합니다. 저자들은 이 도구를 통해, 기존의 골드 스탠다드 소프트웨어가 (오래된 프로그래밍 언어인 포트란(Fortran)으로 작성되었고 이러한 새로운 칩에 맞게 설계되지 않았음에도 불구하고) 이 전자 시뮬레이션을 훨씬 더 빠르고 효율적으로 실행할 수 있는지 확인하고자 했습니다.

연구팀은 jQMC가 현대적인 GPU와 결합했을 때 정말로 속도의 귀재라는 것을 발견했습니다. 시뮬레이션의 '변분 몬테카를로(Variational Monte Carlo, VMC)' 단계—즉, 워커들이 요령을 익히는 연습 단계—에서, GPU를 사용하는 jQMC는 CPU에서 실행되는 전통적인 소프트웨어보다 최대 10배 더 빨랐으며, GPU에서 실행되도록 시도된 전통적인 소프트웨어보다도 더 빨랐습니다. 더 복잡한 '확산 몬테카를로(Diffusion Monte Carlo, LRDMC)' 단계—즉, 최종적인 고난도 단계—에서도 jQMC는 GPU에서 여전히 약 3배 더 빨랐습니다. 연구진은 또한 원자를 밀고 당기는 '힘'(분자가 어떻게 움직이고 반응하는지 이해하는 데 필수적인 요소)을 j-QMC가 얼마나 잘 계산하는지도 테스트했습니다. 그들은 에탄올이라는 분자에 대해 jQMC의 힘 계산이 골드 스탠다드인 '커플드 클러스터(coupled-cluster)' 이론과 약 1.5 kcal/mol/Å 이내로 매우 정확하다는 것을 발견했습니다. 그러나 다른 분자인 말로날데하이드(malonaldehyde)의 경우 결과가 골드 스탠다드와 상당히 달랐는데, 이는 이 강력한 새로운 도구를 사용하더라도 여전히 복잡한 화학적 퍼즐들이 놀라움을 간직하고 있음을 시사합니다.

요약하자면, jQMC는 양자 화학을 현대화하는 데 성공적인 실험입니다. 이는 미래의 하드웨어(GPU)에 맞춰 게임의 규칙을 다시 씀으로써, 과학자들이 이전보다 훨씬 더 크고 복잡한 시스템을 연구할 수 있도록 시뮬레이션을 훨씬 빠르게 실행할 수 있음을 증명합니다. 비록 이것이 모든 화학적 미스터리를 해결하지는 못할지라도, 우리가 이미 알고 있는 미스터리들을 해결할 수 있는 훨씬 더 빠르고 유연한 방법을 제공합니다.

jQMC 기술 요약: GPU 가속 컴퓨팅을 위해 설계된 JAX 기반 ab initio 양자 몬테카를로 패키지

문제 정의
다체 슈뢰딩거 방정식을 푸는 것은 계산 화학 및 물리학의 핵심적인 과제로 남아 있습니다. 밀도 범함수 이론(DFT)이 널리 사용되고 있지만, 근사적인 교환-상관 범함수에 대한 의존성은 강하게 상관된 계(strongly correlated systems)에서의 정확도를 제한합니다. Ab initio 양자 몬테카로(QMC) 방법, 예를 들어 변분 몬테카로(VMC) 및 확산 몬테카로(DMC)는 유리한 스케일링을 갖는 확률론적 대안을 제공하지만, 계산 집약적입니다. 고성능 컴퓨팅(HPC) 시설이 GPU 중심의 아키텍처로 전환됨에 따라, Fortran이나 C로 작성되어 CPU 중심 설계에 의존하는 기존의 레거시 QMC 코드들은 상당한 어려움에 직면해 있습니다. 이러한 코드들을 GPU 병렬성을 활용하도록 리팩토링하는 것은 어렵고, 기존의 GPU 가속 QMC 코드들은 현대적인 프레임워크의 유연성이나 고정밀 ab initio 시뮬레이션에 필요한 특정 알고리즘 기능을 결여하고 있는 경우가 많습니다. 또한, 자동 미분을 활용할 수 있는 효율적인 최적화 기법과 표준화된 파동 함수 형식이 필요합니다.

방법론
저자들은 JAX 라이브러리를 기반으로 구축된, Python 기반의 ab initio QMC 시뮬레이션용 계산 패키지인 jQMC를 제시합니다. JAX는 Just-In-Time (JIT) 컴파일, 자동 미분(AD), 그리고 CPU와 GPU 백엔드 간의 원활한 이식성을 가능하게 합니다.

알고리즘: jQMC는 두 가지 바닥 상태 QMC 알고리즘을 구현합니다:
1. 변분 몬테카로 (VMC): 전자 구성을 샘플링하기 위해 메트로폴리스-헤이스팅스(Metropolis-Hastings) 알고리즘을 사용합니다. 노드 표면 근처의 특이점을 처리하기 위해 Attaccalite와 Sorella (AS) 재가중치(reweighting) 기법을 채택합니다.
2. 격자 정규화 확산 몬테카로 (LRDMC): 수치적 안정성을 개선하기 위해 해밀토니안을 격자 상에 이산화하는 Green's Function Monte Carlo (GFMC)의 변형입니다. jQMC는 병렬 계산에서의 부하 불균형을 완화하기 위해 전통적인 방식과 부하 균형(load-balanced) 브랜칭 기법을 모두 구현합니다.
파동 함수 안사츠 (Wave Function Ansatz): 시도 파동 함수는 Jastrow 인자와 반대칭 부분의 곱으로 구성됩니다. 반대칭 성분은 다음을 지원합니다:
- Antisymmetrized Geminal Power (AGP): 스핀 단일항(spin-singlet) 및 스핀 삼중항(spin-triplet) 쌍을 포함하는 공명 가전자 결합(RVB) 유형의 파동 함수입니다.
- 제한된 분자 궤도 함수를 가진 AGP (AGPn): 페어링 행렬의 랭크를 제한하는 변형입니다.
- 단일 슬레이터 행렬식 (Single Slater Determinant): AGP의 특수한 경우입니다.
  파동 함수는 TREX-IO 라이브러리를 통해 외부 Hartree-Fock (HF) 또는 DFT 계산으로부터 초기화될 수 있습니다.
최적화: 본 패키지는 변분 파라미터를 최적화하기 위해 Stochastic Reconfiguration (SR) 및 **Linear Method (LM)**를 구현합니다. 이 방법들은 파라미터 공간의 메트릭을 활용하여 파동 함수의 진폭과 노드 표면을 동시에 업데이트합니다.
미분 및 힘: JAX의 자동 미분을 활용하여 jQMC는 에너지 미분값을 효율적으로 계산합니다. 본 패키지는 VMC(Space-Warp Coordinate Transformation 사용)와 LRDMC(Reynolds 근사 및 Pathak-Wagner 정규화 사용) 모두에 대한 원자력(atomic force) 계산을 구현합니다.
아키텍처: 코드는 데이터 구조(JAX 호환 dataclass로 구현)와 알고리즘을 분리합니다. vmap을 사용하여 워커(walker)에 대한 벡터화를 수행하며(QMC의 최상위 병렬성), MPI (mpi4py) 및 JAX 샤딩을 통한 멀티 GPU 스케일링을 지원합니다. 또한 정확도를 희생하지 않으면서 계산을 가속화하기 위해 혼합 정밀도 산술(원자 궤도 평가를 위한 FP32, 결정식/제미니얼을 위한 FP64)을 지원합니다.

주요 기여

JAX 기반 구현: jQMC는 알고리즘 미분과 GPU 가속을 위해 JAX를 완전히 활용하는 최초의 QMC 패키지로, JIT 컴파일이 가능하며 CPU와 GPU 아키텍처 간에 이식 가능한 코드를 제공합니다.
고급 파동 함수 지원: 단일 슬레이터 행렬식을 일반화하고 노드 표면의 최적화를 허용하는 Jastrow-Antisymmetrized Geminal Power (JAGP) 안사츠의 견고한 구현을 제공합니다.
효율적인 벡터화: 최상위 수준에서 워커에 대해 벡터화함으로써, jQMC는 높은 intra-GPU 효율성을 달성하고 전통적인 CPU 병렬화 코드보다 GPU 자원을 더 효과적으로 포화시킵니다.
표준화: TREX-IO와의 통합은 다른 전자 구조 패키지(예: PySCF, CP2K)와의 호환성을 보장합니다.
힘 평가: 자동 미분과 특정 정규화 기법(Pathak-Wagner)을 사용한 원자력 구현은 QMC 프레임워크 내에서의 구조 최적화 및 포논(phonon) 계산을 가능하게 합니다.

결과

성능 벤치마크:
- GPU vs. CPU: NVIDIA GPU(A100 및 H100)에서 jQMC는 자체 CPU 구현 대비 2개 차수(100배) 이상의 속도 향상을 보여줍니다.
- TurboRVB와의 비교: 동일한 알고리즘과 파동 함수를 구현하는 Fortran90 코드인 TurboRVB와 비교했을 때, GPU 상의 jQMC는 현저히 빠릅니다. VMC의 경우 실행 시간(wall-time) 이득은 1개 차수(10배)에 달합니다. LRDMC의 경우 이득은 다소 완만하지만 여로 유의미합니다(테스트된 시스템 크기에서 약 3배 빠름).
- 확장성: Miyabi(도쿄 대학) 및 Leonardo(CINECA) 슈퍼컴퓨터에서의 약한 확장성(weak-scaling) 벤치마크 결과, jQMC는 VMC의 경우 1024개 GPU까지, LRDMC(부하 균형 알고리즘 사용 시)의 경우 512개 GPU까지 거의 이상적인 효율성(효율 $\approx$ 1.0)을 유지합니다.
정확도 검증:
- 재현성: 메탄, 물, 메탄-물 이량체에 대한 jQMC의 총 에너지 및 결합 에너지 결과는 통계적 오차 범위 내에서 다른 주요 QMC 코드들(CASINO, QMCPACK, CMQMC, TurboRVB)과 일치하며, 이는 DLA(Determinant Locality Approximation) 및 LRDMC 알고리즘의 구현을 검증합니다.
- 원자력: 논문은 에탄올과 말로날데히드에 대한 원자력을 all-electron CCSD(T) 계산과 비교하여 평가합니다.
  - 에탄올의 경우, LRDMC 힘은 CCSD(T) 대비 $\sim$ 1.5 kcal/mol/Å의 평균 절대 오차(MAE)를 달성하며, 이는 기존 문헌과 일치합니다.
  - 말로날데히드의 경우, LRDMC와 CCSD(T)의 힘 사이에서 $\sim$ 5 kcal/mol/Å(MAE)의 상당한 차이가 관찰되었습니다. 저자들은 노드 표면 최적화 후에도 이 차이가 지속된다는 점을 지적하며, 이것이 단순히 노드 표면 근사에 의한 문제가 아닐 수 있음을 시사하며 추가 조사가 필요하다고 언급했습니다.

의의
본 논문은 jQMC를 고정밀 ab initio QMC 방법과 현대적인 GPU 가속 HPC 아키텍처 사이의 간극을 메우는 현대적이고 다재다능한 플랫폼으로 자리매김합니다. JAX를 활용함으로써, jQ-MC는 새로운 알고리즘의 빠른 프로토타이핑을 위한 유연한 환경을 제공하는 동시에 생산 수준의 성능을 전달합니다. 저자들은 jQMC가 기존의 Fortran 기반 구현보다 GPU 하드웨어에서 더 빠르면서도 확장 및 유지보수가 용이한 도구를 제공함으로써, 레거시 QMC 코드를 GPU로 이식하는 문제를 성공적으로 해결했다고 주장합니다. 힘과 에너지에 대한 성공적인 검증은 jQMC가 특히 강한 상관 효과가 중요한 분자 및 확장 계(extended systems)의 생산적 시뮬레이션을 수행할 준비가 되었음을 입증합니다.

유사한 논문