Accelerating finite-element-based projector augmented-wave density functional theory calculations with scalable GPU-centric computational methods

본 논문은 혼합 정밀도 연산과 체비셰프 필터 부분공간 반복과 같은 알고리즘적 혁신을 활용하여 대규모 화학적으로 정확한 밀도 함수 이론 시뮬레이션에 대해 상당한 속도 향상과 엑사스케일 준비 성능을 달성하는 확장 가능한 GPU 중심 유한 요소 프로젝터 보강 파동 (PAW-FE) 방법을 제시한다.

핵심

복잡한 기계, 예를 들어 자동차 엔진이나 새로운 유형의 배터리가 어떻게 작동할지 예측하려고 상상해 보세요. 이를 정확하게 수행하려면 기계를 구성하는 물질 내부의 모든 단일 전자의 행동을 이해해야 합니다. 이것이 **밀도 범함수 이론 (DFT)**이라는 분야의 역할입니다. 이는 수십억 개의 전자가 동기화되어 움직이는 거대하고 정교한 무용장을 시뮬레이션하려는 것과 같습니다.

오랫동안 과학자들은 다음과 같은 문제에 직면해 왔습니다: 소수의 원자 그룹에 대한 이러한 춤을 시뮬레이션하는 것은 쉽지만, 작은 금속 나노입자나 꼬인 물질 시트와 같은 크고 복잡한 시스템을 시뮬레이션하려고 하면 컴퓨터가 압도당합니다. 이는 100 명을 위해 설계된 방법으로 10 만 명의 춤을 지휘하려는 것과 같습니다; 지시가 꼬이고 메모리가 가득 차며 시뮬레이션이 끝나는 데 영원히 걸립니다.

이 논문은 GPU(고급 비디오 게임과 AI 를 구동하는 동일한 칩) 를 사용하는 현대적이고 강력한 컴퓨터를 위해 특별히 설계된 이러한 시뮬레이션을 실행하는 새로운 초고속 방법을 소개합니다. 그들이 어떻게 했는지 간단한 개념으로 나누어 설명합니다:

1. 구식 방법 vs 새로운 지도

• 구식 방법 (평면파): 모든 평방 인치가 동일한 크기를 가진 거대하고 균일한 격자로 도시를 매핑하려고 상상해 보세요. 건물 하나의 벽돌과 같은 작은 세부 사항을 보고 싶다면, 도시 위의 빈 하늘조차도 전체 격자를 놀라울 정도로 정교하게 만들어야 합니다. 이는 막대한 양의 컴퓨터 성능을 낭비합니다. 이것이 현재 대부분의 소프트웨어가 작동하는 방식입니다.
• 새로운 방법 (유한 요소): 저자들은 "스마트 지도" 접근 방식을 사용합니다. 필요한 곳 (예: 붐비는 도시 중심부) 에만 확대하고 빈 곳 (예: 하늘) 에는 축소하는 지도를 상상해 보세요. 이를 유한 요소 (FE) 이산화라고 합니다. 이를 통해 전자가 흥미로운 일을 하는 정확한 위치에 컴퓨팅 성능을 집중시켜 막대한 시간과 메모리를 절약할 수 있습니다.

2. "PAW" 트릭: 마법 의상

수학을 더 쉽게 만들기 위해 **프로젝터 보강파 (PAW)**라는 방법을 사용합니다.

• 문제: 원자 중심 (원자핵) 근처의 전자는 격하게 떨리고 진동하여 계산하기 어렵습니다.
• 해결책: PAW 는 전자에게 "부드러운 의상"을 입히는 것과 같습니다. 대부분의 계산 동안 전자가 매끄럽고 다루기 쉽다고 가정하지만, 원자핵 근처의 세부 사항을 확인할 때만 전자의 진짜 격렬한 행동을 즉시 드러내주는 비밀 "마법 트릭"을 유지합니다. 이를 통해 정확도를 잃지 않으면서 훨씬 더 거친 (간단한) 지도를 사용할 수 있습니다.

3. GPU 속도 부스트: 조립 라인

저자들은 지도만 바꾼 것이 아니라 현대 GPU에 적합하도록 컴퓨터가 수학을 수행하는 방식을 변경했습니다.

• 병목 현상: 일반적으로 컴퓨터는 메모리와 프로세서 간에 데이터가 이동하는 것을 기다리는 데 많은 시간을 보냅니다.
• 해결책: 컴퓨터가 하나씩이 아니라 여러 계산을 동시에 수행할 수 있도록 (조립 라인처럼) 수학을 재설계했습니다. 또한 체비셰프 필터링이라는 교묘한 기술을 사용했는데, 이는 "중요한" 전자와 "중요하지 않은" 전자를 빠르게 분리하는 체와 같아 컴퓨터가 필요 없는 전자에 시간을 낭비하지 않도록 합니다.

4. "충분한" 단축키 (혼합 정밀도)

이것이 아마도 가장 창의적인 부분일 것입니다.

• 비유: 거대한 벽화를 그리고 있다고 상상해 보세요. 배경 하늘의 경우, 페인트를 미세한 정밀도로 섞을 필요가 없습니다; "충분한" 혼합으로 충분하며 훨씬 빠릅니다. 얼굴의 작은 세부 사항에만 극도의 정밀도가 필요합니다.
• 적용: 저자들은 일반적인 모양을 올바르게 얻기만 하면 되는 계산 부분에서는 낮은 정밀도의 수학(마치 눈금이 적은 자를 사용하는 것과 같음) 을 사용할 수 있음을 깨달았습니다. 이는 현대 칩에서 훨씬 빠릅니다. 최종적이고 중요한 단계에서만 "초정밀" 수학으로 전환합니다.
• 결과: 고정밀도와 저정밀도 수학을 혼합하고 데이터 전송과 계산을 겹쳐서 (동시에 두 가지 일을 수행) 시뮬레이션 속도를 이전보다 8 배에서 20 배까지 높였습니다.

5. 그들이 실제로 달성한 것

이 논문은 이러한 새로운 방법으로 다음과 같은 것을 달성했다고 주장합니다:

• 속도: 슈퍼컴퓨터에서 10,000 개에서 130,000 개의 전자를 가진 시스템을 실용적인 시간 (수분에서 수시간) 내에 시뮬레이션할 수 있습니다.
• 비교: 이 크기의 시스템에 대해 그들의 방법은 선도적인 표준 소프트웨어 (Quantum ESPRESSO) 보다 약 8 배 빠릅니다.
• 규모: 그들은 130,000 개의 전자를 포함하는 "꼬인 이층 (twisted bilayer)" 물질 (두 장의 원자 시트가 서로 꼬인 것) 의 시뮬레이션을 성공적으로 수행했습니다. 이는 표준 방법을 사용하여 이 수준의 정확도로 이전에 시뮬레이션이 불가능했던 규모입니다.

요약

간단히 말해, 저자들은 시뮬레이션에 필요한 새롭고 매우 효율적인 엔진을 구축했습니다. 필요한 곳에만 확대하는 "스마트 지도", 수학을 단순화하는 "마법 의상" 트릭, 그리고 비중요한 단계에서 낮은 정밀도를 사용하는 "빠른 앞당김" 모드를 결합했습니다. 그 결과, 이 도구는 과거에 걸렸던 시간의 일부로 현대 슈퍼컴퓨터에서 거대하고 복잡한 물질을 모델링할 수 있게 되어, 배터리, 전자제품 및 촉매를 위한 새로운 물질을 훨씬 더 빠르게 설계할 수 있는 문을 열었습니다.

기술 요약

"확장 가능한 GPU 중심 계산 방법을 통한 유한 요소 기반 프로젝터 보강 파동 밀도 범함수 이론 계산 가속화" 논문에 대한 상세한 기술적 요약은 다음과 같습니다.

1. 문제 제기

복잡한 물질 시스템 (예: 계면, 결함, 나노 클러스터, 비틀린 이종 구조) 의 정확한 모델링에는 $10^4$에서 $10^5$개의 전자를 포함하는 밀도 범함수 이론 (DFT) 시뮬레이션이 필요합니다. 그러나 기존 구현 방식은 다음과 같은 심각한 병목 현상에 직면해 있습니다.

• 평면파 (PW) 의 한계: 널리 사용되는 PW-PAW 코드 (예: VASP, Quantum ESPRESSO) 는 빠른 푸리에 변환 (FFT) 에 의존하는데, 이는 막대한 모든-대-모든 (all-to-all) 통신을 수반합니다. 이로 인해 대역폭에 제한을 받고 현대의 엑사스케일 GPU 아키텍처에서 확장성이 떨어집니다.
• 실공간 방법의 한계: 기존 실공간 방법들은 종종 노름 보존赝퍼텐셜 (ONCV) 을 사용하는데, 이는 특히 전이 금속의 경우 화학적 정확도를 달성하기 위해 많은 수의 기저 함수 (높은 자유도) 를 필요로 합니다. 이로 인해 메모리 사용량과 계산 비용이 높아집니다.
• 하드웨어 불일치: 현재 고유값 솔버는 현대 GPU 클러스터의 높은 연산 강도와 낮은 지연 시간 요구 사항에 적합하지 않은 경우가 많아, 혼합 정밀도 기능을 활용하거나 통신과 계산을 효과적으로 겹치지 못하게 합니다.

2. 방법론

저자들은 멀티노드 GPU 아키텍처에 최적화된 유한 요소 (FE) 이산화 프로젝터 보강 파동 (PAW) 공식인 PAW-FE를 제시합니다.

A. 수학적 공식화

• PAW-GHEP: 이 방법은 일반화 에르미트 고유값 문제 (GHEP): $H\tilde{\Psi} = S\tilde{\Psi}\Lambda$를 풉니다. 여기서 $H$는 해밀토니안이고 $S$는 PAW 중첩 행렬입니다.
• 공선 스핀 공식: 방정식은 자기 시스템을 처리하기 위해 공선 스핀 프레임워크 내에서 유도됩니다.
• 경계 조건: FE 이산화는 인위적인 주기성 아티팩트 없이 주기적, 반주기적 (슬랩), 비주기적 (나노 클러스터) 경계 조건을 자연스럽게 수용합니다.

B. 계산적 혁신

GPU 에서 PAW-GHEP 를 효율적으로 풀기 위해 저자들은 다음과 같은 주요 알고리즘 전략을 개발했습니다.

1. 잔차 기반 체비셰프 필터링 부분공간 반복 (R-ChFSI):

   • 표준 ChFSI 대신 잔차 기반 공식 ($R = H\tilde{\Psi} - S\tilde{\Psi}\Lambda$) 을 사용합니다.
   • 주요 장점: 이 공식은 부정확한 행렬 - 벡터 곱에 견딜 수 있어, 수렴성을 희생하지 않고 PAW 중첩 행렬 ($S^{-1}$) 에 대한 근사 역행렬과 낮은 정밀도 연산을 사용할 수 있게 합니다.

2. 근사 역행렬 중첩 행렬:

   • 희소 행렬 $S$를 명시적으로 역행렬로 구하는 대신, 질량 행렬의 대각 근사와 국소화된 보정에서 유도된 계산적으로 저렴한 근사 ($\tilde{S}^{-1}$) 를 사용합니다. 이는 값비싼 전역 역행렬 계산을 피합니다.

3. 다중 해상도 구적법:

   • 거친 FE 격자에서 원자 중심 적분 (의사 코어 밀도) 을 처리하기 위해 다중 해상도 구적법 체계가 사용됩니다. 이는 원자의 보강 구 (augmentation spheres) 내에서는 정교한 구적법 규칙을 사용하고 다른 곳에서는 더 거친 규칙을 사용하여 메쉬 정밀화 없이 정확성을 보장합니다.

4. 혼합 정밀도 및 저정밀도 통신:

   • 계산: 체비셰프 필터링 단계 (행렬 - 행렬 곱셈) 에 FP32/TF32 연산을 사용하고, 인접 노드 간 통신에는 BF16 을 사용합니다.
   • 견고성: R-ChFSI 알고리즘이 잔차에 의존하기 때문에, 이러한 정밀도 감소가 바닥 상태의 최종 이중 정밀도 정확도를 저하시키지 않습니다.

5. 계산 - 통신 겹침:

   • 체비셰프 필터링은 블록 단위로 수행됩니다. 파동 함수의 한 블록이 프로세서 간 통신 (MPI) 을 수행하는 동안, 다른 블록은 GPU 에서 계산 (GEMM 연산) 을 수행합니다. 이는 엑사스케일 확장에 중요한 통신 지연 시간을 숨깁니다.

6. 셀 단위 밀집 연산:

   • 전역 희소 행렬을 구성하는 대신, 연산을 셀 단위에서 밀집 행렬 - 행렬 곱셈으로 재구성합니다. 이는 연산 강도와 캐시 지역성을 극대화하여 GPU 실행에 이상적입니다.

3. 주요 기여

• PAW-FE 공식화: 일반적인 경계 조건을 지원하는 공선 스핀 공식 내의 새로운 FE 이산화 PAW 공식.
• R-ChFSI 고유값 솔버: PAW-GHEP 를 풀기 위한 잔차 기반 체비셰프 필터링 부분공간 반복의 확장으로, 근사 역행렬과 혼합 정밀도 사용을 가능하게 함.
• 확장 가능한 GPU 구현: 혼합 정밀도 연산 (FP32/TF32/BF16), 계산 - 통신 겹침, 셀 단위 밀집 선형 대수를 특징으로 하는 포괄적인 구현 전략.
• 다중 해상도 적분: 거친 메쉬에서 정확한 원자 중심 PAW 적분 평가를 가능하게 하는 구적법 체계.

4. 결과 및 성능

이 방법은 AMD, Intel, NVIDIA GPU 를 사용하여 리더십급 슈퍼컴퓨터 (OLCF Frontier, ALCF Aurora, ALCF Polaris) 에서 벤치마크되었습니다.

• 정확도: 분자 ($O_2$, $NO_2$) 와 결정 (BCC Cr) 에 대해 평면파 코드 (Abinit, Quantum ESPRESSO) 와 비교하여 검증되었습니다. 에너지와 힘의 오차는 화학적 정확도 요구 사항보다 수 차수 낮습니다 (에너지의 경우 원자당 $O(10^{-12})$ Ha, 힘의 경우 $O(10^{-6})$ Ha/bohr).
• CPU-GPU 속도 향상: CPU 만 실행된 경우 (Intel 및 AMD 아키텍처) 에 비해 GPU 에서 8 배에서 20 배의 속도 향상을 달성했습니다.
• 평면파 (QE) 와의 비교:
  • 약 10,000 개의 전자를 가진 시스템의 경우, PAW-FE 는 Quantum ESPRESSO 대비 최소 벽시계 시간을 8 배 단축합니다.
  • 더 큰 시스템 (>10,000 개 전자) 의 경우, PW 방법의 전역 통신에 비해 FE 기저 함수의 국소성으로 인해 속도 향상 폭이 더욱 커집니다.
• ONCV-FE 와의 비교:
  • PAW-FE 는 노름 보존赝퍼텐셜 (ONCV) 을 사용하는 FE 방법에 비해 약 6 배 적은 계산 자원을 필요로 하며, PAW 방법이 가능하게 한 감소된 자유도 덕분에 약 5 배 빠른 해결 시간을 달성합니다.
• 엑사스케일 확장성:
  • **130,000 개의 전자 (11,000 개의 원자)**로 구성된 비틀린 이층 WTe2 시스템에 대한 바닥 상태 계산을 성공적으로 시연했습니다.
  • 400 노드에서 SCF 반복당 약 2 분의 해결 시간을 달성하여 현실적인 대규모 시뮬레이션을 위한 이 방법의 타당성을 입증했습니다.

5. 의의

이 연구는 PAW-FE를 1 차 원리 시뮬레이션을 위한 엑사스케일 준비 완료 방법으로 확립합니다. PAW 방법의 높은 정확도와 실공간 유한 요소의 효율성 사이의 간극을 메워 평면파 방법의 통신 병목 현상을 극복합니다. 다양한 GPU 아키텍처에서 혼합 정밀도와 겹침 전략의 성공적인 배포는 $10^5$개의 전자를 가진 복잡한 물질 시스템을 일상적으로 시뮬레이션할 수 있는 경로를 보여주며, 이전에는 계산적으로 처리 불가능했던 촉매, 배터리 계면, 양자 물질 분야에서 새로운 발견을 가능하게 합니다.