Jigsaw: Training Multi-Billion-Parameter AI Weather Models with Optimized Model Parallelism

• URL: http://arxiv.org/abs/2507.05753v1
• Date: Tue, 08 Jul 2025 07:57:08 GMT
• Title: Jigsaw: Training Multi-Billion-Parameter AI Weather Models with Optimized Model Parallelism
• Abstract summary: Accurate modeling of complex atmospheric dynamics requires large neural networks and gigabyte-sized data samples.<n>We introduce WeatherMixer, a multi-layer-perceptron-based architecture whose workload scales linearly with input size.<n>To cope with the computational demand, we propose Jigsaw, a novel model parallelization scheme.
• Score: 0.08388591755871733
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• 4D Parallelism Unlocks Exascale Bayesian Neural Networks for High-Fidelity Atmospheric Modeling [1.0277586860424235]
  We present BEAST, the first-ever Bayesian Swin Transformer for atmospheric forecasting on 0.25$circ$ global resolution.<n>We numerically train BEAST as a 700-million- parameter model with 96 random weight samples on 384 nodes on 40 years of data for nearly one million gradient updates.<n>This model achieves predictive skill scores competitive with state-of-the-art atmospheric AI models and numerical models, and can predict extreme events with exceptional skill.
  arXiv  Detail & Related papers  (2026-09-11T13:08:17Z)
• Piper: Efficient Large-Scale MoE Training via Resource Modeling and Pipelined Hybrid Parallelism [1.0130502379786543]
  Training MoE models on HPC platforms is hindered by large memory footprints, frequent large-scale communication across heterogeneous networks, and severe workload imbalance.<n>We develop a mathematical model that quantifies memory, compute, and communication requirements for MoE configurations under various parallelization schemes.<n>We introduce Piper, a framework that leverages resource modeling to identify efficient training strategies for MoE models on target HPC platforms.
  arXiv  Detail & Related papers  (2026-05-06T15:47:14Z)
• MiniCPM-SALA: Hybridizing Sparse and Linear Attention for Efficient Long-Context Modeling [80.48332380100915]
  MiniCPM-SALA is a hybrid model that integrates the high-fidelity long-context modeling of sparse attention with the global efficiency of linear attention.<n>On a single NVIDIA A6000D GPU, the model achieves up to 3.5x the inference speed of the full-attention model at the sequence length of 256K tokens.
  arXiv  Detail & Related papers  (2026-02-12T09:37:05Z)
• FourCastNet 3: A geometric approach to probabilistic machine-learning weather forecasting at scale [91.84761739154366]
  FourCastNet 3 advances global weather modeling by implementing a scalable, geometric machine learning (ML) approach to probabilistic ensemble forecasting.<n>FourCastNet 3 delivers forecasting accuracy that surpasses leading conventional ensemble models and rivals the best diffusion-based methods.<n>Its computational efficiency, medium-range probabilistic skill, spectral fidelity, and rollout stability at subseasonal timescales make it a strong candidate for improving meteorological forecasting and early warning systems through large ensemble predictions.
  arXiv  Detail & Related papers  (2025-07-16T11:22:18Z)
• Modernizing CNN-based Weather Forecast Model towards Higher Computational Efficiency [5.781137818421603]
  We introduce a modernized CNN-based model for global weather forecasting.<n> KAI-a incorporates a scale-invariant architecture and InceptionNeXt-based blocks within a geophysically-aware design.<n>It is trained on the ERA5 daily dataset with 67 atmospheric variables and completes training in just 12 hours on a single NVIDIA L40s GPU.
  arXiv  Detail & Related papers  (2025-07-15T01:16:32Z)
• Protocol Models: Scaling Decentralized Training with Communication-Efficient Model Parallelism [59.79227116582264]
  Scaling models has led to significant advancements in deep learning, but training these models in decentralized settings remains challenging.<n>We propose a novel compression algorithm that compresses both forward and backward passes, enabling up to 99% compression with no convergence degradation.
  arXiv  Detail & Related papers  (2025-06-02T02:19:22Z)
• Decoupling Spatio-Temporal Prediction: When Lightweight Large Models Meet Adaptive Hypergraphs [12.867023510751787]
  STH-SepNet is a novel framework that decouples temporal and spatial expressiveness to both efficiency and precision.<n>S-SepNet offers a pragmatic and scalable solution for temporal prediction in real-world applications.<n>This work may provide a promising lightweight framework for temporal prediction, aiming to reduce computational demands and while enhancing predictive performance.
  arXiv  Detail & Related papers  (2025-05-26T07:37:39Z)
• Partitioned Neural Network Training via Synthetic Intermediate Labels [0.0]
  GPU memory constraints have become a notable bottleneck in training such sizable models.<n>This study advocates partitioning the model across GPU and generating synthetic intermediate labels to train individual segments.<n>This approach results in a more efficient training process that minimizes data communication while maintaining model accuracy.
  arXiv  Detail & Related papers  (2024-03-17T13:06:29Z)
• KARINA: An Efficient Deep Learning Model for Global Weather Forecast [2.9687381456164004]
  KARINA achieves forecasting accuracy comparable to higher-resolution counterparts with significantly less computational resources. KARINA combines ConvNext, SENet, and Geocyclic Padding to enhance weather forecasting at a 2.5deg resolution. KARINA sets new benchmarks in weather forecasting accuracy, surpassing existing models like the ECMWF S2S reforecasts at a lead time of up to 7 days.
  arXiv  Detail & Related papers  (2024-03-13T06:41:37Z)
• A Dynamical Model of Neural Scaling Laws [79.59705237659547]
  We analyze a random feature model trained with gradient descent as a solvable model of network training and generalization. Our theory shows how the gap between training and test loss can gradually build up over time due to repeated reuse of data.
  arXiv  Detail & Related papers  (2024-02-02T01:41:38Z)
• Towards a Better Theoretical Understanding of Independent Subnetwork Training [56.24689348875711]
  We take a closer theoretical look at Independent Subnetwork Training (IST) IST is a recently proposed and highly effective technique for solving the aforementioned problems. We identify fundamental differences between IST and alternative approaches, such as distributed methods with compressed communication.
  arXiv  Detail & Related papers  (2023-06-28T18:14:22Z)
• A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training [13.346719319555943]
  Mixture-of-Experts (MoE) is a neural network architecture that adds sparsely activated expert blocks to a base model. Current distributed deep learning frameworks are limited in their ability to train high-quality MoE models with large base models. We present DeepSpeed-TED, a novel, three-dimensional, hybrid parallel algorithm that combines data, tensor, and expert parallelism.
  arXiv  Detail & Related papers  (2023-03-11T05:38:15Z)
• Learning Large-scale Subsurface Simulations with a Hybrid Graph Network Simulator [57.57321628587564]
  We introduce Hybrid Graph Network Simulator (HGNS) for learning reservoir simulations of 3D subsurface fluid flows. HGNS consists of a subsurface graph neural network (SGNN) to model the evolution of fluid flows, and a 3D-U-Net to model the evolution of pressure. Using an industry-standard subsurface flow dataset (SPE-10) with 1.1 million cells, we demonstrate that HGNS is able to reduce the inference time up to 18 times compared to standard subsurface simulators.
  arXiv  Detail & Related papers  (2022-06-15T17:29:57Z)
• M6-10T: A Sharing-Delinking Paradigm for Efficient Multi-Trillion Parameter Pretraining [55.16088793437898]
  Training extreme-scale models requires enormous amounts of computes and memory footprint. We propose a simple training strategy called "Pseudo-to-Real" for high-memory-footprint-required large models.
  arXiv  Detail & Related papers  (2021-10-08T04:24:51Z)
• Do We Really Need Deep Learning Models for Time Series Forecasting? [4.2698418800007865]
  Time series forecasting is a crucial task in machine learning, as it has a wide range of applications. Deep learning and matrix factorization models have been recently proposed to tackle the same problem with more competitive performance. In this paper, we try to answer whether these highly complex deep learning models are without alternative.
  arXiv  Detail & Related papers  (2021-01-06T16:18:04Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.