/bionemo-kermt-finetune

name: kermt-finetune
description: Finetune a pretrained KERMT encoder on a labeled CSV. Validate the checkpoint and data, prepare features, and run containerized training. Use a local checkpoint or optionally download a pinned Hugging Face model bundle using HF_TOKEN if configured. Write model bundles, prepared data, logs, and trained models to user-selected host directories.
license: Apache-2.0
compatibility: Requires docker, nvidia-container-toolkit, and a CUDA-capable NVIDIA GPU. Designed for Claude Code, Codex, and Nemotron.
metadata:
  owner: evax@nvidia.com
  classification: workflow-skill
  risk_tier: skill
# Line/token budget: targets ~250 lines / ~3000 tokens — within the
# 500-line / 5000-token cap for skill files.

kermt-finetune

Finetune a pretrained KERMT encoder on a user-supplied labeled CSV. The skill is the workflow orchestrator: validate ckpt, validate data, prepare data, launch the runner detached, return a run directory + container name.

Skill and runtime paths

Set `SKILL_DIR` to the absolute path of this installed skill directory. Export `KERMT_REPO` as the absolute path to the KERMT checkout used for model execution. The bundled container helper mounts that checkout at `/workspace` and this skill at `/skill` (read-only). Commands inside the container use `/skill/scripts/`; defaults are bundled in `config/`. See [Released models](references/released-models.md) for checkpoint bundle requirements.

Downloads and local outputs

The optional released-model branch reads `config/released_model.json` for the Hugging Face repository, pinned revision, and filenames. The bundled `scripts/fetch_released_model.py` downloads the model bundle over HTTPS into the host directory the user selects. Public models work without credentials; if `HF_TOKEN` is set, the container helper forwards it for Hugging Face authentication. Prepared data, logs, and workflow results go into the chosen run directory.

Hardware requirements

• **GPUs**: 1 by default (single-GPU); pass `--gpus 0` (or whichever id) to

select one. For faster training on a multi-GPU host, pass `--num-gpus N` (N>1) to run data-parallel DDP across N GPUs — `--batch-size` is then per-GPU (effective global batch = batch_size × N).

• **VRAM**: ≥ 8 GB for the default `batch_size 32` configuration. Lower VRAM

works at smaller batch sizes — pass `--batch-size N` to override.

• **Disk**: a few GB per run (checkpoint + features + logs).
• **Driver / CUDA**: any host supporting CUDA 12.6 (the kermt image base).

`kermt-setup` validates this up-front.

Inputs

Required:

• `--csv <path>` — labeled CSV. First column is `smiles`; every other column

is a target.

Checkpoint (optional — defaults to the released model if omitted):

• `--ckpt <path>` — input pretrain checkpoint (grover_base / cmim / hybrid).

The validator refuses already-finetuned ckpts with a redirect to `kermt-infer`. **If omitted**, the skill offers to download the released pretrained hybrid model **nvidia/NV-KERMT-70M-v2** and finetune from it — see "Resolve & validate the checkpoint" (workflow step 3).

• `--pretrained-release` — explicit opt-in to use the released model without

the interactive prompt (for non-interactive / agent runs). Mutually exclusive with `--ckpt`.

• `--model-dir <dir>` — where to save the downloaded bundle (default

`$KERMT_REPO/models/NV-KERMT-70M-v2/`). An already-complete bundle there is reused, not re-downloaded.

Optional:

• `--dataset-type {regression | classification | multiclass}` — default

`regression` (from `defaults_finetune.json`). Drives loss, metric defaults, and head initialization. For classification tasks pass `--dataset-type classification`.

• `--targets COL [COL ...]` — explicit target column names. If omitted, the

validator auto-detects numeric non-smiles columns and the skill confirms with the user before proceeding.

• `--val-csv <path>` and `--test-csv <path>` — user-provided val + test

splits. Either pass both or pass neither (the skill auto-splits using the configured `--split-type`).

• `--split-type {random | scaffold_balanced | index_predetermined}` —

default `scaffold_balanced` from `defaults_finetune.json`.

• `random` and `scaffold_balanced`: build the val/test split internally

from the train CSV. No `--val-csv` / `--test-csv` needed.

• `index_predetermined`: **requires** pre-split CSVs passed via

`--val-csv` + `--test-csv` (and, separately, per-fold index files — see `kermt/util/utils.split_data`). Use this when the dataset ships its own canonical split (e.g. `tests/data/Biogen_for_grover/scaffold/ balance/<endpoint>/{train,val,test}.csv`).

• `--metric NAME` — `mae` (regression default), `auc` (classification default),

or any name `kermt.util.metrics.get_metric_func` accepts.

• `--epochs N` / `--batch-size N` / `--init-lr F` / `--max-lr F` /

`--final-lr F` / `--warmup-epochs F` / `--weight-decay F` / `--dropout F` / `--bond-drop-rate F` / `--dist-coff F` / `--early-stop-epoch N` / `--seed N` — training-hyperparameter overrides. Anything not given is filled from `config/defaults_finetune.json`.

• `--ffn-hidden-size N` / `--ffn-num-layers N` — shared FFN trunk dims.
• `--ffn-num-task-specific-layers N` / `--ffn-task-specific-hidden-size H` —

per-target FFN heads (default 0 = off; useful for heterogeneous multi-target finetunes). Both must be set together when N > 0.

• `--ensemble-size N` / `--num-folds N` — multi-model / k-fold CV. Default 1

each.

• `--gpus 0` — single GPU id for single-process finetune (default 0). Ignored

when `--num-gpus > 1`.

• `--num-gpus N` — number of GPUs for data-parallel DDP finetune. Default 1

(single-process, unchanged). N>1 runs `main.py finetune` with `WORLD_SIZE=N` (one process per GPU); `--batch-size` is per-GPU.

• `--from-prepare <dir>` — skip the prepare step and reuse an existing

`prepare_data.json` in `<dir>`. Useful when iterating on hyperparameters.

Workflow

Let `$