/bionemo-kermt-continue-pretrain

name: kermt-continue-pretrain
description: Continue KERMT pretraining on a custom SMILES corpus with a grover_base, cmim, or hybrid checkpoint. Use a local checkpoint or optionally download a pinned Hugging Face model bundle using HF_TOKEN if configured. Run containerized training and write model bundles, prepared data, logs, and checkpoints 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: this file is targeted at ~250 lines / ~3000 tokens —
# well within the 500-line / 5000-token cap. Long examples live in
# /skill/scripts/run_pretrain_local.py's docstring.

kermt-continue-pretrain

Continue pretraining from a user-supplied KERMT checkpoint (grover_base / cmim / hybrid). The skill is the workflow orchestrator: it validates inputs, prepares the corpus, launches the runner, and returns a run directory.

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–N CUDA-capable NVIDIA GPUs. The runner auto-detects via

`torch.cuda.device_count()`; `--gpus 0,2` overrides. On a single GPU the runner falls back to `--batch_size 32 --save_interval 500`; on multi-GPU it uses the `defaults_pretrain.json` values (currently `batch_size 256`). Note: `--gpus N` uses **torch.cuda** indexing, which can differ from `nvidia-smi`'s display order on multi-GPU hosts (PCI bus vs. CUDA enumeration). To target a specific physical GPU, set `CUDA_VISIBLE_DEVICES` before invoking, or run `python -c "import torch; print([torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())])"` to confirm which device you're picking.

• **VRAM**: the default `--batch-size 256` is sized for A100-class hardware

(80 GB VRAM). On smaller GPUs, downscale to avoid OOM:

| GPU class | VRAM | Suggested `--batch-size` | |--------------------------|------------|--------------------------| | L4, T4, V100 16 GB | 16–24 GB | 32–64 | | A100 40 GB, L40, A40 | 40–48 GB | 128 | | A100 80 GB, H100, H200 | 80 GB | 256 (default) |

These are rough starting points — pass `--batch-size N` to override.

• **Disk**: tens of GB depending on corpus size + epochs (each checkpoint

is several hundred MB).

• **Driver / CUDA**: any host supporting CUDA 12.6 (the kermt image base).

`kermt-setup` validates this up-front.

Inputs

Required:

• `--csv <path>` — the pretrain CSV (single column `smiles`). If you have

separate train/val CSVs, pass `--val-csv <path>` too.

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

• `--ckpt <path>` — the input pretrain checkpoint to continue from. Must be

a grover_base (with vocab heads), cmim, or hybrid ckpt; the validator rejects everything else with a redirect to the correct workflow. **If omitted**, the skill offers to download the released pretrained hybrid model **nvidia/NV-KERMT-70M-v2** and continue-pretrain from it — see "Resolve & validate the checkpoint" (workflow step 3). The released bundle ships its three vocab files alongside the ckpt, so the authoritative-vocab pass-through (step 5) works automatically.

• `--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:

• `--val-csv <path>` — separate validation CSV. Without it, the prep step

auto-splits the input by `--val-frac 0.1` (random shuffle with `--seed`).

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

`--final-lr F` / `--warmup-epochs F` / `--weight-decay F` / `--dropout F` / `--save-interval N` / `--seed N` — training-hyperparameter overrides. Anything not given is filled from `config/defaults_pretrain.json`.

• `--vocab-loss-weight F` (hybrid only) / `--latent-dim N` /

`--contrastive-temperature F` (cmim and hybrid only) — loss / decoder overrides.

• `--wandb-project NAME` / `--wandb-run-name NAME` — optional Weights & Biases

logging. When `--wandb-project` is set, rank 0 logs train/val losses; the run name is honored only alongside a project. Off by default. (Independent of the ckpt's `wandb_run_id` continuity handling under `--resume`.)

• `--resume` — see "Modes" section below.
• `--gpus 0,2` — restrict to a GPU subset. Default uses all visible GPUs.
• `--from-prepare <dir>` — skip the prepare step and reuse an existing

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

Modes

The runner has two modes for ingesting the input ckpt, dispatched on whether `--resume` is set. Pick based on intent:

Default (fresh-schedule continue