ORBIT: Oak Ridge Base Foundation Model for Earth System Predictability

  • 2024-08-07 17:19:15
  • Xiao Wang, Siyan Liu, Aristeidis Tsaris, Jong-Youl Choi, Ashwin Aji, Ming Fan, Wei Zhang, Junqi Yin, Moetasim Ashfaq, Dan Lu, Prasanna Balaprakash
  • 0

Abstract

Earth system predictability is challenged by the complexity of environmentaldynamics and the multitude of variables involved. Current AI foundation models,although advanced by leveraging large and heterogeneous data, are oftenconstrained by their size and data integration, limiting their effectiveness inaddressing the full range of Earth system prediction challenges. To overcomethese limitations, we introduce the Oak Ridge Base Foundation Model for EarthSystem Predictability (ORBIT), an advanced vision transformer model that scalesup to 113 billion parameters using a novel hybrid tensor-data orthogonalparallelism technique. As the largest model of its kind, ORBIT surpasses thecurrent climate AI foundation model size by a thousandfold. Performance scalingtests conducted on the Frontier supercomputer have demonstrated that ORBITachieves 684 petaFLOPS to 1.6 exaFLOPS sustained throughput, with scalingefficiency maintained at 41% to 85% across 49,152 AMD GPUs. These breakthroughsestablish new advances in AI-driven climate modeling and demonstrate promise tosignificantly improve the Earth system predictability.

 

Quick Read (beta)

loading the full paper ...