ORBIT: Oak Ridge Base Foundation Model for Earth System Predictability

📄 arXiv: 2404.14712v5 📥 PDF

作者: Xiao Wang, Siyan Liu, Aristeidis Tsaris, Jong-Youl Choi, Ashwin Aji, Ming Fan, Wei Zhang, Junqi Yin, Moetasim Ashfaq, Dan Lu, Prasanna Balaprakash

分类: physics.ao-ph, cs.AI, cs.DC, eess.IV, physics.geo-ph

发布日期: 2024-04-23 (更新: 2024-08-19)

💡 一句话要点

提出ORBIT模型以提升地球系统可预测性

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 地球系统预测 AI基础模型 视觉变换器 混合张量数据 气候建模 高性能计算 环境动态 数据整合

📋 核心要点

  1. 现有AI基础模型在处理地球系统预测时受限于模型规模和数据整合能力，无法全面应对复杂的环境动态。
  2. 本文提出的ORBIT模型采用混合张量数据正交并行技术，显著提升了模型规模和计算能力，达到1130亿参数。
  3. 在Frontier超级计算机上进行的测试表明，ORBIT模型的性能达到684 petaFLOPS至1.6 exaFLOPS，扩展效率高达85%。

📝 摘要（中文）

地球系统的可预测性受到环境动态复杂性和多变量的挑战。现有的AI基础模型虽然利用了大量异构数据，但由于模型规模和数据整合的限制，无法有效应对地球系统预测的全面挑战。为此，本文提出了奥克里奇基础模型（ORBIT），这是一个先进的视觉变换器模型，规模达到1130亿参数，采用新颖的混合张量数据正交并行技术。作为同类模型中最大的模型，ORBIT的规模超越了当前气候AI基础模型的千倍。通过在Frontier超级计算机上的性能扩展测试，ORBIT实现了684 petaFLOPS至1.6 exaFLOPS的持续吞吐量，扩展效率保持在41%至85%之间。这些突破为AI驱动的气候建模带来了新的进展，并显示出显著提升地球系统可预测性的潜力。

🔬 方法详解

问题定义：本文旨在解决现有AI基础模型在地球系统预测中的局限性，特别是在模型规模和数据整合方面的不足。现有方法无法充分利用复杂的环境动态和多样化的数据源，导致预测效果不佳。

核心思路：ORBIT模型通过引入混合张量数据正交并行技术，显著提升了模型的参数规模和计算效率。这种设计使得模型能够处理更复杂的环境数据，从而提高预测的准确性和可靠性。

技术框架：ORBIT模型采用视觉变换器架构，整体流程包括数据预处理、模型训练和性能评估等主要模块。通过并行计算，模型能够在多个GPU上高效运行，确保大规模数据的实时处理。

关键创新：ORBIT的最大创新在于其规模和计算能力，成为同类模型中最大的，参数量超越现有气候AI基础模型千倍。这一创新使得模型在处理复杂环境动态时具备更强的能力。

关键设计：在模型设计中，ORBIT采用了高效的损失函数和优化算法，以确保在大规模训练时的收敛性和稳定性。此外，模型结构经过精心设计，以支持高效的并行计算和数据处理。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在Frontier超级计算机上，ORBIT模型实现了684 petaFLOPS至1.6 exaFLOPS的持续吞吐量，扩展效率保持在41%至85%。这一性能显著高于现有气候AI基础模型，展示了其在处理复杂环境数据时的卓越能力。

🎯 应用场景

ORBIT模型的潜在应用领域包括气候预测、环境监测和自然灾害预警等。通过提升地球系统的可预测性，该模型能够为政策制定者和科学家提供更准确的决策支持，促进可持续发展和环境保护。未来，ORBIT有望在气候变化研究和生态系统管理中发挥重要作用。

📄 摘要（原文）

Earth system predictability is challenged by the complexity of environmental dynamics and the multitude of variables involved. Current AI foundation models, although advanced by leveraging large and heterogeneous data, are often constrained by their size and data integration, limiting their effectiveness in addressing the full range of Earth system prediction challenges. To overcome these limitations, we introduce the Oak Ridge Base Foundation Model for Earth System Predictability (ORBIT), an advanced vision transformer model that scales up to 113 billion parameters using a novel hybrid tensor-data orthogonal parallelism technique. As the largest model of its kind, ORBIT surpasses the current climate AI foundation model size by a thousandfold. Performance scaling tests conducted on the Frontier supercomputer have demonstrated that ORBIT achieves 684 petaFLOPS to 1.6 exaFLOPS sustained throughput, with scaling efficiency maintained at 41% to 85% across 49,152 AMD GPUs. These breakthroughs establish new advances in AI-driven climate modeling and demonstrate promise to significantly improve the Earth system predictability.