A Survey of Neural Network Variational Monte Carlo from a Computing Workload Characterization Perspective

이 논문은 네 가지 대표적인 신경망 변분 몬테카를로 (NNVMC) 모델의 실행 특성을 GPU 워크로드 관점에서 체계적으로 분석하여, 저강도 연산과 데이터 이동이 성능 병목이 됨을 규명하고 확장 가능한 NNVMC 시스템을 위한 알고리즘 - 하드웨어 공동 설계 방향을 제시합니다.

핵심

이 논문은 **"양자 화학을 계산하는 인공지능 (NNVMC) 이 왜 컴퓨터 그래픽 카드 (GPU) 에서 그렇게 느리고 메모리를 많이 먹는지"**를 분석한 연구입니다.

쉽게 말해, "아주 정교한 시뮬레이션을 돌리는 AI 가 왜 슈퍼컴퓨터를 쓰는데도 '막히는' 현상이 발생하는지" 그 원인을 찾아내고, 앞으로 어떻게 더 빠르게 만들 수 있을지 제안하는 보고서입니다.

이 복잡한 내용을 일상적인 비유로 설명해 드릴게요.

---

1. 이 연구가 해결하려는 문제: "거대한 주방의 혼란"

양자 화학 계산은 원자와 전자가 어떻게 움직이는지 시뮬레이션하는 일입니다. 기존 방법은 계산이 너무 복잡해서 한계가 있었지만, 최근 **인공지능 (NNVMC)**을 쓰면 더 정확하고 빠르게 계산할 수 있다고 합니다.

하지만 실제 실행해 보니 문제가 생겼습니다.

비유: 마치 거대한 요리를 하려고 고급 주방 (GPU) 을 샀는데, 요리사 (AI) 가 재료를 다듬는 데만 90% 의 시간을 쓰고, 실제 요리 (계산) 는 10% 만 하는 상황입니다.

요리사가 재료를 다듬는 것 (데이터 이동) 이 너무 많아서, 아무리 강력한 가스레인지 (계산 능력) 가 있어도 전체 속도가 느려지는 것입니다. 이 논문은 바로 **"왜 재료를 다듬는 시간이 이렇게 오래 걸리는지"**를 네 가지 다른 요리법 (AI 모델) 을 비교하며 분석했습니다.

2. 네 가지 요리법 (AI 모델) 비교

연구진은 네 가지 다른 AI 모델 (PauliNet, FermiNet, Psiformer, Orbformer) 을 비교했습니다. 각 모델은 양자 세계를 설명하는 방식이 조금씩 다릅니다.

• PauliNet & FermiNet (전통적인 요리법):
  • 이 방식은 재료를 다듬고 (데이터 이동), 요리하고, 다시 재료를 다듬고, 다시 요리하는 과정을 매우 반복합니다.
  • 결과: 계산 능력 (FLOP) 은 충분하지만, 재료를 옮기는 손길이 너무 빨라 주방이 혼잡해집니다. 메모리 대역폭 (데이터 이동 속도) 이 병목이 됩니다.
• Psiformer (현대적인 요리법):
  • 이 방식은 'Transformer'라는 기술을 써서 재료를 한 번에 많이 다듬고, 요리하는 과정에 집중합니다.
  • 결과: 계산량이 늘어나서 GPU 의 계산 능력을 더 잘 쓰지만, 그래도 재료를 옮기는 시간이 여전히 많습니다.
• Orbformer (최신 요리법):
  • 이 방식은 'FlashAttention'이라는 기술을 써서 재료를 더 효율적으로 다듬으려 합니다.
  • 결과: 계산량이 줄어든 대신, 재료를 옮기는 작업 (데이터 이동) 의 비중이 다시 커졌습니다. 결국 다시 메모리 이동 속도에 걸려서 느려집니다.

3. 핵심 발견: "계산 능력보다 '이동'이 문제다"

이 논문이 발견한 가장 중요한 사실은 다음과 같습니다.

"인공지능이 양자 화학을 계산할 때, 가장 큰 병목은 '계산이 안 되는 것'이 아니라 '데이터를 옮기는 것'이다."

• 일반적인 오해: "컴퓨터가 느리니까 더 강력한 칩을 만들면 되겠지?"
• 실제 현실: 칩이 아무리 빨라도, 재료를 주방 한구석에서 다른 구석으로 옮기는 시간이 너무 길어서 전체 속도가 느려집니다.
• 비유: F1 레이서 (GPU) 가 있어도, 길이 막혀서 (데이터 이동 병목) 시속 10km 로 달리는 것과 같습니다.

4. 해결책 제안: "주방을 어떻게 재설계할까?"

이 문제를 해결하기 위해 연구진은 다음과 같은 아이디어를 제안합니다.

1. 주방을 재배치하라 (PIM - 메모리 내 처리):
   • 재료를 옮기는 시간을 줄이기 위해, 재료가 있는 곳 (메모리) 에서 바로 다듬는 작업을 하도록 주방을 개조해야 합니다. 계산기 (GPU) 로 옮기지 않고, 냉장고 (메모리) 옆에서 바로 요리를 시작하는 것입니다.
2. 상황에 따라 주방을 바꿔라 (하이브리드 시스템):
   • 요리 단계마다 필요한 작업이 다릅니다. 재료를 많이 다듬을 때는 메모리 옆에서, 요리를 할 때는 강력한 가스레인지에서 하도록 작업장을 동적으로 나누는 것이 좋습니다.
3. 메모리 부족 문제 해결 (오프로딩):
   • 요리 재료가 너무 많아서 냉장고가 꽉 차면, **냉장고 옆에 있는 창고 (CPU 메모리/SSD)**에 재료를 잠시 옮겨두고 필요할 때 가져오는 시스템을 만들어야 합니다.

5. 결론: "단순한 속도 향상이 아닌, '흐름'의 최적화"

이 논문은 단순히 "더 빠른 컴퓨터를 만들어라"가 아니라, **"양자 화학 계산이라는 특수한 작업의 흐름 (Stage) 을 이해하고, 그에 맞춰 하드웨어와 알고리즘을 함께 설계해야 한다"**고 말합니다.

• 지금까지: "계산이 느리니까 계산기를 더 빠르게 만들자."
• 이제부터: "재료를 옮기는 시간이 너무 길다. 재료를 옮기지 않고 바로 처리할 수 있는 주방 구조로 바꾸자."

이 연구는 양자 화학을 연구하는 과학자들과 컴퓨터를 설계하는 엔지니어들이 함께 손잡고, 데이터가 흐르는 길을 막지 않는 새로운 주방 (하드웨어) 을 만들 수 있는 길잡이가 되어줍니다.

기술 요약

1. 문제 정의 (Problem)

• 배경: 전통적인 양자 화학 방법 (CCSD(T), FCI 등) 은 정확도와 계산 비용 간의 트레이드오프가 존재합니다. NNVMC 는 신경망을 파동함수 (wavefunction) 에 적용하여 이 문제를 해결할 수 있는 유망한 대안으로 부상했습니다.
• 현실적 한계: NNVMC 는 이론적으로 우수한 확장성을 가지지만, 실제 GPU 환경에서는 높은 실행 시간과 메모리 비용으로 인해 대규모 시스템에 적용하기 어렵습니다.
• 핵심 과제: 언어 (NLP) 나 비전 (Vision) 워크로드와 달리, NNVMC 는 물리 법칙에 기반한 고유한 실행 단계 (Markov Chain Monte Carlo 샘플링, 파동함수 구성, 미분/라플라시안 평가 등) 를 가집니다. 이로 인해 단순한 연산량 (FLOPs) 만으로는 성능을 예측하기 어렵고, 이질적인 커널 행동과 비선형적인 병목 현상이 발생합니다.

2. 방법론 (Methodology)

저자들은 네 가지 대표적인 NNVMC 모델 (PauliNet, FermiNet, Psiformer, Orbformer) 을 대상으로 통합 프로파일링 프로토콜을 적용하여 워크로드를 분석했습니다.

• 대상 모델 및 구현체:
  • DEEPQMC 코드베이스: PauliNet, FermiNet (행렬식 기반, JVP 기반 라플라시안 평가 사용).
  • ONEQMC 코드베이스: Psiformer, Orbformer (Transformer 기반, Hutchinson 스타일 라플라시안 평가 사용).
• 실험 환경: NVIDIA RTX A5000, A100, H200 GPU 사용. 다양한 분자 크기 (LiH, CH4, C2H6, C4H4) 에 대해 훈련 및 추론 워크로드를 측정.
• 분석 지표:
  • 엔드 - 투 - 엔드 (End-to-End): 실행 시간 및 메모리 사용량.
  • 커널 레벨: 연산 강도 (Arithmetic Intensity, AI), 루프라인 (Roofline) 모델 분석, 하드웨어 활용도 (SM 활용도, Tensor Core, L2 캐시 히트율 등).
  • 프로파일링 도구: NVIDIA Nsight Systems 및 Nsight Compute.

3. 주요 기여 (Key Contributions)

1. 워크로드 중심의 NNVMC 개관: DEEPQMC 와 ONEQMC 코드베이스에 구현된 네 가지 대표 모델을 포괄하는 체계적인 리뷰 제공.
2. 연산자 및 커널 레벨 특성화: 융합된 요소별 (elementwise) 연산과 데이터 이동 커널이 GEMM(일반 행렬 곱) 커널과 병행될 때에도 실행 시간을 지배하는 이유를 경험적 연산 강도와 루프라인 분석을 통해 규명.
3. 하드웨어 행동 보고 및 공동 설계 시사점: 스트리밍 멀티프로세서 (SM) 활용도, 메모리 처리량, 캐시 히트율 등 하드웨어 지표를 분석하여 확장 가능한 NNVMC 시스템을 위한 알고리즘 - 하드웨어 공동 설계 방향을 제시.

4. 주요 결과 (Key Results)

A. 실행 시간 및 메모리 특성

• 모델별 스케일링 차이: 실행 시간은 분자 크기에 따라 증가하지만, 그 기울기는 모델 아키텍처에 따라 다릅니다.
  • PauliNet/FermiNet: JVP(자코비안 - 벡터 곱) 기반의 라플라시안 평가로 인해 미세한 요소별 커널이 반복 실행되어 실행 시간이 급격히 증가합니다.
  • Psiformer/Orbformer: Hutchinson 스타일 라플라시안으로 재연산 오버헤드가 줄어 상대적으로 완만한 스케일링을 보이지만, 샘플링 단계에서의 GEMM/어텐션 연산 비중이 큽니다.
• 메모리 사용: 구현체별 메모리 할당 전략 (ONEQMC 의 동적 할당 등) 에 따라 메모리 사용 패턴이 크게 달라집니다.

B. 커널 구성 및 병목 현상

• 메모리 바운드 (Memory-Bound) 우세: 대부분의 NNVMC 워크로드는 **낮은 연산 강도 (Low Arithmetic Intensity)**를 가진 융합된 요소별 (fused elementwise) 연산과 데이터 이동 커널에 의해 지배됩니다. 이는 GPU 의 대역폭 한계에 부딪히게 만듭니다.
• 단계별 이질성:
  • Stage E (미분/라플라시안 평가): PauliNet/FermiNet 에서 이 단계는 이전 단계의 계산을 재실행 (Replay) 하여 수많은 미세 커널을 생성하며, 전체 실행 시간의 큰 비중을 차지합니다.
  • GEMM 의 역할: FermiNet 은 PauliNet 대비 GEMM 비중이 높지만, 여전히 메모리 트래픽과 커널 세분화 문제로 인해 전체 성능은 메모리 바운드 상태입니다.
  • Orbformer: FlashAttention 도입으로 GEMM 비중이 감소하고 메모리 바운드 특성이 더욱 강화되었습니다.

C. 하드웨어 활용도

• 낮은 활용도: 대부분의 커널이 루프라인 모델의 메모리 바운드 영역에 위치하며, 피크 연산 처리량 (Peak Throughput) 활용도는 낮습니다 (예: PauliNet 의 경우 약 26%).
• Tensor Core 활용: 일부 커널에서 높은 Tensor Core 활용도를 보이지만, 전체적으로는 평균 활용도가 낮아 (약 19%) 대역폭 효율성이 주요 제한 요인입니다.

5. 시사점 및 공동 설계 방향 (Significance & Co-design Directions)

이 연구는 NNVMC 가속화를 위해 단순한 커널 최적화보다는 위상 인식 (Phase-aware) 및 메모리 중심 (Memory-centric) 접근이 필요함을 강조하며 다음과 같은 방향을 제안합니다.

1. 메모리 바운드 커널을 위한 PIM (Processing-in-Memory):
   • 낮은 연산 강도를 가진 요소별/레이아웃 커널의 데이터 이동 오버헤드를 줄이기 위해 메모리 내부 연산 (PIM) 기술 도입이 효과적입니다.
2. 협력형 GPU-PIM 이종 시스템:
   • GEMM 중심 단계는 GPU 에서, 메모리 바운드 단계는 PIM 에서 처리하는 등 위상별 (Phase-dependent) 동적 할당 전략이 필요합니다.
3. 재구성 가능한 가속기 (Reconfigurable Acceleration):
   • 고정된 하드웨어 아키텍처는 모든 단계를 효율적으로 처리하지 못하므로, 연산자 혼합 (Operator Mix) 에 따라 계산/메모리 자원을 동적으로 재구성할 수 있는 아키텍처가 필요합니다.
4. 어텐션 외의 커널 최적화:
   • Orbformer 와 같은 모델에서 어텐션 커널만 최적화하는 것은 전체 성능 향상에 한계가 있으며, 요소별 연산 및 데이터 이동 패턴 최적화가 필수적입니다.
5. CPU/SSD 메모리 오프로딩:
   • 대규모 분자 시스템에서 GPU 메모리 부족 문제를 해결하기 위해, 접근 빈도가 낮은 텐서를 CPU DRAM 또는 SSD 로 비동기적으로 오프로딩하는 전략이 필요합니다.

결론

이 논문은 NNVMC 가 단순한 딥러닝 워크로드가 아니라 물리 법칙에 의해 제약받은 이질적인 워크로드임을 증명했습니다. 따라서 NNVMC 의 확장성을 높이기 위해서는 연산량 중심의 접근을 넘어, 메모리 대역폭 효율성과 실행 단계별 특성에 맞춘 하드웨어 - 알고리즘 공동 설계가 필수적임을 시사합니다.