[논문 리뷰] Model-Parallel Fourier Neural Operators as Learned Surrogates for Large-Scale Parametric PDEs

Thomas J. Grady, Rishi Khan외 6명|arXiv (Cornell University)|2022. 04. 04.

Seismic Imaging and Inversion Techniques인용 수 9

한 줄 요약

이 논문은 도메인 분할을 사용하여 모델 병렬 FNO(Fourier Neural Operator) 프레임워크를 제안하여 FNO를 대규모 3D 및 4D 파arametric PDE에 확장함으로써, 26억 개 이상의 변수를 가진 문제에서 추론과 학습을 가능하게 한다. Perlmutter에서 최대 512개의 A100 GPU를 사용한 약한 스케일링을 달성하여, 지금까지 가장 큰 FNO 추론 및 기울기 계산을 실현한다.

ABSTRACT

Fourier neural operators (FNOs) are a recently introduced neural network architecture for learning solution operators of partial differential equations (PDEs), which have been shown to perform significantly better than comparable deep learning approaches. Once trained, FNOs can achieve speed-ups of multiple orders of magnitude over conventional numerical PDE solvers. However, due to the high dimensionality of their input data and network weights, FNOs have so far only been applied to two-dimensional or small three-dimensional problems. To remove this limited problem-size barrier, we propose a model-parallel version of FNOs based on domain-decomposition of both the input data and network weights. We demonstrate that our model-parallel FNO is able to predict time-varying PDE solutions of over 2.6 billion variables on Perlmutter using up to 512 A100 GPUs and show an example of training a distributed FNO on the Azure cloud for simulating multiphase CO$_2$ dynamics in the Earth's subsurface.

연구 동기 및 목표

• 고차원 입력과 가중치 제약으로 인해 FNO의 확장성에 한계가 있는 문제를 해결하기 위해, 대규모 3D 및 4D 파arametric PDE를 다룰 수 있도록 FNO의 확장성 문제를 해결한다.
• 분산 GPU 자원을 활용하여 26억 개 이상의 변수를 가진 문제에서 FNO의 효율적 추론 및 학습을 가능하게 한다.
• FNO의 인덕티브 바이어스와 성능을 유지하면서, 도메인 분할 기반의 모델 병렬 FNO 프레임워크를 개발하여 다수의 GPU에 걸쳐 확장 가능하도록 한다.
• 지하 CO2 누출 시뮬레이션 및 역문제와 같은 복잡한 실제 응용 분야에 대해 분산 FNO의 가능성을 입증한다.
• 기존의 수치적 해법으로는 구현이 어려웠던 대규모 최적화 및 불확실성 정량화 작업을 위한 빠르고 미분 가능한 PDE 해법을 제공한다.

제안 방법

• 다수의 GPU에 걸쳐 입력 데이터와 네트워크 가중치를 분할하는 도메인 분할 전략을 구현하여 모델 병렬성을 실현한다.
• FNO의 선형 대수적 공식화를 활용하여 파artition 간 수학적 일관성을 유지하는 분산 스펙트럴 컨볼루션 연산을 유도한다.
• 전방 및 역방향 전파 중 효율적인 GPU 간 통신을 위해 CUDA-인식 MPI를 활용하여 통신 오버헤드를 최소화한다.
• PyTorch의 DistDL 라이브러리를 사용하여 기존 FNO 아키텍처(4개의 스펙트럴 블록, 임bedding 차원 20)를 분산 환경에 적합하게 변형한다.
• 공간 차원을 최대 512×512×512로, 시간 차원을 최대 10,240개 타임스텝으로 확장하여 Perlmutter에서 약한 스케일링 실험을 수행한다.
• 자동 미분를 활용한 분산 학습 및 추론을 수행하여 기울기 계산을 지원하고, 역문제 워크플로우를 구현한다.

실험 결과

연구 질문

• RQ1 모델 병렬성을 사용하여 64³ 공간 한계를 초월한 대규모 3D 및 4D 파arametric PDE에 FNO를 확장할 수 있는가?
• RQ2 제안된 도메인 분할 기반 FNO가 최대 512개의 A100 GPU에서 약한 스케일링 실험에서 어떻게 성능을 발휘하는가?
• RQ3 기존의 해법 대비 분산 FNO가 대규모 PDE의 추론 및 기울기 계산을 얼마나 빠르게 처리할 수 있는가?
• RQ4 모델 병렬 FNO 프레임워크가 고차원 출력을 가진 실제 지하 CO2 유동 시뮬레이션에 효과적으로 적용될 수 있는가?
• RQ5 현대의 HPC 시스템인 Perlmutter에서 분산 FNO의 문제 크기 및 확장성에 대한 성능 한계는 무엇인가?

주요 결과

• 모델 병렬 FNO는 Perlmutter에서 최대 512개의 A100 GPU까지 약한 스케일링을 달성하여, 최대 512×512×512×20 공간-시간 차원을 가진 문제에 대한 확장성을 입증한다.
• 이 프레임워크는 26억 개 이상의 변수를 가진 문제에서 추론 및 기울기 계산을 가능하게 하여, 지금까지 가장 큰 FNO 응용 사례로 기록된다.
• FNO의 전방향 통신 특성으로 인해 약간의 비완전한 스케일링이 발생하지만, 테스트된 모든 스케일에서 뛰어난 성능을 유지한다.
• 해당 방법은 Azure 클라우드에서 지하에서의 다상태 CO2 역학을 시뮬레이션하기 위해 분산 FNO를 성공적으로 학습시켜 실제 적용 가능성에 대한 검증을 완료했다.
• 역문제 및 최적화 작업에 적합한 빠르고 미분 가능한 PDE 해법을 제공하며, 기존의 수치적 해법보다 추론 시간이 수 개의 지수 차수만큼 더 빠르다.
• 이 작업은 분산 연산자 학습을 통해 현실적인 부피 데이터 기반의 복합한 역문제 및 통계 문제를 해결하는 기반을 마련한다.