Skip to content

Training

What a run needs

  • a curated corpus directory (train.smi, val.smi, test.smi, vocab.json, corpus_meta.json) produced by pic2smiles corpus
  • PyTorch with a working accelerator (CUDA or Apple MPS; CPU works but is slow)
  • CPU cores for rendering. Data is generated on the fly, so render workers are part of the hardware requirement, not an optimization. A rule of thumb: one worker per ~40 images/second of target throughput.

Sizing a run

The unit of progress is optimizer steps; the sample budget is steps x batch_size. Measure before committing:

pic2smiles train --corpus data/corpus/chembl --output runs/probe \
  --preset base --steps 300 --batch-size 64 --eval-every 0 --checkpoint-every 0

The img/s in the log is the number that matters. Total wall clock is steps * batch_size / img_per_second.

Local (Apple Silicon)

Use --preset tiny or small, and keep --amp none: measured bf16 gain on MPS is about 17%, which does not justify autocast numerics debugging.

pic2smiles train --corpus data/corpus/chembl --output runs/molmini_small \
  --preset small --steps 20000 --batch-size 24 --num-workers 5 --device mps

Cluster (Slurm, one GPU)

sbatch --gres=gpu:RTX_6000_Ada:1 --nodelist=peacock06 \
  scripts/slurm/train_a100.sbatch <corpus_dir> <output_dir> <steps> <batch> <preset>

Use the RTX 6000 Ada, not an A100: measured 12-13% faster end-to-end on this model, and it is the card nobody else is queuing for. The reasoning and the numbers are in D19; re-measure with scripts/slurm/probe_roofline.sbatch if the model size changes materially, because the result turns on the model being small enough that the A100's scale goes unused.

The job requests 16 CPUs for 14 render workers. Below ~10 workers the run becomes data-bound and the card stops mattering. --amp bf16 is on: unlike MPS, CUDA bf16 goes through tensor cores and is worth 1.9-2.2x.

Sync code to the cluster with:

scripts/sync_to_remote.sh peacock ~/pic2smiles/repo

Resuming

Checkpoints land in <output>/last.pt and <output>/best.pt (best held-out exact match). Resume with --resume runs/.../last.pt. Because the schedule is WSD, a run extended past its original budget only changes where the decay starts; it does not invalidate the earlier steps.

Reading the log

<output>/train_log.jsonl has one record per logging interval:

  • loss, token_accuracy -- teacher-forced, so they move early and saturate long before exact match does. Token accuracy above 0.9 with exact match near zero is normal mid-run: one wrong token invalidates a whole molecule.
  • grad_norm -- should settle after warmup. Sustained growth means the LR is too high for the current width.
  • images_per_second -- if this drops when num_workers is raised, rendering is contending for cores with the training process.
  • eval records -- greedy exact match and validity on held-out molecules. This is the number that decides best.pt.

Failure modes worth recognizing

  • Exact match stuck at zero while validity is high. The model has learned SMILES grammar but not the image. Expected early; a real problem if it persists past the first few million samples.
  • Validity high, Tanimoto high, exact match low. Predictions are chemically close but not identical -- usually stereo or a substituent position. Check the exact vs exact/no-stereo gap in the benchmark report.
  • Throughput falling over time. Render workers are being starved, or a memory leak in the worker pool. Compare images_per_second early and late.