kermt-continue-pretrain

Files (15)

• BENCHMARK.md
• SKILL.md
• defaults_pretrain.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_pretrain_local.py · scripts/
• skill-card.md
• skill.oms.sig

SKILL.md (first lines)

# 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-pretrain)

**Use when**: you have a finished pretrain ckpt and want to continue training
it — on a new corpus, with a different objective, or just for more epochs
than its original plan. The previous training's step counter and schedule
shape are no longer relevant; you want a new learning-rate schedule for the
new run.

**What gets loaded from the ckpt**:
- ✓ Model weights (encoder + vocab heads + contrast head + decoder, whatever
  is there)
- ✓ Optimizer state (Adam's running m1/m2 moments — warm-starts the new
  schedule so the first few hundred steps aren't dominated by noisy
  gradient-estimate startup)
- ✗ Scheduler step counter (reset to 0)
- ✗ Epoch counter (reset to 0)
- ✗ Batch counter (reset to 0)
- ✗ wandb run id (new wandb run, not a continuation)

From NVIDIA/skills, Apache-2.0.

More from NVIDIA/skills