Makani: 大规模并行训练基于机器学习的天气和气候模型

Makani（夏威夷语中"风"的意思 🍃🌺）是一个实验性库，旨在使用PyTorch进行基于机器学习的天气和气候模型的研究和开发。Makani用于正在进行的研究。稳定的功能会定期移植到NVIDIA Modulus框架中，这是一个用于训练科学和工程领域物理-ML模型的框架。

概述

Makani由NVIDIA和NERSC的工程师和研究人员发起，用于训练FourCastNet，这是一个基于深度学习的天气预报模型。

Makani是一个研究代码，用于在100多个GPU上大规模并行训练天气和气候预测模型，并推动下一代天气和气候模型的发展。除此之外，Makani还被用于在ERA5数据集上训练球面傅里叶神经算子(SFNO) [1]和自适应傅里叶神经算子(AFNO) [2]。Makani使用PyTorch编写，支持各种形式的模型和数据并行、异步数据加载、未预测通道、自回归训练等功能。

入门

可以通过运行以下命令安装Makani：

git clone git@github.com:NVIDIA/makani.git
cd makani
pip install -e .

训练：

通过调用train.py并传递必要的CLI参数来指定配置文件--yaml_config和配置目标--config来启动训练：

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

:warning: 具有复值权重的架构目前会失败。更多信息请参见已知问题。

Makani支持各种优化，以将大型模型装入GPU内存并实现计算效率高的训练。下表提供了这些功能及相应CLI参数的概述：

功能	CLI参数	选项
自动混合精度	`--amp_mode`	`none`, `fp16`, `bf16`
即时编译	`--jit_mode`	`none`, `script`, `inductor`
CUDA图	`--cuda_graph_mode`	`none`, `fwdbwd`, `step`
激活检查点	`--checkpointing_level`	0,1,2,3
数据并行	`--batch_size`	1,2,3,...
通道并行	`--fin_parallel_size`, `--fout_parallel_size`	1,2,3,...
空间模型并行	`--h_parallel_size`, `--w_parallel_size`	1,2,3,...
多步训练	`--multistep_count`	1,2,3,...

特别是大型模型通过使用这些技术的组合来实现。空间模型并行将模型和数据分割到多个GPU上，从而减少模型的内存占用和IO负载，因为每个进程只需要读取部分数据。可以通过运行以下命令来启动SFNO的典型"大规模"训练：

mpirun -np 256 --allow-run-as-root python -u makani.train --amp_mode=bf16 --cuda_graph_mode=fwdbwd --multistep_count=1 --run_num="ngpu256_sp4" --yaml_config="config/sfnonet.yaml" --config="sfno_linear_73chq_sc3_layers8_edim384_asgl2" --h_parallel_size=4 --w_parallel_size=1 --batch_size=64

在这里，我们在256个GPU上训练模型，水平分割为4个进程，批量大小为64，相当于本地批量大小为1/4。通过使用bf16自动混合精度进一步减少内存需求。

推理：

与训练类似，可以通过从CLI调用inference.py来执行推理，并由inferencer.py处理。要在样本外数据集上启动推理，我们可以调用：

mpirun -np 256 --allow-run-as-root python -u makani.inference --amp_mode=bf16 --cuda_graph_mode=fwdbwd --multistep_count=1 --run_num="ngpu256_sp4" --yaml_config="config/sfnonet.yaml" --config="sfno_linear_73chq_sc3_layers8_edim384_asgl2" --h_parallel_size=4 --w_parallel_size=1 --batch_size=64

默认情况下，推理脚本将在指定的样本外数据集上执行推理。

关于Makani的更多信息

项目结构

项目结构如下：

makani
├── ...
├── config                      # 配置文件，也称为recipes
├── data_process                # 数据预处理，如统计计算
├── datasets                    # 数据集实用工具脚本
├── docker                      # 用于构建训练用Docker镜像的脚本
├── makani                      # 包含主要包的目录
│   ├── inference               # 包含推理器
│   ├── mpu                     # 模型并行化的实用工具
│   ├── networks                # 网络，包含各种ML模型的定义
│   ├── third_party/climt       # 第三方模块
│   │   └── zenith_angle.py     # 计算天顶角
│   ├── utils                   # 实用工具
│   │   ├── dataloaders         # 包含各种数据加载器
│   │   ├── metrics             # 指标文件夹包含评分和基准测试的例程
│   │   ├── ...
│   │   ├── comm.py             # 正交通信器基础设施的通信模块
│   │   ├── dataloader.py       # 数据加载器接口
│   │   ├── metric.py           # 集中的指标处理程序
│   │   ├── trainer_profile.py  # 用于分析的trainer.py副本
│   │   └── trainer.py          # 处理训练的主文件
│   ├── ...
│   ├── inference.py            # 用于启动推理的CLI脚本
│   ├── train.py                # 用于启动训练的CLI脚本
├── tests                       # 测试文件
└── README.md                   # 本文件

模型和训练配置

Makani中的模型训练通过使用位于config文件夹中的.yaml文件来指定。相应的模型位于networks中，并在networks/models.py的get_model例程中注册。下表列出了最重要的配置选项。

配置键	描述	选项
`nettype`	网络架构	`SFNO`, `FNO`, `AFNO`, `ViT`
`loss`	损失函数	`l2`, `geometric l2`, ...
`optimizer`	要使用的优化器	`Adam`, `AdamW`
`lr`	初始学习率	浮点数 > 0.0
`batch_size`	批量大小	整数 > 0
`max_epochs`	要训练的epoch数	整数
`scheduler`	要使用的学习率调度器	`None`, `CosineAnnealing`, `ReduceLROnPlateau`, `StepLR`
`lr_warmup_steps`	学习率调度器的预热步数	整数 >= 0
`weight_decay`	权重衰减	浮点数
`train_data_path`	包含训练数据的目录路径	字符串
`test_data_path`	包含测试数据的目录路径	字符串
`exp_dir`	输出（如模型检查点）的目录路径	字符串
`metadata_json_path`	元数据文件`data.json`的路径	字符串
`channel_names`	用于训练的通道	List[string]

要获得更全面的概述，我们建议查看现有的.yaml配置。有关可用配置的更多详细信息，可以在此文件中找到。

训练数据

Makani期望训练/测试数据采用HDF5格式，每个文件包含一整年的数据。然后，Makani中的数据加载器将加载输入inp和目标tar，它们分别对应于给定时间点和稍后时间点的大气状态。输入和目标之间的时间差由参数dt确定，它决定了两者相隔多少步。实际的物理时间差由数据集的时间分辨率dhours决定。

Makani需要一个名为data.json的元数据文件，该文件描述了数据集的重要属性，如包含数据的HDF5变量名。另一个例子是要在数据加载器中加载的通道，它们通过通道名称指定。元数据文件具有以下结构：

{
    "dataset_name": "给这个数据集一个名称",      # 数据集的名称
    "attrs": {                                 # 可选属性，可以包含任何你想要的内容
        "decription": "数据集的描述",
        "location": "你的数据集的位置"
    },
    "h5_path": "fields",                       # hdf5文件内数据的变量名
    "dims": ["time", "channel", "lat", "lon"], # 数据集中包含的字段的维度
    "dhours": 6,                               # 时间分辨率（小时）
    "coord": {                                 # 坐标和通道描述
        "grid_type": "equiangular",            # 数据集中使用的网格类型：当前支持的选项是'equiangular'和'legendre-gauss'
        "lat": [0.0, 0.1, ...],                # 纬度网格坐标
        "lon": [0.0, 0.1, ...],                # 经度网格坐标
        "channel": ["t2m", "u10", "v10", ...]  # 数据集中包含的通道名称
    }
}

ERA5数据集可以在这里下载。

模型包

默认情况下，Makani在开始训练时会保存一个模型包。模型包方便地包含了运行模型所需的所有必要数据。这包括用于归一化输入和输出的统计数据、未预测的静态通道，甚至是附加天体特征（如太阳天顶角的余弦）的代码。在这里了解更多关于模型包的信息。

已知问题

:warning: 具有复值权重的架构：训练某些具有复值权重的架构需要尚未发布的PyTorch补丁。解决这些问题的临时修复方案可在makani/third_party/torch文件夹中找到。覆盖PyTorch安装中的相应文件将解决这些问题。

贡献

感谢您有兴趣做出贡献。有很多方式可以为这个项目做出贡献。

如果您发现了bug，请告诉我们并提出问题。更好的是，如果您愿意修复它并提出pull request，我们会非常感激。🙏
如果您想添加新功能，我们建议您先与我们讨论，这样我们可以指导您如何最好地实现它。

虽然这是一个研究项目，但我们的目标是拥有功能性的单元测试，并有良好的覆盖率。如果您添加了可以测试的新功能，我们恳请您实现单元测试。

作者

此代码由Thorsten Kurth、Boris Bonev、Jean Kossaifi、Animashree Anandkumar、Kamyar Azizzadenesheli、Noah Brenowitz、Ashesh Chattopadhyay、Yair Cohen、David Hall、Peter Harrington、Pedram Hassanzadeh、Christian Hundt、Alexey Kamenev、Karthik Kashinath、Zongyi Li、Morteza Mardani、Jaideep Pathak、Mike Pritchard、David Pruitt、Sanjeev Raja、Shashank Subramanian开发。

参考文献

<a id="#sfno_paper">[1]</a> Bonev B., Kurth T., Hundt C., Pathak, J., Baust M., Kashinath K., Anandkumar A.; 球面傅里叶神经算子：学习球面上的稳定动力学; arXiv 2306.0383, 2023.

<a id="1">[2]</a> Pathak J., Subramanian S., Harrington P., Raja S., Chattopadhyay A., Mardani M., Kurth T., Hall D., Li Z., Azizzadenesheli K., Hassanzadeh P., Kashinath K., Anandkumar A.; FourCastNet：使用自适应傅里叶神经算子的全球数据驱动高分辨率天气模型; arXiv 2202.11214, 2022.

引用

如果您使用此包，请引用

@InProceedings{bonev2023sfno,
    title={球面傅里叶神经算子：学习球面上的稳定动力学},
    author={Bonev, Boris and Kurth, Thorsten and Hundt, Christian and Pathak, Jaideep and Baust, Maximilian and Kashinath, Karthik and Anandkumar, Anima},
    booktitle={第40届国际机器学习会议论文集},
    pages={2806--2823},
    year={2023},
    volume={202},
    series={机器学习研究论文集},
    month={7月23--29日},
    publisher={PMLR},
}

@article{pathak2022fourcastnet,
    title={FourCastNet：使用自适应傅里叶神经算子的全球数据驱动高分辨率天气模型},
    author={Pathak, Jaideep and Subramanian, Shashank and Harrington, Peter and Raja, Sanjeev and Chattopadhyay, Ashesh and Mardani, Morteza and Kurth, Thorsten and Hall, David and Li, Zongyi and Azizzadenesheli, Kamyar and Hassanzadeh, Pedram and Kashinath, Karthik and Anandkumar, Animashree},
    journal={arXiv预印本 arXiv:2202.11214},
    year={2022}
}