Examples¶

The scripts below live in the repository and are included here verbatim.

Minimal training run¶

examples/training/minimal.sh runs debug-single-gpu, a quick check: the HealDA v2 model at reduced depth and time length 1, on one GPU. It needs the ERA5 and NNJA archives prepared.

examples/training/minimal.sh

#!/usr/bin/env bash
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

set -euo pipefail

# Quick check: the HealDA v2 model at reduced depth and time length 1, on one GPU.
# Needs the ERA5 and NNJA data ETL'd (docs/etl.md).
: "${ERA5_HOURLY_ROOT:?set the 0.25 degree ERA5 HDF5 root}"
: "${ERA5_HOURLY_EXT_ROOT:?set the ERA5 extra-variable root}"
: "${NNJA_ROOT:?set the NNJA observation archive root}"

OUTPUT_DIR="${OUTPUT_DIR:-./training-runs}"

torchrun --standalone --nproc-per-node=1 \
  -m healda.cli.train \
  --name debug-single-gpu \
  --output_dir "${OUTPUT_DIR}"

# Run this command again with the same name and output directory to resume the
# newest valid training-state checkpoint automatically.

Launch from a preset file¶

Presets are YAML files in src/healda/configs/loops/, each extending another and listing the fields it changes. --config takes a preset name or a file; a file may extend any shipped preset. --name names the run directory and defaults to the preset or file name. Floats need a decimal point and a signed exponent (2.0e-4, 1.0e+5).

# my-debug.yaml
extends: debug-single-gpu
loop:
  steps_per_tick: 10

torchrun --standalone --nproc-per-node=1 -m healda.cli.train \
  --config my-debug.yaml --output_dir ./training-runs --verbose

--verbose prints the resolved loop with every field, in the same YAML form.

Fine-tune from a checkpoint¶

examples/fine_tuning/from_checkpoint.sh fine-tunes the v2 checkpoint with --finetune_from. Optimizer and iterator state are not restored, so give a fine-tuning preset its own learning-rate schedule.

examples/fine_tuning/from_checkpoint.sh

#!/usr/bin/env bash
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

set -euo pipefail

# Experimental only: this workflow has not been scientifically validated.
# The HealDA v2 configuration runs 8-way time parallel: 8 GPUs, about 150 GB each (README,
# Hardware).
: "${ERA5_HOURLY_ROOT:?set the 0.25 degree ERA5 HDF5 root}"
: "${ERA5_HOURLY_EXT_ROOT:?set the ERA5 extra-variable root}"
: "${NNJA_ROOT:?set the NNJA observation archive root}"
: "${BASE_CHECKPOINT:?set the v2 training-state checkpoint}"

OUTPUT_DIR="${OUTPUT_DIR:-./fine-tuning-runs}"
RUN_NAME="${RUN_NAME:-v2-nnja-latlon-dense}"

torchrun --standalone --nproc-per-node=8 \
  -m healda.cli.train \
  --name "${RUN_NAME}" \
  --output_dir "${OUTPUT_DIR}" \
  --finetune_from "${BASE_CHECKPOINT}"

# If OUTPUT_DIR/RUN_NAME already contains a valid checkpoint, normal resume
# takes precedence and BASE_CHECKPOINT is not loaded.