[论文解读] FourCastNet: Accelerating Global High-Resolution Weather Forecasting using Adaptive Fourier Neural Operators

Thorsten Kurth, Shashank Subramanian+7 more|arXiv (Cornell University)|Aug 8, 2022

Meteorological Phenomena and Simulations被引用 24

一句话总结

FourCastNet 使用自适应傅里叶神经算子来模拟地球系统动力学并在高分辨率下预测全球天气，与传统数值天气预报相比实现约 80,000x 更快的推断速度且具有强准确性。培训可扩展到 3,808 个 GPU，混合精度下达到 140.8 petaFLOPS。

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.

研究动机与目标

• 在气候变化导致极端天气增加的背景下，强调需要更快的高分辨率天气预报。
• 开发一个数据驱动的地球系统模拟器，以超越传统数值天气预报（NWP）模型的吞吐量。
• 在多台超级计算机上利用 GPU 加速展示可扩展的训练和推理性能。
• 表明模型能够高效生成准确的大型集合预测和更高分辨率的预报。

提出的方法

• 采用自适应傅里叶神经算子（AFNO）架构来建模地球系统状态张量在时间上的映射。
• 在 AFNO 块内使用基于光谱的空间混合（FFT 基）以实现高分辨率的全球预测。
• 在原生 0.25 度分辨率的 ERA5 再分析数据上进行训练，包含 20 个预报变量，预测 6 小时步长（t -> t+6h），包含预训练和微调阶段。
• 实现模型与数据并行（特征与数据并行），通过 NCCL 通信和 CUDA 图来最大化吞吐并最小化 CPU 开销。
• 使用 NVIDIA DALI 进行数据预处理，并通过 Just-In-Time 融合优化以减少内核启动开销；在 AFNO Transformer 块中采用基于 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,000x。
• 推理可对长达一周的即时预报给出准确的瞬时结果，并支持大型集合；在 Selene 节点上使用 8 个 A100 GPU 的 100 成员 24 小时集合预报在 12.41 节点秒内完成，而 IFS（L91 18 km）需要 984,000 节点秒。
• 能耗效率显著提升， FourCastNet 集合运行的能耗估计比 IFS 集合低约 10,000x。
• FourCastNet 实现了高分辨率（比以往 DL 模型高出 8 倍）并在精度方面优于先前的 DL 天气代理模型，包括解析热带气旋和大气河流等细小现象。