[논문 리뷰] Highly-scalable, physics-informed GANs for learning solutions of stochastic PDEs

Liu Yang, Sean Treichler외 9명|arXiv (Cornell University)|2019. 10. 29.

Model Reduction and Neural Networks참고 문헌 29인용 수 16

한 줄 요약

이 논문은 한포드 사이트의 지하류 모델링에서 고차원 스토케스틱 PDE를 해결하기 위한 매우 확장 가능한 물리 기반 GAN 프레임워크를 제안한다. 도메인 분할, 다중 GPU 병렬 처리 및 물리 제약 학습을 활용하여, 27,500장의 볼타 GPU에서 93.1%의 스케일링 효율을 달성하였으며, 반정밀도 성능으로 1,207 PF/s를 지속적으로 달성하였다. 이는 스토케스틱 차원이 최대 1,000까지 가능한 혼합 SPDE 문제를 위한 첫 번째 엑사스케일 대응 GAN이다.

ABSTRACT

Uncertainty quantification for forward and inverse problems is a central challenge across physical and biomedical disciplines. We address this challenge for the problem of modeling subsurface flow at the Hanford Site by combining stochastic computational models with observational data using physics-informed GAN models. The geographic extent, spatial heterogeneity, and multiple correlation length scales of the Hanford Site require training a computationally intensive GAN model to thousands of dimensions. We develop a hierarchical scheme for exploiting domain parallelism, map discriminators and generators to multiple GPUs, and employ efficient communication schemes to ensure training stability and convergence. We developed a highly optimized implementation of this scheme that scales to 27,500 NVIDIA Volta GPUs and 4584 nodes on the Summit supercomputer with a 93.1% scaling efficiency, achieving peak and sustained half-precision rates of 1228 PF/s and 1207 PF/s.

연구 동기 및 목표

• 한포드 사이트의 복잡한 고차원 문제이자 광범위한 오염과 데이터 부족을 겪는 지하류 모델링에서의 불확실성 정량화를 해결한다.
• 공간적 이질성과 다중 척도 상관 길이로 인한 스토케스틱 PDE(SPDE) 해결 시 차원의 극복 문제를 해결한다.
• 스토케스틱 차원이 최대 1,000까지 가능한 확장 가능하고 물리 기반의 GAN(PI-GAN) 아키텍처를 개발한다.
• 도메인 분할과 최적화된 통신을 사용하여 대규모 HPC 시스템에서 GAN의 효율적인 분산 학습을 가능하게 한다.
• 1,000~10,000개의 우회원에서 수집된 관측 데이터를 통합된 혼합 문제 설정에 통합하여 압력 헤드와 수리 전도도 필드를 스토케스틱 과정으로 추론한다.

제안 방법

• 단일 생성자와 수백 개의 판별기를 다수의 GPU에 맵핑하는 계층적 도메인 분할 기반 학습 기반의 확장 가능한 학습을 가능하게 하는 기법을 제안한다.
• 스토케스틱 PDE 제약 조건을 직접 GAN 학습 과정에 통합한 물리 기반 손실 함수를 구현하여 생성된 해의 물리적 일관성을 보장한다.
• 고차원 스토케스틱 과정에서의 분포 정합성 향상을 위해 기울기 페널티를 적용한 WGAN-GP를 사용하여 학습을 안정화시킨다.
• 비동기 및 겹침 스케줄링을 사용한 통신 최적화 파이프라인을 설계하여 유휴 시간을 최소화하고 GPU 활용도를 극대화한다.
• 볼타 GPU를 대상으로 텐서플로우에서 생성자 및 판별기 아키텍처를 최적화하여 노드당 최대 59.3 TF/s 및 38.6 TF/s 성능을 달성한다.
• 동적 로드 밸런싱을 적용한 프리런 학습 스케줄링을 도입하여 4,584개의 Summit 노드에서 높은 병렬 효율을 유지한다.

실험 결과

연구 질문

• RQ1 1,000개 이상의 스토케스틱 차원을 가진 지하류에서 유래한 고차원 스토케스틱 PDE를 효과적으로 모델링할 수 있는 물리 기반 GAN 프레임워크는 가능한가?
• RQ2 단일 글로벌 생성자와 다수의 로컬 판별기를 가진 GAN에 대해 도메인 분할을 어떻게 효과적으로 적용하여 엑사스케일 학습을 가능하게 할 수 있는가?
• RQ3 반정밀도 산술을 사용하여 27,500개의 GPU에서 PI-GAN을 학습할 경우 도달할 수 있는 스케일링 효율성과 성능 수준은 어느 정도인가?
• RQ4 수천 개의 우회원에서 수집된 관측 데이터의 통합은 유도된 수리 전도도 필드와 압력 헤드 분포의 정확성과 신뢰성에 어떤 영향을 미치는가?
• RQ5 대규모 분산 환경에서 정밀도 감소 효과와 불균형한 GPU 워크로드가 존재하더라도 제안된 GAN 아키텍처는 학습 안정성과 수렴성을 유지할 수 있는가?

주요 결과

• 제안된 PI-GAN 프레임워크는 4,584개의 Summit 노드에서 93.1%의 스케일링 효율을 달성하여 초대규모에서 GAN 학습의 거의 최적의 강력 스케일링을 입증하였다.
• 27,500장의 볼타 GPU에서 반정밀도 성능으로 1,207 PF/s의 지속 성능과 최대 1,228 PF/s를 달성하였으며, 이는 스토케스틱 차원이 최대 1,000까지 가능한 SPDE 문제를 위한 첫 번째 엑사스케일 대응 GAN이다.
• 이 방법은 이전의 GAN 기반 SPDE 솔버보다 훨씬 높은 스토케스틱 차원(최대 1,000)을 가진 혼합 SPDE 문제를 성공적으로 해결하였다.
• 프리런 학습 스케줄링은 가장 높은 지속 성능(1,207 PF/s)과 병렬 효율(93.1%)을 달성하여 순차적 및 겹침 스케줄링을 모두 능가하였다.
• 제한된 센서 데이터에서 분포를 학습하는 방식을 통해 밀리온 분의 일 수준의 상관 길이를 가진 고차원 다중 척도 스토케스틱 과정을 효과적으로 처리하였다. 브루트 포스 샘플링이 아닌 방식이다.
• 정밀도 감소 및 워크로드 불균형 영향으로 특정 GPU 수(예: 10,500개)에서 약간의 효율성 저하가 관측되었지만, 테스트된 모든 구성에서 전체적으로 스케일링이 강인하게 유지되었다.