在运行时配置 NVIDIA Earth-2 FourCastNet NIM#

使用本文档详细了解如何在运行时配置 NVIDIA Earth-2 FourCastNet NIM。

GPU 选择#

在具有 1 个或多个相同 GPU 的同构环境中，将 --gpus all 传递给 docker run 是可以接受的。在某些环境中，在特定 GPU 上运行容器是有益的。通过以下任一方式在容器内公开特定 GPU

  • --gpus 标志，例如 --gpus='"device=1"'。

  • 环境变量 NVIDIA_VISIBLE_DEVICES，例如 -e NVIDIA_VISIBLE_DEVICES=1。

用作输入的设备 ID 是 nvidia-smi -L 的输出。

GPU 0: Tesla H100 (UUID: GPU-b404a1a1-d532-5b5c-20bc-b34e37f3ac46)
GPU 1: NVIDIA GeForce RTX 3080 (UUID: GPU-b404a1a1-d532-5b5c-20bc-b34e37f3ac46)

请参阅 NVIDIA Container Toolkit 文档以获取更多说明。

共享内存标志#

FourCastNet NIM 使用 Triton 的 Python 后端 功能，该功能可随着可用 CPU 核心数量的增加而扩展。您可能需要增加提供给微服务容器的共享内存。

提供 4g 共享内存的示例

docker run ... --shm-size 4g ...

模型配置文件#

FourCastNet NIM 具有以下可以使用的模型配置文件

SFNO ERA5 73 通道 微调#

NIM_MODEL_PROFILE: eb0719154e9c48106f33177da4913ace10a2ac03d18424d0f2ba37542b719140

用于预测大气动力学的球谐傅里叶神经算子模型，在 ERA5 的 73 个变量上训练。针对确定性精度进行了微调。

SFNO ERA5 73 通道#

NIM_MODEL_PROFILE: ef1d33ffd6c1f0ba4acaf96801b79564d943481c36445625cb85cc9fa9916f12

用于预测大气动力学的球谐傅里叶神经算子模型，在 ERA5 的 73 个变量上训练。未使用微调过程。

环境变量#

FourCastNet NIM 允许在容器启动时引用的一些自定义设置。以下变量可用于更改 NIM 行为

变量

默认值

描述

NGC_API_KEY

您的 NGC API 密钥，具有对您正在使用的模型配置文件的模型注册表的读取权限。

NIM_MODEL_PROFILE

9f9…c23a

启动时要加载到 NIM 中的模型包。 假设您具有正确的权限，这将从 NGC 下载。

NIM_HTTP_API_PORT

8000

将 NIM 服务发布到容器内指定的端口。 确保调整传递给 docker run 的 -p/--publish 标志的端口以反映这一点。

NIM_DISABLE_MODEL_DOWNLOAD

在容器启动时禁用模型下载。

挂载卷#

可以挂载容器内的以下路径以增强 NIM 的运行时

容器路径

是否必需

描述

示例

/opt/nim/.cache

否

这是模型在容器内下载的目录。 此目录必须可以从容器内部访问。 这可以通过在 docker run 命令中添加选项 -u $(id -u) 来实现。

-v ~/.cache/nim:/opt/nim/.cache
在本页中