Makani：大规模并行训练机器学习天气和气候模型

项目简介

Makani，源于夏威夷语中的“风”，是一个专为在PyTorch中研究和开发基于机器学习的天气和气候预测模型而设计的实验性库。它起源于NVIDIA工程师和NERSC研究人员对FourCastNet的训练工作，并已用于前沿的研究项目。稳定的功能会定期移植到NVIDIA Modulus框架中，该框架用于在科学和工程领域训练Physics-ML模型。

技术剖析

Makani支持大规模并行训练，可在100+个GPU上运行，旨在推动下一代天气和气候模型的发展。这个库采用PyTorch编写，提供各种模型并行和数据并行方案，如异步数据加载、不可预见通道、自回归训练等。特别是，它被用来训练如Spherical Fourier Neural Operators (SFNO) 和 Adaptive Fourier Neural Operators (AFNO) 这样的模型。

应用场景

Makani的用途广泛，尤其适合于需要高效处理大量气象和气候数据的应用。它可以用于：

• 实时或短期天气预报
• 长期气候模式模拟
• 大规模环境影响评估
• 灾害预警系统
• 能源管理（如风电场优化）

项目特点

• 高度并行化：支持多达数百个GPU的并行训练，加速模型训练。
• 灵活的架构：适应不同的机器学习模型，包括先进的神经网络架构。
• 自动混合精度：通过自动混合精度训练减少内存需求，提高运算速度。
• 异步数据加载：减少I/O瓶颈，提升训练效率。
• 多种优化策略：支持CUDA图、激活检查点、多步训练等多种优化方式，以适应大型模型。

开始使用

要安装Makani，只需执行以下命令：

git clone [e-mail]:NVIDIA/makani.git
cd makani
pip install -e .

训练模型可以运行train.py，并通过命令行参数指定配置文件和目标：

mpirun -np 8 --allow-run-as-root python -u makani.train --yaml_config="config/sfnonet.yaml" --config="sfno_linear_73chq_sc3_layers8_edim384_asgl2"

对于更复杂的模型，可通过组合上述优化策略来有效利用资源。

结论

Makani是机器学习应用于气象学领域的强大工具，其创新的并行计算能力和高效的训练方法为天气和气候建模带来了革命性的变革。无论是研究人员还是开发者，都可以借助Makani快速构建和训练大规模预测模型，探索地球大气的新理解。立即加入，开启您的天气与气候智能预测之旅吧！