| name | kermt-finetune |
| description | Finetune a pretrained KERMT encoder on a labeled CSV. The skill validates the input checkpoint (must be a pretrain ckpt — grover_base / cmim / hybrid), validates the labeled CSV, prepares the data (clean + features + optional split), then launches main.py finetune inside the kermt container (detached for hours-scale runs). Hyperparameters come from agent/config/defaults_finetune.json with per-flag CLI override. |
| 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"} |
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.
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 agent/config/defaults_finetune.json.
-
--ffn-hidden-size N / — shared FFN trunk dims.
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.
-
Pre-flight: ensure container + system probe.
$KERMT_REPO/agent/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.
-
Compute run directory.
RUN_DIR=$KERMT_REPO/runs/finetune_$(date -u +%Y-%m-%dT%H-%M-%SZ)
-
Resolve & validate the checkpoint.
Resolve — only if --ckpt was omitted. Default to the released
pretrained hybrid model nvidia/NV-KERMT-70M-v2:
Validate the resolved (or user-provided) ckpt:
$KERMT_REPO/agent/scripts/kermt_container.sh run --ckpt <user-ckpt> -- \
"python agent/scripts/check_checkpoint.py --mode finetune_init --ckpt /ckpt"
Parse the JSON. Abort on ok: false. The validator rejects already-
finetuned ckpts (has_task_ffn: true) with a redirect to kermt-infer.
-
Validate the data.
$KERMT_REPO/agent/scripts/kermt_container.sh run --data <user-csv> -- \
"python agent/scripts/check_data.py --mode finetune --csv /data/<basename> [--targets COL1 COL2 ...]"
Hard rules
- Never download the released model without consent. When
--ckpt is
omitted, download nvidia/NV-KERMT-70M-v2 only after an explicit user "yes"
or an explicit --pretrained-release flag. --ckpt and
--pretrained-release are mutually exclusive.
- Never modify the user's input ckpt. The runner passes its path via
--checkpoint_path; task/train.py loads it read-only into the model and
attaches a new FFN head. The source file stays untouched.
- Arch comes from the ckpt, not from CLI/defaults. The runner extracts
hidden_size, depth, num_attn_head, activation, embedding_output_type,
self_attention (+ attn_hidden / attn_out when applicable) from the
ckpt's saved_args. There is no --hidden-size flag on this runner.
- Never block on the long-running finetune. The skill launches via
run_detached and returns immediately after step 9. Use kermt-monitor.
- Echo applied defaults back to the user. The
args_applied field of
run.json records every flag's value + source (user / default-config).
Surface a one-line summary of every filled-from-default flag so the user
knows what was assumed.
Common errors
finetune_init requires a pretrain ckpt (grover_base / cmim / hybrid) →
the ckpt you passed is already finetuned (has task FFN heads). Pick a
pretrain ckpt instead, or use kermt-infer if you want to run
predictions with the existing finetuned model. To resume a finetune on
the SAME dataset, bypass the skill and call
python main.py finetune --checkpoint_path <ckpt> ... directly — the
agent skill doesn't support resume because saved-task identity
can't be machine-verified against the new training data.
prepare_data manifest reports ok=False → check errors for the failed
step (typically clean_smiles or save_features). Fix and re-run.
ffn_num_task_specific_layers=N>0 but ffn_task_specific_hidden_size is unset
→ MTL heads need an explicit hidden size. Pass --ffn-task-specific-hidden-size H.
finetune is single-GPU (from --gpus 0,1) → --gpus selects one device
for single-process finetune. For multi-GPU, use --num-gpus N (DDP) instead.
Replayability
The run.json cmd_replay field is a single-line command that re-runs the
finetune with the same inputs, hyperparameters, and arch. To replay inside
the kermt container:
$(jq -r .cmd_replay $RUN_DIR/run.json)
If ok_to_replay: false in the manifest (because the kermt repo working
tree was dirty at launch time), the replay may not be bit-exact — pin the
exact commit via the repo.commit field and git checkout it
first.