Combining Domain and Tensor Parallelism to Train Multi-billion-parameter AI Weather Models
Published in Proceedings of the Platform for Advanced Scientific Computing Conference (PASC 26), 2026
Training AI weather foundation models at high spatial resolution and billion-parameter scale is bottlenecked by accelerator memory and I/O bandwidth. We introduce Jigsaw, a novel parallelization scheme that combines domain parallelism and tensor parallelism to eliminate memory redundancy during training. Applied to the WeatherMixer architecture, Jigsaw exceeds state-of-the-art performance in compute-communication-limited settings and achieves superscalar weak scaling in I/O-bandwidth-limited settings, scaling to 256 GPUs and multi-billion-parameter models.
Recommended citation: Kieckhefen, D., Götz, M., Heyen, L. H., Streit, A., & Debus, C. (2026). "Combining Domain and Tensor Parallelism to Train Multi-billion-parameter AI Weather Models." In: Proceedings of the Platform for Advanced Scientific Computing Conference (PASC 26).
Download Paper