TERRA：用于基于人工智能的高分辨率地球建模的分层并行训练和内存编排框架

训练基于人工智能的高分辨率地球预测模型需要大量内存。基于窗口的 Swin Transformer 降低了全局注意力的二次成本，但现有的分布式系统（例如 AERIS）主要针对像素级模型，并且不共同支持卷积采样模块和移位窗口执行。长预测时距轨迹微调进一步增加了激活记忆。为了应对这些挑战，我们提出了 TERRA，这是一个用于高分辨率地球预测的分层并行训练框架。 TERRA 引入了采样感知窗口、序列和张量并行 (SAWSTP)，它保留了采样模块的空间连续布局，并将词元路由到拓扑感知的参差不齐的窗口布局中以供 Transformer 执行。对于长期微调，内存编排 (MO) 提供了可感知的检查点规划，并将输入缓冲与预算受限的激活卸载相结合。基于 $1/12^\circ$ GLORYS 的 Wenhai 工作负载上的实验表明，TERRA 在 96 个 H200 GPU 上支持高达 11.4B 参数的模型，并维持高达 $39.76$ PFLOPS，实现 $65.0\%$ 强扩展和 $94.1\%$ 弱扩展效率。与仅检查点策略相比，MO 进一步将峰值分配的 GPU 内存减少了 $32.2\%$--$51.8\%$，步长开销最多为 $20.0\%$，这使得可以通过更小的补丁大小和更长的预测轨迹进行微调，从而提高预测准确性。