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

• URL: http://arxiv.org/abs/2608.15211v1
• Date: Sat, 15 Aug 2026 12:52:48 GMT
• Title: TERRA: A Hierarchical Parallel Training and Memory Orchestration Framework for High-Resolution AI-based Earth Modeling
• Abstract summary: Training high-resolution AI-based Earth forecasting models is memory-intensive.<n>We present TERRA, a hierarchical parallel training framework for high-resolution Earth forecasting.
• Score: 9.997208121319169
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Event-Driven Language Models with Sparse Neural Activity for Neuromorphic Hardware [5.066906328188419]
  Inference with transformer-based large language models (LLMs) is often limited by the memory-bound KV cache and quadratic attention cost.<n>We introduce a method that induces sparse neural activity in heavily quantized linear-attention models with minimal performance loss.
  arXiv  Detail & Related papers  (2026-08-31T08:28:01Z)
• Fine-Tuning of Transformer models with Frames [47.891438890695014]
  We propose a Fusion Frame-based solution for fine-tuning pre-trained models.<n>We show that FrameFT achieves performance on par with/exceeding state-of-the-art PEFT techniques.<n>Our experiments show that FrameFT achieves performance on par with/exceeding state-of-the-art PEFT techniques, but needs far fewer trainable parameters.
  arXiv  Detail & Related papers  (2026-08-26T22:13:45Z)
• Efficient Reinforcement Learning for Long-Horizon Tool-Use Agentic Tasks [3.5732010112106507]
  Long-horizon tool-using agents must reason over user goals, domain policies, tool calls, simulator state, and delayed verifiable rewards.<n>We present SINKFLEX-RL, a modular training system for RL in dual-control tool-use environments.
  arXiv  Detail & Related papers  (2026-08-11T01:30:56Z)
• AESOP: Adversarial Execution-path Selection to Overload Deep Learning Pipelines [6.773252452930085]
  Modern machine learning deployments increasingly compose specialized models into dynamic inference pipelines.<n>We show that this structure creates an efficiency-attack surface that existing methods targeting single models cannot exploit.<n>We present AESOP, a framework combining vulnerability-guided path ranking with adaptive loss weighting.
  arXiv  Detail & Related papers  (2026-05-09T09:16:49Z)
• DexWorldModel: Causal Latent World Modeling towards Automated Learning of Embodied Tasks [54.32016216994156]
  We introduce the Causal Latent World Model (CLWM), which employs DINOv3 features as generative targets to disentangle interaction semantics from visual noise.<n>To overcome memory scaling, CLWM features a Dual-State Test-Time Training (TTT) Memory that guarantees a strict $mathcalO(1)$ footprint for long-horizon tasks.<n>We present EmbodiChain, an online framework that establishes the Efficiency Law by injecting an infinite flow of physics-grounded trajectories during training.
  arXiv  Detail & Related papers  (2026-04-13T03:19:36Z)
• One-for-All: A Lightweight Stabilized and Parameter-Efficient Pre-trained LLM for Time Series Forecasting [4.364999214109123]
  One-for-All is a framework for adapting pre-trained Large Language Models for time-series analysis.<n>rsLoRA introduces a mathematically grounded rank-stabilization mechanism that enables provable gradient stability at low ranks.<n>One-for-All achieves state-of-the-art efficiency-accuracy trade-offs.
  arXiv  Detail & Related papers  (2026-03-31T13:54:43Z)
• SIMPLER: Efficient Foundation Model Adaptation via Similarity-Guided Layer Pruning for Earth Observation [3.1313965020437546]
  We introduce SIMPLER, a pre-fine-tuning architecture selection method that reduces inference and deployment costs.<n>On Prithvi-EO-2, SIMPLER prunes up to 79% of parameters while retaining 94% of baseline performance, yielding a 2.1x training speedup and 2.6x inference speedup.
  arXiv  Detail & Related papers  (2026-03-20T11:38:32Z)
• PFGNet: A Fully Convolutional Frequency-Guided Peripheral Gating Network for Efficient Spatiotemporal Predictive Learning [27.26429269735324]
  PFGNet is a fully convolutional framework that dynamically modulates receptive fields through pixel-wise frequency-guided gating.<n> PFGNet delivers SOTA or near-SOTA forecasting performance with substantially fewer parameters and FLOPs.
  arXiv  Detail & Related papers  (2026-02-24T04:31:12Z)
• ProAct: Agentic Lookahead in Interactive Environments [56.50613398808361]
  ProAct is a framework that enables agents to internalize accurate lookahead reasoning through a two-stage training paradigm.<n>We introduce Grounded LookAhead Distillation (GLAD), where the agent undergoes supervised fine-tuning on trajectories derived from environment-based search.<n>We also propose the Monte-Carlo Critic (MC-Critic), a plug-and-play auxiliary value estimator designed to enhance policy-gradient algorithms.
  arXiv  Detail & Related papers  (2026-02-05T05:45:16Z)
• Efficient Multi-Instance Generation with Janus-Pro-Dirven Prompt Parsing [53.295515505026096]
  Janus-Pro-driven Prompt Parsing is a prompt- parsing module that bridges text understanding and layout generation.<n>MIGLoRA is a parameter-efficient plug-in integrating Low-Rank Adaptation into UNet (SD1.5) and DiT (SD3) backbones.<n>The proposed method achieves state-of-the-art performance on COCO and LVIS benchmarks while maintaining parameter efficiency.
  arXiv  Detail & Related papers  (2025-03-27T00:59:14Z)
• Flat-LoRA: Low-Rank Adaptation over a Flat Loss Landscape [52.98187034726091]
  We introduce Flat-LoRA, which aims to identify a low-rank adaptation situated in a flat region of the full parameter space.<n>We show that Flat-LoRA improves both in-domain and out-of-domain generalization.
  arXiv  Detail & Related papers  (2024-09-22T11:24:10Z)
• HiRE: High Recall Approximate Top-$k$ Estimation for Efficient LLM Inference [68.59839755875252]
  HiRE comprises of two novel components: (i) a compression scheme to cheaply predict top-$k$ rows/columns with high recall, followed by full computation restricted to the predicted subset, and (ii) DA-TOP-$k$: an efficient multi-device approximate top-$k$ operator. We demonstrate that on a one billion parameter model, HiRE applied to both the softmax as well as feedforward layers, achieves almost matching pretraining and downstream accuracy, and speeds up inference latency by $1.47times$ on a single TPUv5e device.
  arXiv  Detail & Related papers  (2024-02-14T18:04:36Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.