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

Ruohan Wu; Ziqi Zhu; Yang Zhao; Jiarui Tang; Yingzhe Cui; Junshi Chen; Zhao Jing; Jun Shi

cs.CVAug 15, 2026

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

Authors: Ruohan Wu, Ziqi Zhu, Yang Zhao, Jiarui Tang, Yingzhe Cui, Junshi Chen, Zhao Jing, Jun Shi,

Organizations: School of Computer Science and Technology, University of Science and Technology of China, Hefei 230026, China · School of Artificial Intelligence and Data Science, University of Science and Technology of China, Hefei 230026, China · Department of Ocean Big Data and Artificial Intelligence, Laoshan Laboratory, Qingdao, China · Frontiers Science Center for Deep Ocean Multispheres and Earth System and the Key Laboratory of Physical Oceanography/Academy of Future Ocean, Ocean University of China, Qingdao, China

Abstract

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.

Figures & tables

Panel of Figure 1 — Fig. 1: High-resolution Earth forecasting model and training stages.

Notation	Description
$H,W$	Spatial resolution of the input field
$\hat{H},\hat{W}$	Padded spatial resolution of the input field
$C$	Number of input physical variables
$p$	Patch size
$\alpha$	Stride of each down- or up-sampling module
$d$	Embedding dimension

TABLE I: Notation for the Wenhai model.

Panel of Figure 3 — Fig. 2: Memory pressure in training high-resolution Earth forecasting models.

	Model setting		Parallelism support
Framework	Model architecture	Resolution	Parallel sampling	WP	SP	TP	PP	DP	FSDP
ORBIT	Vision Transformer	$1.4^{\circ}$	$\times$	$\times$	$\times$	$\checkmark$	$\times$	$\checkmark$	$\checkmark$
AERIS	Pixel-level Swin Transformer	$0.25^{\circ}$	$\times$	$\checkmark$	$\checkmark$	$\times$	$\checkmark$	$\checkmark$	$\times$
TERRA	Hierarchical Swin Transformer	$1/12^{\circ}$	$\checkmark$	$\checkmark$	$\checkmark$	$\checkmark$	$\times$	$\checkmark$	$\checkmark$

TABLE II: Comparison of model settings and parallelism support in ORBIT [ 8 ] , AERIS [ 13 ] , and TERRA.

Fig. 3: Overview of TERRA.

Fig. 4: SAWSTP Training Framework for the hierarchical Swin Transformer.

Fig. 5: Multi-level Transformer topology with DMP and shift token routes in SAWSTP.

Fig. 6: Fixed, ragged, and density-aware window assignment.

Strategy	Communication
Fixed $\mathrm{Topo}(m,n)$	Two-dimensional WP P2P
Ragged $\mathrm{Topo}(mn,1)$	Boundary-only P2P
Density-aware $\mathrm{Topo}(mn/s,1,s)$	Boundary-only P2P + SP all-to-all

TABLE III: Communication characteristics of window assignments.

Panel of Figure 10 — Fig. 7: Allocated-memory timelines for ten-lead finetuning under different policies.

Component	Configuration
Cluster scale	Up to 12 nodes and 96 GPUs
GPU	$8\times$ NVIDIA H200 GPUs per node
CPU	$2\times$ Intel Xeon Platinum 8558, 96 cores per node
Host memory	1.8 TB per node
Intra-node interconnect	NVLink
Inter-node interconnect	Mellanox RDMA network

TABLE IV: Experimental platform.

Model	$L$	$d$	$N_{\mathrm{h}}$	$p$	$\alpha$	$w$	Number of Parameters
Small	10	1024	8	4	2	8	0.18B
Medium	10	4096	32	4	2	8	2.85B
Large	10	8192	64	4	2	32	11.4B

TABLE V: Model configurations.

Panel of Figure 13 — Fig. 8: Training-loss consistency between serial execution and TERRA parallel configurations.

Panel of Figure 14 — Fig. 9: Strong scaling performance.

Panel of Figure 15 — Fig. 10: Weak scaling performance.

Fig. 11: Largest trainable model size for the AERIS-style baseline and TERRA configurations with TP.

Panel of Figure 17 — Fig. 12: Training-step time under fixed AERIS-style topologies and TERRA’s ragged window assignments.

Fig. 13: Peak allocated GPU memory of checkpoint-only policies for ten-lead finetuning.

Fig. 14: Peak allocated GPU memory and step time for input buffering and budget-constrained offloading.

Patch/DMP/DP	$M_{b}$ (GB)	Offload Volume (GB)
		Sampling	Transformer
2/4/2	40	9.7	16.7
	30	72.5	29.9
2/8/1	20	4.3	9.9
	16	36.2	14.9
4/1/8	80	22.8	10.8

TABLE VI: Breakdown of per-GPU activation offload volume per training step.

Panel of Figure 21 — Fig. 15: Mean RMSE of Wenhai forecasts at different lead times.

Fig. 16: Day-10 SST forecasts over the western North Pacific.

Distributed Training

Distributed training coordinates computation across multiple devices or machines to train models that exceed single-device limits.

See all 130 papers

4D Parallelism Unlocks Exascale Bayesian Neural Networks for High-Fidelity Atmospheric Modeling
Sep 11, 2026Deifilia Kieckhefen, Juan Pedro Gutiérrez Hermosillo Muriedas, Lars Helge Heyen +12Bayesian Neural Networks Uncertainty Quantification
Karlsruhe Institute of Technology, Germany · Forschungszentrum Jülich, Germany · CSC – IT Center for Science Ltd., Finnland +1
TESSERA v2: Scaling Pixel-wise Earth Foundation Models
Jul 4, 2026Zhengpeng Feng, Sadiq Jaffer, Ira Shokar +10Geospatial Foundation Models AI Scaling Laws
University of Cambridge · 2NVIDIA · 3dClimate Labs
Global Attention with Linear Complexity for Exascale Generative Data Assimilation in Earth System Prediction
Apr 17, 2026Xiao Wang, Zezhong Zhang, Isaac Lyngaas +10Climate Science High-Performance Computing
Oak Ridge National Laboratory, Oak Ridge, TN, USA · Auburn University, Auburn, AL, USA · Florida State University, Tallahassee, FL, USA +2

Sep 11, 2026J/K move · Enter open · S save

4D Parallelism Unlocks Exascale Bayesian Neural Networks for High-Fidelity Atmospheric Modeling

Deifilia Kieckhefen, Juan Pedro Gutiérrez Hermosillo Muriedas, Lars Helge Heyen, Mathis Bode+11

Karlsruhe Institute of Technology, Germany · Forschungszentrum Jülich, Germany · CSC – IT Center for Science Ltd., Finnland+1

Bayesian Neural Networks Uncertainty Quantification High-Performance Computing Probabilistic Forecasting Distributed Training Weather Forecasting

Open paper AskPDF arXiv

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.

Abstract

Distributed Training

Explore similar work

4D Parallelism Unlocks Exascale Bayesian Neural Networks for High-Fidelity Atmospheric Modeling

TESSERA v2: Scaling Pixel-wise Earth Foundation Models

Global Attention with Linear Complexity for Exascale Generative Data Assimilation in Earth System Prediction