Kermt Add Cmim Pretrain

kermt-add-cmim-pretrain

Convert a grover_base checkpoint (legacy original-GROVER grover.encoders.* or modern kermt.encoders.*, with or without vocab heads) into a fully-formed hybrid (cMIM + vocab) checkpoint, then continue pretraining on the user's corpus as hybrid.

This is a thin wrapper: upgrade_to_hybrid.py produces a new ckpt that classifies as model_type: hybrid via check_checkpoint.py, and the rest of the workflow is identical to kermt-continue-pretrain.

Status: experimental. This workflow is functional end-to-end but has not been benchmarked against the manuscript's from-scratch hybrid training (which produces the released checkpoint). Use as an experimental alternative to kermt-pretrain-scratch when you want to extend an existing grover_base checkpoint rather than restart from random init. Validate downstream performance on your own benchmark before relying on the upgraded ckpt for production work.

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/.

Hardware requirements

Same as kermt-continue-pretrain (the cMIM decoder adds parameters but not substantially; VRAM headroom should be fine). The upgrade step itself is fast (~5 s) and CPU-only — only the subsequent continue-pretrain consumes GPU.

When to invoke

• User has a grover_base checkpoint (encoder-only or with vocab heads) and wants to extend it into a hybrid (vocab + cMIM contrastive) pretrain.
• Useful for adding the SMILES-reconstruction contrastive objective to a pretrained encoder without restarting pretraining from scratch (which kermt-pretrain-scratch would do at days-scale).

For continuing an existing hybrid or cmim ckpt: use kermt-continue-pretrain directly. For training a fresh model on a custom corpus: use kermt-pretrain-scratch.

Inputs

Required:

• --ckpt <path> — grover_base ckpt to upgrade. Validated via check_checkpoint.py --mode upgrade_to_hybrid; rejected if the ckpt already has a contrast head or task FFN.
• --csv <path> — pretrain corpus CSV. Same shape as kermt-continue-pretrain's --csv input.

Optional (same as kermt-continue-pretrain):

• --val-csv <path> — separate validation CSV. Without it, prepare_data auto-splits by --val-frac 0.1.
• Training-hyperparameter overrides (--epochs N, --batch-size N, lr triple, --warmup-epochs F, etc.).
• --vocab-loss-weight F / --latent-dim N / --contrastive-temperature F.
• --wandb-project NAME / --wandb-run-name NAME — optional Weights & Biases logging (run name honored only alongside a project). Off by default.
• --gpus 0,2.

Workflow

Let $KERMT_REPO be the path to your kermt repo checkout.

1. Pre-flight: check_system (same as kermt-continue-pretrain step 1).

2. Compute run directory:

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

3. Validate the input ckpt with check_checkpoint --mode upgrade_to_hybrid. Abort on ok: false. The validator rejects ckpts that already have contrast head (suggest kermt-continue-pretrain) or task FFN heads (the ckpt has been finetuned; suggest using the original pretrain checkpoint).

4. Validate the corpus via check_data --mode pretrain. Abort on ok: false.

5. Prepare the data with --mode pretrain — without --vocab-dir. The upgrade builds fresh vocab heads sized to the corpus's vocab, so we want prepare_data to produce a new vocab from the corpus rather than passing through the ckpt's old vocab (which may not even exist for encoder-only legacy grover_base ckpts):

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

   The output manifest has vocab_source: "built_fresh" and includes a smiles_vocab (built from the corpus, needed for the new decoder).

6. Upgrade the ckpt.

   "$SKILL_DIR/scripts/kermt_container.sh" run --ckpt <user-ckpt> --run-dir $RUN_DIR -- \
       "python /skill/scripts/upgrade_to_hybrid.py \\
            --ckpt /ckpt \\
            --prepare-manifest /runs/data/prepare_data.json \\
            --out /runs/upgraded.pt"

   Surface the JSON summary to the user — especially warnings[], which includes any encoder-arch drift notes (e.g. legacy GROVER had two extra act_func_* keys that modern KERMTEmbedding doesn't) and the pretrain_ddp.py --backbone argparse-restriction note if the upgraded ckpt's backbone is anything other than gtrans.

7. Estimate runtime + confirm with the user. Same heuristic as kermt-continue-pretrain (corpus size × epochs × GPU count → wall time).

8. Launch the runner detached.

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

   The runner sees the upgraded ckpt as model_type: hybrid, so it auto-dispatches --pretrain_mode hybrid --vocab_loss_weight 1.0 with smiles_vocab plumbed through.

9. Report to the user with the upgraded ckpt path + the same run.json pointer / log path / tensorboard URL pattern as kermt-continue-pretrain.