发表机构
Karlsruhe Institute for Technology; CSC – IT Center for Science Ltd.; NVIDIA（卡尔斯鲁厄理工学院; CSC信息技术中心有限公司; 英伟达）
机构由 AI 辅助整理，请以论文原文为准。AI 中文总结
本文提出BEAST，首个贝叶斯Swin Transformer，通过4D并行化在百亿亿次超算上高效训练，实现高分辨率大气建模与不确定性量化，性能超越现有AI模型。
AI 中文摘要
我们提出了BEAST，这是首个能够在0.25°全球分辨率下进行大气预报的贝叶斯Swin Transformer，能够准确量化偶然不确定性和认知不确定性。为克服相关的计算瓶颈，我们设计了一种正交的4D并行化方案，引入了独特的域张量并行策略和一种新颖的不确定性并行方法，使我们能够充分利用GPU能力并高效扩展模型训练。对于一个24亿参数的模型，我们在JUPITER超级计算机上的20,480个NVIDIA GH200 GPU上实现了3.96 EFLOP/s的峰值性能。我们将BEAST训练为一个7亿参数的模型，使用96个随机权重样本，在384个节点上基于40年的数据进行了近一百万次梯度更新。该模型的预测技能得分与最先进的概率性大气AI模型和数值模型相当，并且能够以卓越的技能预测极端事件，同时生成大型集合的速度比当前最佳的AI模型快3到4倍。我们的贡献解锁了大气AI模型中高保真不确定性量化的潜力，预示着基于AI的模型在气候和地球系统科学中的新时代。
英文摘要
We present BEAST, the first-ever Bayesian Swin Transformer for atmospheric forecasting on 0.25$^\circ$ global resolution able to accurately quantify both aleatoric and epistemic uncertainty. To overcome the associated computational bottlenecks, we devise an orthogonal 4D-parallelization scheme that introduces a unique domain-tensor-parallelism strategy and a novel uncertainty parallel method, enabling us to fully leverage GPU capacity and efficiently scale model training. For a 2.4-billion-parameter model, we achieve a peak performance of 3.96 EFLOP/s on 20,480 NVIDIA GH200 GPUs on the JUPITER supercomputer. We train BEAST as a 700-million-parameter model with 96 random weight samples on 384 nodes on 40 years of data for nearly one million gradient updates. This model achieves predictive skill scores competitive with state-of-the-art probabilistic atmospheric AI models and numerical models, and can predict extreme events with exceptional skill, while generating large ensembles 3 to 4 times faster than the current-best AI model. Our contribution unlocks the potential of high-fidelity uncertainty quantification in atmospheric AI models, heralding a new era for AI-based models in climate and Earth system sciences.
