[논문 리뷰] SciAI4Industry -- Solving PDEs for industry-scale problems with deep learning

Philipp Witte, Russell J. Hewett외 3명|arXiv (Cornell University)|2022. 11. 23.

Reservoir Engineering and Simulation Methods인용 수 4

한 줄 요약

이 논문은 과학적 딥러닝을 위한 대규모 학습 데이터를 클라우드에서 생성하기 위한 Julia 기반 분산 API인 SciAI4Industry를 소개한다. 이는 도메인 분할을 통한 모델 병렬 딥러닝과 결합되어 산업 규모의 PDE 문제에 대한 3D 신경망 학습을 가능하게 하며, 200만 셀의 CO₂ 유동 시뮬레이션에 적용되어 기존 솔버 대비 5개 항목의 속도 향상과 3,200배의 비용 절감을 달성했으며, 병렬 효율성이 90% 이상을 기록했다.

ABSTRACT

Solving partial differential equations with deep learning makes it possible to reduce simulation times by multiple orders of magnitude and unlock scientific methods that typically rely on large numbers of sequential simulations, such as optimization and uncertainty quantification. Two of the largest challenges of adopting scientific AI for industrial problem settings is that training datasets must be simulated in advance and that neural networks for solving large-scale PDEs exceed the memory capabilities of current GPUs. We introduce a distributed programming API in the Julia language for simulating training data in parallel on the cloud and without requiring users to manage the underlying HPC infrastructure. In addition, we show that model-parallel deep learning based on domain decomposition allows us to scale neural networks for solving PDEs to commercial-scale problem settings and achieve above 90% parallel efficiency. Combining our cloud API for training data generation and model-parallel deep learning, we train large-scale neural networks for solving the 3D Navier-Stokes equation and simulating 3D CO2 flow in porous media. For the CO2 example, we simulate a training dataset based on a commercial carbon capture and storage (CCS) project and train a neural network for CO2 flow simulation on a 3D grid with over 2 million cells that is 5 orders of magnitudes faster than a conventional numerical simulator and 3,200 times cheaper.

연구 동기 및 목표

• 기존 시뮬레이터가 너무 느리고 HPC 자원이 필요로 하는 산업 환경에서 과학적 AI를 위한 대규모 학습 데이터셋을 생성하는 도전 과제를 해결한다.
• 단일 GPU의 메모리 용량을 초과하는 큰 3D PDE 문제를 학습하기 위해 도메인 분할을 통한 모델 병렬화를 통해 GPU 메모리 제약을 극복한다.
• 실제 산업 문제, 예를 들어 탄소 포집 및 저장(CCS)에서의 3D CO₂ 유동과 같은 대규모 딥러닝 대체 모델의 학습 가능성을 입증한다.
• 최적화 및 불확실성 정량화와 같은 많은 시뮬레이션을 요구하는 응용 분야에서 비용 효율적이고 고성능의 추론을 가능하게 한다.
• Julia에서의 고수준, 클라우드 네이티브 프로그래밍 인터페이스를 통해 HPC 수준의 시뮬레이션 능력을 대중화한다.

제안 방법

• 사용자가 클러스터를 직접 관리하지 않더라도 병렬적으로 대규모 학습 데이터셋을 생성할 수 있도록 클라우드 HPC 인프라를 추상화한 고수준 Julia API를 개발한다.
• 공간적 및 시간적 도메인을 여러 GPU에 나누어 분할하는 방식으로 모델 병렬 딥러닝을 구현함으로써, 큰 규모의 PDE 문제를 메모리 제약 내에서 수용할 수 있도록 한다.
• 시간에 따라 변화하는 PDE를 해결하기 위해 푸리에 신경 연산자(FNO) 아키텍처를 사용하며, GPU 간의 효율적 통신을 위해 텐서를 분할한다.
• 단일 노드 내의 GPU 간 고대역폭 NVLink 인터커넥트를 활용하여 학습 및 추론 중에도 90% 이상의 높은 병렬 효율성을 유지한다.
• 클라우드 VM(Azure ND96amsr 등)과 분산 학습 파이프라인을 통합하여 비용과 관리 오버헤드를 최소화하면서도 학습 및 추론의 확장성을 확보한다.
• 향후 더 큰 배치 크기 및 다중 노드 확장에 대비해 하이브리드 병렬화 원리를 적용하지만, 현재 구현은 배치 크기가 2인 모델 병렬화만을 사용한다.

실험 결과

연구 질문

• RQ1 고수준의 클라우드 네이티브 프로그래밍 인터페이스를 통해 HPC 인프라나 클라우드 인프라를 수동으로 관리할 필요 없이 대규모 학습 데이터셋의 생성을 단순화할 수 있는가?
• RQ2 도메인 분할을 통한 모델 병렬 딥러닝이 단일 GPU의 메모리 용량을 초과하는 큰 3D PDE 문제를 해결하는 데 실현 가능한가?
• RQ3 단일 노드 내에서 여러 GPU로 확장할 경우에도 높은 병렬 효율성(예: 90% 이상)을 달성할 수 있는가?
• RQ4 해당 접근 방식이 실제 산업 규모의 PDE 문제에서 기존 수치적 솔버 대비 속도와 비용 면에서 뛰어난 성능을 내는가?
• RQ5 이 프레임워크는 불확실성 정량화 및 최적화와 같은 고성능 과학 워크플로우를 산업 응용 분야에서 얼마나 효과적으로 지원할 수 있는가?

주요 결과

• 제안된 클라우드 기반 API는 HPC 클러스터나 클라우드 인프라를 수동으로 관리할 필요 없이 대규모 PDE 문제에 대한 스케일러블하고 자동화된 학습 데이터 생성을 가능하게 한다.
• 도메인 분할을 통한 모델 병렬 FNO 학습은 최대 8개의 A100 GPU에서 90% 이상의 병렬 효율성을 달성했으며, 200만 개 이상의 셀을 포함하는 3D 격자에서의 학습을 가능하게 했다.
• 학습된 FNO 모델은 기존 OPM 시뮬레이터 대비 1회 시뮬레이션당 0.12초로, 6.8시간이 소요되는 기존 솔버 대비 약 5개 항목의 속도 향상을 기록했다.
• 클라우드 VM 비용을 고려할 때 FNO 기반 시뮬레이션은 OPM 솔버 대비 1회 시뮬레이션당 3,200배 더 저렴하며, 비용은 각각 $0.11과 $3.4였다.
• FNO 대체 모델은 단지 1,848회의 시뮬레이션 후 비용 면에서 손익이 맞춰지며, 웰 배치 최적화와 같은 고성능 응용 분야에서 경제적으로 타당한 솔루션이 된다.
• 정적 PDE 문제를 해결할 경우 3D에서 최대 512³, 2D에서 최대 12,000²의 격자에서 학습이 가능함을 입증하여 상용 규모 문제에 대한 확장성을 보였다.