TERRA：面向高分辨率AI地球建模的分层并行训练与内存编排框架
TERRA: A Hierarchical Parallel Training and Memory Orchestration Framework for High-Resolution AI-based Earth Modeling
浏览论文内容
中文总结 AI 辅助
TERRA是面向高分辨率AI地球预报的分层并行训练与内存编排框架，通过SAWSTP和MO技术，在96个H200 GPU上支持114亿参数模型，实现高算力与内存优化，提升预报精度。
中文摘要 AI 辅助
训练高分辨率AI地球预报模型内存需求极大。基于窗口的Swin Transformer降低了全局注意力的二次成本，但现有分布式系统（如AERIS）主要面向像素级模型，未同时支持卷积采样模块与移位窗口执行；长期rollout微调进一步增加了激活内存。为应对这些挑战，本文提出TERRA，一种面向高分辨率地球预报的分层并行训练框架。TERRA引入采样感知窗口、序列与张量并行（SAWSTP），为采样模块保留空间连续布局，并将令牌路由为拓扑感知的不规则窗口布局以用于Transformer执行。针对长期rollout微调，内存编排（MO）提供rollout感知的检查点规划，结合输入缓冲与预算约束的激活卸载。在基于1/12° GLORYS的Wenhai工作负载上的实验表明，TERRA可在96个H200 GPU上支持参数规模达114亿的模型，维持最高39.76 PFLOPS的算力，实现65.0%的强扩展效率与94.1%的弱扩展效率。与仅检查点策略相比，MO可进一步将峰值已分配GPU内存降低32.2%至51.8%，且步骤时间开销最多为20.0%，这使得采用更小patch尺寸与更长rollout的微调成为可能，以提升预报精度。
英文摘要
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.
发表机构
- School of Computer Science and Technology, University of Science and Technology of China（中国科学技术大学计算机科学与技术学院）
- School of Artificial Intelligence and Data Science, University of Science and Technology of China（中国科学技术大学人工智能与数据科学学院）
- Laoshan Laboratory（崂山实验室）
- Ocean University of China（中国海洋大学）
机构由 AI 辅助整理，请以论文原文为准。
