# ORBIT: Oak Ridge Base Foundation Model for Earth System Predictability

arXiv: https://arxiv.org/abs/2404.14712 · HTML: https://mlsanity.com/paper/2404.14712/

## Key facts

- Cited by 17 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 30 references that are in the map.
- First posted on arXiv 23 Apr 2024.

## Builds on (most cited first, 19 of 19)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [FourCastNet: A Global Data-driven High-resolution Weather Model using Adaptive Fourier Neural Operators](https://mlsanity.com/paper/2202.11214/) (2022)
- [Scaling Vision Transformers to 22 Billion Parameters](https://mlsanity.com/paper/2302.05442/) (2023)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [ClimaX: A foundation model for weather and climate](https://mlsanity.com/paper/2301.10343/) (2023)
- [GraphCast: Learning skillful medium-range global weather forecasting](https://mlsanity.com/paper/2212.12794/) (2022)
- [WeatherBench 2: A benchmark for the next generation of data-driven global weather models](https://mlsanity.com/paper/2308.15560/) (2023)
- [Pangu-Weather: A 3D High-Resolution Model for Fast and Accurate Global Weather Forecast](https://mlsanity.com/paper/2211.02556/) (2022)
- [GSPMD: General and Scalable Parallelization for ML Computation Graphs](https://mlsanity.com/paper/2105.04663/) (2021)
- [Scaling transformer neural networks for skillful and reliable medium-range weather forecasting](https://mlsanity.com/paper/2312.03876/) (2023)
- [PyTorch Distributed: Experiences on Accelerating Data Parallel Training](https://mlsanity.com/paper/2006.15704/) (2020)
- [Memory-Efficient Backpropagation Through Time](https://mlsanity.com/paper/1606.03401/) (2016)
- [torchgpipe: On-the-fly Pipeline Parallelism for Training Giant Models](https://mlsanity.com/paper/2004.09910/) (2020)
- [PipeTransformer: Automated Elastic Pipelining for Distributed Training of Transformers](https://mlsanity.com/paper/2102.03161/) (2021)
- [Optimizing Distributed Training on Frontier for Large Language Models](https://mlsanity.com/paper/2312.12705/) (2023)

## Built on by (most cited first, 17 of 17)

- [PANGAEA: A Global and Inclusive Benchmark for Geospatial Foundation Models](https://mlsanity.com/paper/2412.04204/) (2024)
- [Can AI weather models predict out-of-distribution gray swan tropical cyclones?](https://mlsanity.com/paper/2410.14932/) (2024)
- [Improving Parallel Program Performance with LLM Optimizers via Agent-System Interfaces](https://mlsanity.com/paper/2410.15625/) (2024)
- [ORBIT-2: Scaling Exascale Vision Foundation Models for Weather and Climate Downscaling](https://mlsanity.com/paper/2505.04802/) (2025)
- [AERIS: Argonne Earth Systems Model for Reliable and Skillful Predictions](https://mlsanity.com/paper/2509.13523/) (2025)
- [Dynamical errors in machine learning forecasts](https://mlsanity.com/paper/2504.11074/) (2025)
- [Finetuning a Weather Foundation Model with Lightweight Decoders for Unseen Physical Processes](https://mlsanity.com/paper/2506.19088/) (2025)
- [Data Readiness for Scientific AI at Scale](https://mlsanity.com/paper/2507.23018/) (2025)
- [On Neural Scaling Laws for Weather Emulation through Continual Training](https://mlsanity.com/paper/2603.25687/) (2026)
- [Global Attention with Linear Complexity for Exascale Generative Data Assimilation in Earth System Prediction](https://mlsanity.com/paper/2604.16590/) (2026)
- [Distributed Cross-Channel Hierarchical Aggregation for Foundation Models](https://mlsanity.com/paper/2506.21411/) (2025)
- [Pixel-Resolved Long-Context Learning for Turbulence at Exascale: Resolving Small-scale Eddies Toward the Viscous Limit](https://mlsanity.com/paper/2507.16697/) (2025)
- [Automated Data Readiness for Scientific AI](https://mlsanity.com/paper/2607.02771/) (2026)
- [Comprehensive Performance Modeling and System Design Insights for Foundation Models](https://mlsanity.com/paper/2410.00273/) (2024)
- [BioAnalyst: A Foundation Model for Biodiversity](https://mlsanity.com/paper/2507.09080/) (2025)
- [TERRA: A Hierarchical Parallel Training and Memory Orchestration Framework for High-Resolution AI-based Earth Modeling](https://mlsanity.com/paper/2608.15211/) (2026)
- [4D Parallelism Unlocks Exascale Bayesian Neural Networks for High-Fidelity Atmospheric Modeling](https://mlsanity.com/paper/2609.12815/) (2026)

## Abstract (arXiv)

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.