# TERRA: A Hierarchical Parallel Training and Memory Orchestration Framework for High-Resolution AI-based Earth Modeling

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 27 of its 41 references that are in the map.
- First posted on arXiv 15 Aug 2026.

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

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [Reducing Activation Recomputation in Large Transformer Models](https://mlsanity.com/paper/2205.05198/) (2022)
- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://mlsanity.com/paper/2104.07857/) (2021)
- [DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models](https://mlsanity.com/paper/2309.14509/) (2023)
- [FuXi: A cascade machine learning forecasting system for 15-day global weather forecast](https://mlsanity.com/paper/2306.12873/) (2023)
- [FourCastNet: Accelerating Global High-Resolution Weather Forecasting using Adaptive Fourier Neural Operators](https://mlsanity.com/paper/2208.05419/) (2022)
- [Scaling transformer neural networks for skillful and reliable medium-range weather forecasting](https://mlsanity.com/paper/2312.03876/) (2023)
- [Memory-Efficient Backpropagation Through Time](https://mlsanity.com/paper/1606.03401/) (2016)
- [Dynamic Tensor Rematerialization](https://mlsanity.com/paper/2006.09616/) (2020)
- [ORBIT: Oak Ridge Base Foundation Model for Earth System Predictability](https://mlsanity.com/paper/2404.14712/) (2024)
- [PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management](https://mlsanity.com/paper/2108.05818/) (2021)
- [The Case for Strong Scaling in Deep Learning: Training Large 3D CNNs with Hybrid Parallelism](https://mlsanity.com/paper/2007.12856/) (2020)
- [ORBIT-2: Scaling Exascale Vision Foundation Models for Weather and Climate Downscaling](https://mlsanity.com/paper/2505.04802/) (2025)
- [High Resolution Medical Image Analysis with Spatial Partitioning](https://mlsanity.com/paper/1909.03108/) (2019)
- [AERIS: Argonne Earth Systems Model for Reliable and Skillful Predictions](https://mlsanity.com/paper/2509.13523/) (2025)
- [SuperOffload: Unleashing the Power of Large-Scale LLM Training on Superchips](https://mlsanity.com/paper/2509.21271/) (2025)
- [Rockmate: an Efficient, Fast, Automatic and Generic Tool for Re-materialization in PyTorch](https://mlsanity.com/paper/2307.01236/) (2023)
- [Jigsaw: Training Multi-Billion-Parameter AI Weather Models with Optimized Model Parallelism](https://mlsanity.com/paper/2507.05753/) (2025)
- [ShardTensor: Domain Parallelism for Scientific Machine Learning](https://mlsanity.com/paper/2605.11111/) (2026)

## Abstract (arXiv)

Training high-resolution AI-based Earth forecasting models is memory-intensive. Window-based Swin Transformers reduce the quadratic cost of global attention, but existing distributed systems such as AERIS primarily target pixel-level models and do not jointly support convolutional sampling modules and shifted-window execution. Long-lead rollout finetuning further increases activation memory. To address these challenges, we present TERRA, a hierarchical parallel training framework for high-resolution Earth forecasting. TERRA introduces Sampling-Aware Window, Sequence, and Tensor Parallelism (SAWSTP), which preserves spatially contiguous layouts for sampling modules and routes tokens into topology-aware ragged window layouts for Transformer execution. For long-lead finetuning, Memory Orchestration (MO) provides rollout-aware checkpoint planning and combines input buffering with budget-constrained activation offloading. Experiments on the 1/12^circ GLORYS-based Wenhai workload show that TERRA supports models with up to 11.4B parameters on 96 H200 GPUs and sustains up to 39.76 PFLOPS, achieving 65.0% strong-scaling and 94.1% weak-scaling efficiency. Compared with checkpoint-only policies, MO further reduces peak allocated GPU memory by 32.2%--51.8% with at most 20.0% step-time overhead, which makes finetuning with smaller patch sizes and longer rollouts feasible for improved forecasting accuracy.