Kermt Continue Pretrain

Workflow

Let $KERMT_REPO be the path to your kermt repo checkout, and assume kermt-setup has already 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); {d[\"disk\"][\"free_gb\"]} GB free; CUDA via container toolkit')
   "

   Surface any gaps to the user. Refuse to proceed if ok: false.

2. Compute run directory.

   RUN_DIR=$KERMT_REPO/runs/continue-pretrain_$(date -u +%Y-%m-%dT%H-%M-%SZ)

3. Resolve & validate the checkpoint.

   Resolve — only if --ckpt was omitted. Default to the released pretrained hybrid model nvidia/NV-KERMT-70M-v2:

   • Consent gate. Unless --pretrained-release was passed, ask the user: "No checkpoint given — download the released model nvidia/NV-KERMT-70M-v2 (NVIDIA Open Model License, https://huggingface.co/nvidia/NV-KERMT-70M-v2) and continue-pretrain from it? [y/N]". Never download without an explicit yes (or --pretrained-release). If both --ckpt and --pretrained-release are given, abort — they conflict.
   • Save location. Default $KERMT_REPO/models/NV-KERMT-70M-v2/; honor --model-dir <dir> if given. An already-complete bundle is reused.
   • Download (foreground; ~282 MB on first fetch):

     "$SKILL_DIR/scripts/kermt_container.sh" run --model-dir <save-dir> -- \
         "python /skill/scripts/fetch_released_model.py --out /model"

     Parse the JSON; abort on ok: false (surface errors). On success set <user-ckpt> = <save-dir>/kermt_contrastive_v2.0.pt. The bundle's three vocab files land in <save-dir> too, so step 5's --vocab-dir auto-detection (which looks in the ckpt's parent directory) finds them with no extra work.

   Validate the resolved (or user-provided) ckpt:

   "$SKILL_DIR/scripts/kermt_container.sh" run --ckpt <user-ckpt> -- \
       "python /skill/scripts/check_checkpoint.py --mode continue_pretrain --ckpt /ckpt"

   Parse the JSON. Abort on ok: false, showing the error verbatim. The error message redirects the user to kermt-add-cmim-pretrain for encoder-only ckpts, or to kermt-finetune for finetuned ckpts.

4. Validate the data.

   "$SKILL_DIR/scripts/kermt_container.sh" run --data <user-csv> -- \
       "python /skill/scripts/check_data.py --mode pretrain --csv /data/<basename>"

   Abort on ok: false.

5. Prepare the data (skip if --from-prepare given). Pass the ckpt's vocab through. Look in the ckpt's parent directory for the conventional pretrain_atom_vocab.{json,pkl}, pretrain_bond_vocab.{json,pkl}, and pretrain_smiles_vocab.pkl files (the bundling convention for released models; see references/released-models.md). If all three are present, auto-pass via --vocab-dir <ckpt_parent_dir>. If only some are present, pass them via explicit flags (--atom-vocab, --bond-vocab, --smiles-vocab). If none are present, ask the user for --vocab-dir — or refuse to proceed, because rebuilding a fresh vocab from the new corpus would silently mismatch the ckpt's vocab heads (the ckpt's vocab is authoritative for continue-pretrain).

   Note the two-layer mount pattern: pass the host directory to kermt_container.sh --vocab-dir (which mounts it at /vocab inside the container), and reference /vocab from the inner prepare_data.py command. The same pattern applies to every host path the inner command needs to read (--data <host-csv> → /data/<basename>, --ckpt <host-ckpt> → /ckpt).

   VOCAB_DIR=$(dirname <user-ckpt>)
   "$SKILL_DIR/scripts/kermt_container.sh" run \
       --data <user-csv> --vocab-dir $VOCAB_DIR --run-dir $RUN_DIR -- \
       "python /skill/scripts/prepare_data.py --mode pretrain \\
            --csv /data/<basename> --out /runs/data \\
            --vocab-dir /vocab \\
            [--val-csv /data/<val-basename>] [--val-frac 0.1] [--seed 0]"

   Outputs land at $RUN_DIR/data/prepare_data.json with vocab_source: "user_provided". The runner step 7 will verify the vocab files' entry counts match the ckpt's vocab-head sizes and refuse to launch on mismatch.

6. Estimate runtime + confirm with user.

   • Pretrain wall time depends on corpus size × epochs × GPU count.
   • Tell the user the estimate; ask "proceed?" unless --yes flag was given (agent-non-interactive case).
   • Example estimate template: ~N hours on K GPUs for E epochs over M molecules (~steps/epoch × seconds/step).

7. Launch the runner detached.

   "$SKILL_DIR/scripts/kermt_container.sh" run_detached \\
       --name kermt-continue-pretrain-<ts> \\
       --ckpt <user-ckpt> --run-dir $RUN_DIR -- \\
       "python /skill/scripts/run_pretrain_local.py \\
            --ckpt /ckpt \\
            --prepare-manifest /runs/data/prepare_data.json \\
            --out /runs \\
            [--epochs N --batch-size N --init-lr F ...]"

   Returns the container name + id + log file path.

8. Report to the user. Output a short summary:

   • Container name + id
   • $RUN_DIR/run.json (the manifest with cmd_replay + image digest)
   • Log file: $RUN_DIR/logs/pretrain_ddp.log
   • TensorBoard: $RUN_DIR/logs/tb (open with tensorboard --logdir $RUN_DIR/logs/tb)
   • Suggest invoking kermt-monitor <RUN_DIR> to check progress.

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 symlinks it into the save_dir; the symlink is what pretrain_ddp.py auto-resumes from. The source file stays untouched.
• Never silently override arch. If the user passes a --hidden-size etc. that doesn't match the ckpt-derived value, the runner aborts loudly. Arch params come from the ckpt, period.
• Never block on the long-running pretrain itself. The runner is invoked via run_detached; the skill returns immediately after step 8. Use kermt-monitor for progress.
• Echo applied defaults back to the user. The args_applied field of run.json records every flag's value + source (user / default-config / auto-1gpu / auto-multi-gpu). Skill should surface a summary of any flag not user-specified so the user knows what was assumed.

Common errors

• model_type='finetuned' rejected → the ckpt is a downstream finetune, not a pretrain. The error redirects to the relevant workflow.
• grover_base ckpt has no vocab head → encoder-only ckpt (e.g. the original-grover grover_base.pt). The error redirects to kermt-add-cmim-pretrain.
• prepare_data manifest is missing required outputs → user passed --from-prepare to a directory where prepare was run with --skip-vocab or --skip-split. Re-run prepare without those flags.
• --gpus all not available → install nvidia-container-toolkit; check kermt_container.sh check_system.

Replayability

The run.json cmd_replay field is a single-line command that re-runs the pretrain 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.