kermt-pretrain-scratch

kermt-pretrain-scratch

Pretrain a brand-new KERMT model from scratch on a user-provided corpus. Useful when you want to retrain a model on a custom chemistry domain rather than extending one of the released checkpoints. Significantly more expensive than kermt-continue-pretrain — no warm start, so the loss curves need to descend from scratch over many epochs.

Hardware requirements

Same as kermt-continue-pretrain:

• GPUs: 1–N CUDA-capable. The runner auto-detects via torch.cuda.device_count(); --gpus 0,2 overrides. Single-GPU fallback: --batch_size 32 --save_interval 500. Multi-GPU keeps defaults (--batch_size 256 etc.). 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 for shards + vocab + checkpoints, scaled by epochs.

• Wall time: this is the big difference. Pretraining from scratch on an 11M-mol corpus at 100 epochs typically takes days even on a multi-GPU box. The skill prints an estimate before launching; confirm with the user.

When to invoke

• User wants to train a new model on a custom corpus (e.g. domain-specific chemistry that the released ckpts don't cover).
• User wants to reproduce a pretrain config end-to-end without depending on a released ckpt.

For continuing an existing released ckpt, use kermt-continue-pretrain. For adding a cMIM decoder to an encoder-only grover_base ckpt, use kermt-add-cmim-pretrain.

Inputs

Required:

• --csv <path> — the pretrain corpus CSV with a smiles column. Single file by convention; multi-file corpora deferred. Use --val-csv for a separate validation set.
• --pretrain-target-mode {vocab|cmim|hybrid} — which pretrain objective to use. No default — must be set explicitly so the user makes an informed choice:
  • vocab — original GROVER-style atom + bond vocab prediction (encoder-only output, lightweight).
  • cmim — contrastive + SMILES reconstruction objective. Requires building a SMILES vocab from the corpus.
  • hybrid — both vocab and contrastive objectives jointly (the state-of-the-art config from the KERMT manuscript).

Optional:

• --val-csv <path> — separate validation CSV. Without it, prepare_data auto-splits the input by --val-frac 0.1 (random shuffle with --seed).
• Training-hyperparameter overrides: --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. Anything not given is filled from agent/config/defaults_pretrain.json.
• --vocab-loss-weight F (hybrid only) / --latent-dim N / --contrastive-temperature F (cmim and hybrid only).
• --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.
• --gpus 0,2 — restrict to a GPU subset.

Workflow

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

1. Pre-flight: ensure container + system probe (same as kermt-continue-pretrain step 1). Refuse to proceed if check_system reports gaps.

2. Compute run directory.

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

3. Validate the corpus (no ckpt to validate, so this is the only input check):

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

   Abort on ok: false.

4. Prepare the data — no vocab pass-through (we want fresh vocab from corpus):

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

   Outputs land at $RUN_DIR/data/prepare_data.json with vocab_source: "built_fresh".

5. Estimate runtime + warn loudly. This is critical for pretrain-from-scratch:

   • "Pretraining from scratch is days-scale even on multi-GPU; the released KERMT checkpoints were each trained on millions of molecules for hundreds of GPU-hours. If you mainly want to leverage existing knowledge for a downstream task, consider kermt-continue-pretrain from a released ckpt instead, which converges in hours instead of days."
   • Show the corpus size × epochs × GPU count → estimated wall time.
   • Ask for explicit confirmation unless --yes was given.

6. Launch the runner detached.

   $KERMT_REPO/agent/scripts/kermt_container.sh run_detached \\
       --name kermt-pretrain-scratch-<ts> \\
       --run-dir $RUN_DIR -- \\
       "python agent/scripts/run_pretrain_local.py \\
            --from-scratch --pretrain-target-mode <vocab|cmim|hybrid> \\
            --prepare-manifest /runs/data/prepare_data.json \\
            --out /runs \\
            [--epochs N --batch-size N ...]"
   

   Note: NO --ckpt flag (the runner refuses if both --from-scratch and --ckpt are given). The runner uses the arch group from agent/config/defaults_pretrain.json to size the model.

7. Report to the user. Always include all of the following — do not omit the TensorBoard line under output-length pressure:

   • Container name + id
   • $RUN_DIR/run.json (the manifest with workflow: pretrain-scratch, from_scratch: true, vocab_check: null, arch from defaults, full cmd_replay)
   • Log file: $RUN_DIR/logs/pretrain_ddp.log
   • TensorBoard: $RUN_DIR/logs/tb (open with tensorboard --logdir $RUN_DIR/logs/tb)
   • Suggest kermt-monitor <RUN_DIR> for progress.

Hard rules

• Never accept a --ckpt flag. From-scratch is exclusive with input ckpt — the runner enforces this; the skill should too.
• Never silently default --pretrain-target-mode. This is a significant architectural choice (vocab = lightweight, hybrid = SOTA). Prompt the user if not given on the CLI.
• Strong warning before launching. From-scratch pretrain is the most expensive workflow. The user needs to know what they're committing to.

Common errors

• --pretrain-target-mode is required when --from-scratch is set → user forgot the mode flag. Prompt.
• --from-scratch is incompatible with --ckpt → user provided both; ask which one they meant.
• defaults_pretrain.json has no arch group → repo state issue (should never happen on a fresh clone); points the user at running kermt-setup again.

What's in the manifest after a from-scratch run

Same reproducibility fields as continue-pretrain (repo.commit, kermt_image, cmd_replay, args_applied), plus:

• workflow: "pretrain-scratch"
• from_scratch: true
• inputs.ckpt: null
• ckpt_symlink: null
• vocab_check: null (not verified — vocab built from corpus is authoritative for from-scratch)
• arch: the values pulled from agent/config/defaults_pretrain.json's arch group (with any future CLI overrides applied).

Replayability

Same as continue-pretrain: cmd_replay is a copy-pasteable command. If ok_to_replay: false, the kermt repo working tree was dirty at launch time — check repo.commit and git checkout it first.