Jigsaw：针对大规模气象模型优化模型并行的训练方法
机器学习
2025-07-09 v1
摘要
AI 方法已彻底变革大气预测, 最近在中程预测上的成功促进了气候基础模型的开发。在高空间分辨率和更长预报时限上准确建模复杂大气动力学 requires 大型神经网络和 gigabytes 大小的数据样本, 这使得加速器内存和 I/O 带宽成为模型训练的瓶颈。我们引入 WeatherMixer, 一个基于多层感知器的架构, 其工作负载随输入大小线性比例扩展, 允许模型在类似数值天气预测的准确度下学习全局气象现象。为应对计算需求, 我们提出了 Jigsaw, 一种新颖的模型并行方案, 采用域并行和张量并行, 消除内存冗余。Jigsaw 在计算-通信受限系统中超越了最先进的强比例性能, 并在 I/O 带宽受限系统中实现超标量弱比例性能。我们将训练扩展到 256 个 GPU, 实现峰值性能达到 9 和 11 PFLOPs, 分别达到理论峰值的 23% 和 28%, 实现 68% 和 72% 的扩展效率, 相较于无模型并行的 51% 有显著提升。
引用
@article{arxiv.2507.05753,
title = {Jigsaw: Training Multi-Billion-Parameter AI Weather Models with Optimized Model Parallelism},
author = {Deifilia Kieckhefen and Markus Götz and Lars H. Heyen and Achim Streit and Charlotte Debus},
journal= {arXiv preprint arXiv:2507.05753},
year = {2025}
}
备注
12 pages, 10 figures