用新型并行策略训练千亿参数气象大模型，突破内存与带宽瓶颈。
Jigsaw: Training Multi-Billion-Parameter AI Weather Models with Optimized Model Parallelism
- 提出WeatherMixer架构，输入规模增大时计算量线性增长。
- 在256张GPU上实现9~11 PFLOPs峰值性能，达到理论值的23%~28%。
- 适合需要高分辨率长期天气预测的研究者和气象AI开发者。
基于AI的方法已彻底改变大气预报，近期中短期预报的成功推动了气候基础模型的发展。在高空间分辨率和长预报时效下精确建模复杂大气动力学，需要大型神经网络和吉字节级数据样本，导致加速器内存和I/O带宽成为训练瓶颈。我们提出WeatherMixer，一种基于多层感知机的架构，其工作负载随输入大小线性增长，使模型在精度上可媲美数值天气预报，能学习全球天气现象。为应对计算需求，我们提出Jigsaw，一种结合域并行与张量并行的新模型并行方案，消除内存冗余。Jigsaw在计算-通信受限系统中实现最优强缩放，在I/O带宽受限系统中实现超标量弱缩放。我们将训练扩展至256张GPU，达到9和11 PFLOPs的峰值性能，分别为理论峰值的23%和28%，缩放效率达68%和72%（无模型并行时仅为51%）。
原文摘要 · Abstract (English)
AI-based methods have revolutionized atmospheric forecasting, with recent successes in medium-range forecasting spurring the development of climate foundation models. Accurate modeling of complex atmospheric dynamics at high spatial resolutions and longer lead times requires large neural networks and gigabyte-sized data samples, making accelerator memory and I/O-bandwidth the bottlenecks for model training. We introduce WeatherMixer, a multi-layer-perceptron-based architecture whose workload scales linearly with input size, allowing the model to learn global weather phenomena at accuracies similar to numerical weather prediction. To cope with the computational demand, we propose Jigsaw, a novel model parallelization scheme that employs both domain and tensor parallelism, eliminating memory redundancy. Jigsaw exceeds state-of-the-art performance in strong scaling in compute-communication-limited systems and achieves superscalar weak scaling in I/O-bandwidth-limited systems. We scale training to 256 GPUs, reaching peak performances of 9 and 11 PFLOPs, 23% and 28% of theoretical peaks, achieving 68% and 72% scaling efficiency versus 51% without model parallelism.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。