kermt-finetune

Files (15)

• BENCHMARK.md
• SKILL.md
• defaults_finetune.json · config/
• released_model.json · config/
• evals.json · evals/
• released-models.md · references/
• _utils.py · scripts/
• check_checkpoint.py · scripts/
• check_data.py · scripts/
• fetch_released_model.py · scripts/
• kermt_container.sh · scripts/
• prepare_data.py · scripts/
• run_finetune_local.py · scripts/
• skill-card.md
• skill.oms.sig

SKILL.md (first lines)

# 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 `$KERMT_REPO` be the path to your kermt repo checkout, and assume
`kermt-setup` has built `kermt:latest`. All paths below are on the host; the
helper bind-mounts them at known container paths.

1. **Pre-flight: ensure container + system probe.**
   ```
   "$SKILL_DIR/scripts/kermt_container.sh" check_system | python -c "
   import json, sys; d = json.load(sys.stdin)
   if not d['ok']:
       print('System check failed:', d['gaps']); sys.exit(1)
   print(f'OK: {len(d[\"gpus\"])} GPU(s); CUDA via container toolkit')
   "
   ```
   Refuse to proceed if `ok: false`.

2. **Compute run directory.**

From NVIDIA/skills, Apache-2.0.

More from NVIDIA/skills