Scaling GPU-Accelerated AI Models for Advanced Numerical Weather Prediction

AI-driven Numerical Weather Prediction (NWP) digital twins offer significant speed improvements over traditional physics-based models, enabling more frequent updates and longer-range forecasts. However, these AI-NWP models, particularly those utilizing advanced architectures like transformers and graph neural networks, significantly increase GPU utilization and maintain substantial storage requirements. This shift necessitates novel approaches to efficiently utilize high-performance computing (HPC) infrastructure. The Community Runnable Earth Digital Intelligence Twin (CREDIT) framework, developed by the MILES group at NSF NCAR, addresses these challenges by facilitating efficient training, execution, and analysis of AI-NWP models.

This study investigates advanced GPU scaling techniques to optimize AI-NWP model performance within CREDIT and similar platforms. We leverage PyTorch's Fully Sharded Data Parallel (FSDP) framework, implementing memory-efficient strategies including mixed precision training and activation checkpointing. Additionally, our optimization efforts encompass the development of effective shard policies for model distribution across GPUs and the implementation of combined auto-wrap policies tailored for transformer and graph architectures. Furthermore, we examine various parallelization approaches, such as NCCL, MPI, and Gloo, to enhance communication efficiency across multiple GPUs and nodes. Special attention is given to the complexities of NCCL implementation in HPC environments like Derecho, providing solutions for proper configuration and execution.

These technical advancements have resulted in stable and scalable weather predictions, extending to two weeks at 0.28-degree resolution and up to one year at 1-degree resolution using a single model. This presentation will elucidate the technical challenges and solutions involved in scaling AI-NWP models, demonstrating the capacity of CREDIT and similar platforms to fully leverage GPU resources for advanced weather prediction. The findings contribute to the ongoing development of efficient, large-scale AI models for atmospheric science applications.