Prerequisites for NVIDIA Earth-2 Correction Diffusion NIM#

The following are the prerequisites for using NVIDIA Earth-2 Correction Diffusion (CorrDiff) NIM.

Support Matrix#

Models#

Model Name

Publisher

CorrDiff US GEFS - HRRR

NVIDIA

Hardware Support#

Optimized configuration#

The GPU Memory and Disk Space values are in Gb; Disk Space is for both the container and the model; Profile is for what the model is optimized.

GPU

GPU Memory (GB)

Precision

# of GPUs

Disk Space

H100

80

FP32

1

64

A100

40 & 80

FP32

1

64

RTX A6000

48

FP32

1

64

L40S

48

FP32

1

64

Non-optimized configuration#

The GPU Memory and Disk Space values are in Gb; Disk Space is for both the container and the model.

GPU

GPU Memory (GB)

Precision

# of GPUs

Disk Space

Any NVIDIA GPU with sufficient GPU memory and compute capability ≥ 8.0

40

FP32

1

64

For models with lower GPU memory, some additional configuration may be required during runtime. Namely, the adjustment of the target sample batch size of the NIM. See the configuration documentation for additional details.

Software#

  • With a Docker-supported operating system, install Docker - minimum version: 23.0.1

  • Install NVIDIA Drivers - minimum version: 545. However, if you are running on a data center GPU (for example, T4 or any other data center GPU), you can use NVIDIA driver release 470.57 (or later R470), 525.85 (or later R525), 535.86 (or later R535), or 545.23 (or later R545).

  • Install and configure the NVIDIA Container Toolkit - minimum version: 1.13.5

  • Verify your container runtime supports NVIDIA GPUs by running

docker run --rm --runtime=nvidia --gpus all ubuntu nvidia-smi

Example output:

+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 560.35.03              Driver Version: 560.35.03      CUDA Version: 12.6     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA A100-PCIE-40GB          Off |   00000000:41:00.0 Off |                    0 |
| N/A   37C    P0             35W /  250W |   40423MiB /  40960MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
                                                                                         

For more information on enumerating multi-GPU systems, please see the NVIDIA Container Toolkit’s GPU Enumeration Docs.

NGC Account#

You must have an authenticated NGC (NVIDIA GPU Cloud) account with access to both the respective NIM and model checkpoints. Use the following procedure to log in to NGC, and set the NGC_API_KEY environment variable, so that you can pull images.

  1. Create an account on NGC.

  2. Generate an API Key. The following steps require your NGC API key.

  3. Authenticate local Docker with NGC by running the following code. For more details, see the NGC authentication documentation.

    docker login nvcr.io
    Username: $oauthtoken
    Password: <NGC API key>
    
  4. Set the NGC API key environment variable in your shell. You need to set the NGC_API_KEY variable.

    export NGC_API_KEY=<NGC API key>
    

Client Side Python Dependencies#

This NIM requires an input containing several weather fields at the desired time-stamp in the form of a numpy array. This input data can be easily fetched using the Earth2Studio package as shown in the quickstart guide. For this, setup an isolated Python environment (Python ≥ 3.10, < 3.12) and install the required dependencies using:

pip install --upgrade pip
pip install earth2studio

Important

This NIM requires Earth2Studio version >= 0.3.0

In the provided guide, several other packages will be used to interact with the NIM and post process the results:

pip install requests
pip install matplotlib

It is recommended to create an empty working directory for client code / input data:

mkdir corrdiff
cd corrdiff

Next Steps#

After you install all prerequisites, to get started, see the quickstart guide.