NVIDIA/makani

Massively parallel training of machine-learning based weather and climate models

What it solves

Makani is a PyTorch-based library designed for the massively parallel training and inference of machine-learning weather and climate prediction models. It addresses the challenges of training large-scale models on hundreds of GPUs, managing massive datasets (like ERA5), and optimizing GPU memory usage to prevent out-of-memory (OOM) errors.

How it works

Makani provides a modular framework for developing novel weather models, supporting various architectures such as FourCastNet3, SFNO, and ViT. It employs several parallelism strategies to scale training across 100+ GPUs, including spatial model parallelism, ensemble parallelism, and channel parallelism. To optimize performance, it supports automatic mixed precision (AMP), just-in-time (JIT) compilation, and activation checkpointing. Training and inference are fully configurable via .yaml files, and the library includes a comprehensive suite of data processing scripts to handle HDF5 formatted weather data and ensure compatibility with WeatherBench2.

Who it’s for

It is primarily aimed at researchers working on ML-based weather and climate prediction who need to scale their training runs to high-performance computing (HPC) infrastructure.

Highlights

  • Massive Scalability: Supports training on 100+ GPUs using multiple forms of model and data parallelism.
  • Flexible Configuration: Fully configurable through YAML files for models, losses, optimizers, and schedulers.
  • Modular Design: Modular metrics and losses to support custom training and inference recipes.
  • Comprehensive Data Pipeline: Includes tools for HDF5 data conversion, statistics computation, and WeatherBench2 integration.
  • Memory Optimizations: Integrated support for activation checkpointing and specific CUDA allocator configurations to mitigate buffer fragmentation.

Related

  • Project
    sktime/pytorch-forecasting

    A PyTorch-based library for time series forecasting that provides a high-level API and state-of-the-art deep learning architectures like Temporal Fusion Transformers and N-BEATS.

  • Project
    timeseriesAI/tsai

    A state-of-the-art deep learning library built on PyTorch and fastai for time series tasks including classification, regression, and forecasting.

  • Project
    PersiaML/PERSIA

    A PyTorch-based parallel training system designed to scale deep learning recommendation models up to 100 trillion parameters on commodity hardware.

  • Project
    openclimatefix/graph_weather

    A PyTorch library implementing various graph neural networks for weather forecasting and data assimilation, including models like GenCast and Aurora.

  • Project
    Project-MONAI/MONAI

    A PyTorch-based open-source framework for deep learning in healthcare imaging that provides standardized workflows and domain-specific tools for medical data.