arXiv:2608.15211cs.CVcs.DC2026-08

TERRA框架让高分辨率地球建模训练更高效,突破内存瓶颈。

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

论文配图:TERRA: A Hierarchical Parallel Training and Memory Orchestration Framework for High-Resolution AI-based Earth Modeling
图 1 · 摘自论文原文
  • 分层并行设计兼顾采样模块与Transformer的计算需求
  • 支持114亿参数模型在96张H200上训练,达39.76 PFLOPS算力
  • 动态内存调度使长周期微调内存降低超三成,适合精准气候预测

高分辨率人工智能地球预报模型训练面临巨大内存压力。基于窗口的Swin Transformer虽降低了全局注意力的二次复杂度,但现有分布式系统如AERIS主要针对像素级模型,无法同时支持卷积采样模块和移位窗口执行。长周期滚动微调进一步加剧激活内存占用。为此,我们提出TERRA,一种面向高分辨率地球预报的分层并行训练框架。TERRA引入采样感知的窗口、序列与张量并行(SAWSTP),保持采样模块的空间连续布局,并将令牌路由至拓扑感知的不规则窗口布局以供Transformer执行。针对长周期微调,内存编排(MO)提供滚动感知的检查点规划,并结合输入缓冲与预算约束的激活卸载。在$1/12^\ abla$ GLORYS-based Wenhai工作负载上的实验表明,TERRA可在96张H200 GPU上支持高达11.4B参数的模型,维持最高39.76 PFLOPS算力,强弱扩展效率分别达65.0%和94.1%。相较于仅检查点策略,MO将峰值分配显存减少32.2%–51.8%,步时开销最多增加20.0%,使得采用更小块尺寸和更长滚动周期的微调成为可能,从而提升预报精度。

原文摘要 · Abstract (English)

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.

地球建模并行训练内存优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。