[논문 리뷰] FourCastNet: Accelerating Global High-Resolution Weather Forecasting using Adaptive Fourier Neural Operators

Thorsten Kurth, Shashank Subramanian외 7명|arXiv (Cornell University)|2022. 08. 08.

Meteorological Phenomena and Simulations인용 수 24

한 줄 요약

FourCastNet은 Adaptive Fourier Neural Operators를 사용하여 지구 시스템 역학을 모사하고 고해상도 글로벌 기상 예측을 수행하며 전통적인 NWP와 비교하여 추론을 약 80,000배 빠르고 강력한 정확도를 달성합니다. 혼합 정밀도에서 3,808 GPUs로 학습이 확장됩니다.

ABSTRACT

Extreme weather amplified by climate change is causing increasingly devastating impacts across the globe. The current use of physics-based numerical weather prediction (NWP) limits accuracy due to high computational cost and strict time-to-solution limits. We report that a data-driven deep learning Earth system emulator, FourCastNet, can predict global weather and generate medium-range forecasts five orders-of-magnitude faster than NWP while approaching state-of-the-art accuracy. FourCast-Net is optimized and scales efficiently on three supercomputing systems: Selene, Perlmutter, and JUWELS Booster up to 3,808 NVIDIA A100 GPUs, attaining 140.8 petaFLOPS in mixed precision (11.9%of peak at that scale). The time-to-solution for training FourCastNet measured on JUWELS Booster on 3,072GPUs is 67.4minutes, resulting in an 80,000times faster time-to-solution relative to state-of-the-art NWP, in inference. FourCastNet produces accurate instantaneous weather predictions for a week in advance, enables enormous ensembles that better capture weather extremes, and supports higher global forecast resolutions.

연구 동기 및 목표

• climate change로 인한 극심한 기상 현상의 증가에 대응하여 더 빠르고 고해상도 기상 예측의 필요성을 제시합니다.
• 전통적 수치 예보(NWP) 모델의 처리량을 능가하기 위해 데이터 기반의 지구 시스템 에뮬레이터를 개발합니다.
• GPU 가속화를 이용한 다중 슈퍼컴퓨터에서 확장 가능한 학습 및 추론 성능을 보여줍니다.
• 모델이 대규모 엔셈블리와 높은 해상도 예측을 효율적으로 생성할 수 있음을 보입니다.

제안 방법

• 지구 시스템 상태 텐서 간의 매핑을 모델링하기 위해 Adaptive Fourier Neural Operator(AFNO) 아키텍처를 채택합니다.
• AFNO 블록 내의 스펙트럴(FFT 기반) 공간 혼합을 사용하여 고해상도 글로벌 예측을 가능하게 합니다.
• native 0.25도 해상도에서 ERA5 재분석 데이터를 사용하고 20개 예측 변수를 통해 6시간 간격의 예측(t -> t+6h)을 예비 학습 및 미세 조정 단계로 수행합니다.
• NCCL 통신과 CUDA 그래프를 활용한 모델 및 데이터 병렬성(특징 및 데이터 병렬성)을 구현하여 처리량을 극대화하고 CPU 오버헤드를 최소화합니다.
• 데이터 전처리를 NVIDIA DALI로 수행하고 JIT 융합을 통해 커널 실행 오버헤드를 줄이며 AFNO 트랜스포머 블록에서 FFT 기반 스펙트럴 합성 및 MLP를 적용합니다.

실험 결과

연구 질문

• RQ1 AFNO 기반 아키텍처가 고해상도 글로벌 해상도에서 최첨단 기상 예측 정확도에 근접하거나 달성할 수 있는가?
• RQ2 FourCastNet은 리더십급 HPC 시스템에서 얼마나 빠르고 확장 가능하게 학습 및 추론이 가능한가?
• RQ3 대규모 DL 대리모델의 글로벌 기상 예측 및 엔셈블리에 대한 에너지 및 처리량 영향은 무엇인가?
• RQ4 모델이 리드 타임이 최대 일주일에 이르는 대규모 엔셈블리와 다단계 예측을 지원하는가?

주요 결과

• 다수의 HPC 시스템에서 학습이 3,808 NVIDIA A100 GPU로 확장되며 140.8 petaFLOPS(피크의 11.9%)에 도달합니다.
• JUWELS Booster에서 3,072 GPU를 사용하는 예비 학습 시간은 67.4분으로, 최신 NWP(IFS)에 비해 추론 시간-솔루션이 약 80,000배 빠르게 가능합니다.
• 추론은 최대 일주일까지의 instantaneous forecast를 생성하고 대규모 엔셈블리를 지원합니다; 100-member 24-hour 엔셈블리 예측은 Selene 노드의 8 A100 GPU에서 12.41 node-seconds로 실행되며, IFS(L91 18 km) 대비 984,000 node-seconds입니다.
• 에너지 효율성은 현저히 높아 FourCastNet 엔셈블리 실행의 에너지가 IFS 엔셈블리 대비 약 10,000배 낮은 것으로 추정됩니다.
• FourCastNet은 이전 DL 기상 대리모델보다 고해상도(이전 DL 모델의 8배)와 정확도를 달성하며, 열대 싸이클론 및 대기 강설류와 같은 미세 규모 현상을 해결하는 능력을 개선했습니다.