kermt-continue-pretrain

Par nvidia · skills

Continuez le pré-entraînement KERMT sur un corpus SMILES personnalisé à partir d'un checkpoint grover_base, cmim ou hybrid. Utilisez un checkpoint local ou téléchargez optionnellement un bundle de modèles épinglé depuis Hugging Face via HF_TOKEN si configuré. Exécutez l'entraînement dans un conteneur et écrivez les bundles de modèles, les données préparées, les logs et les checkpoints dans les répertoires hôtes sélectionnés par l'utilisateur.

npx skills add https://github.com/nvidia/skills --skill kermt-continue-pretrain

kermt-continue-pretrain

Continue-pretraining à partir d'un checkpoint KERMT fourni par l'utilisateur (grover_base / cmim / hybrid). La skill est l'orchestrateur du workflow : elle valide les entrées, prépare le corpus, lance le runner et retourne un répertoire d'exécution.

Chemins de la skill et du runtime

Définissez SKILL_DIR sur le chemin absolu du répertoire de cette skill installée. Exportez KERMT_REPO comme le chemin absolu du checkout KERMT utilisé pour l'exécution du modèle. L'assistant conteneur fourni monte ce checkout à /workspace et cette skill à /skill (lecture seule). Les commandes à l'intérieur du conteneur utilisent /skill/scripts/ ; les valeurs par défaut sont regroupées dans config/. Voir Released models pour les exigences du bundle de checkpoint.

Téléchargements et sorties locales

La branche released-model optionnelle lit config/released_model.json pour le repository Hugging Face, la révision épinglée et les noms de fichiers. Le scripts/fetch_released_model.py fourni télécharge le bundle du modèle sur HTTPS dans le répertoire de l'hôte que l'utilisateur sélectionne. Les modèles publics fonctionnent sans identifiants ; si HF_TOKEN est défini, l'assistant conteneur le transmet pour l'authentification Hugging Face. Les données préparées, les logs et les résultats du workflow vont dans le répertoire d'exécution choisi.

Exigences matérielles

  • GPUs : 1–N GPU NVIDIA compatibles CUDA. Le runner auto-détecte via torch.cuda.device_count() ; --gpus 0,2 remplace. Sur un seul GPU, le runner revient à --batch_size 32 --save_interval 500 ; sur multi-GPU, il utilise les valeurs de defaults_pretrain.json (actuellement batch_size 256). Remarque : --gpus N utilise l'indexation torch.cuda, qui peut différer de l'ordre d'affichage de nvidia-smi sur les hôtes multi-GPU (bus PCI vs énumération CUDA). Pour cibler un GPU physique spécifique, définissez CUDA_VISIBLE_DEVICES avant l'invocation, ou exécutez python -c "import torch; print([torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())])"pour confirmer quel appareil vous sélectionnez.

  • VRAM : le --batch-size 256 par défaut est dimensionné pour le matériel de classe A100 (80 GB VRAM). Sur les GPU plus petits, réduisez l'échelle pour éviter OOM :

    Classe GPU VRAM --batch-size suggéré
    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)

    Ce sont des points de départ approximatifs — passez --batch-size N pour remplacer.

  • Disque : des dizaines de GB selon la taille du corpus + epochs (chaque checkpoint est plusieurs centaines de MB).

  • Driver / CUDA : tout hôte supportant CUDA 12.6 (la base d'image kermt). kermt-setup valide ceci d'avance.

Entrées

Requis :

  • --csv <path> — le CSV de pretraining (colonne unique smiles). Si vous avez des CSV train/val séparés, passez également --val-csv <path>.

Checkpoint (optionnel — par défaut au modèle released si omis) :

  • --ckpt <path> — le checkpoint de pretraining d'entrée à continuer. Doit être un ckpt grover_base (avec vocab heads), cmim ou hybrid ; le validateur rejette tout le reste avec une redirection vers le workflow correct. Si omis, la skill propose de télécharger le modèle pretrained released hybrid nvidia/NV-KERMT-70M-v2 et continuer-pretrain à partir de celui-ci — voir « Résoudre et valider le checkpoint » (étape 3 du workflow). Le bundle released embarque ses trois fichiers vocab à côté du ckpt, donc le passage authoritative-vocab (étape 5) fonctionne automatiquement.
  • --pretrained-release — opt-in explicite pour utiliser le modèle released sans l'invite interactive (pour les exécutions non-interactive / agent). Mutuellement exclusif avec --ckpt.
  • --model-dir <dir> — où sauvegarder le bundle téléchargé (par défaut $KERMT_REPO/models/NV-KERMT-70M-v2/). Un bundle déjà complet là-bas est réutilisé, non re-téléchargé.

Optionnel :

  • --val-csv <path> — CSV de validation séparé. Sans lui, l'étape prep auto-divise l'entrée par --val-frac 0.1 (mélange aléatoire avec --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 — remplacements d'hyperparamètres d'entraînement. Tout ce qui n'est pas donné est rempli depuis config/defaults_pretrain.json.
  • --vocab-loss-weight F (hybrid uniquement) / --latent-dim N / --contrastive-temperature F (cmim et hybrid uniquement) — remplacements de loss / decoder.
  • --wandb-project NAME / --wandb-run-name NAME — logging Weights & Biases optionnel. Quand --wandb-project est défini, le rang 0 enregistre les losses train/val ; le nom de run est honoré seulement aux côtés d'un project. Désactivé par défaut. (Indépendant de la gestion de continuité wandb_run_id du ckpt sous --resume.)
  • --resume — voir la section « Modes » ci-dessous.
  • --gpus 0,2 — restreindre à un sous-ensemble de GPU. Par défaut, utilise tous les GPU visibles.
  • --from-prepare <dir> — sauter l'étape prepare et réutiliser un prepare_data.json existant dans <dir>. Utile lors de l'itération sur les hyperparamètres.

Modes

Le runner a deux modes pour ingérer le ckpt d'entrée, dispatché selon si --resume est défini. Choisissez selon l'intention :

Par défaut (fresh-schedule continue-pretrain)

Utilisez quand : vous avez un ckpt pretrain terminé et voulez continuer son entraînement — sur un nouveau corpus, avec un objectif différent, ou juste pour plus d'epochs que son plan original. Le compteur d'étape et la forme de l'emploi du temps du pretraining précédent ne sont plus pertinents ; vous voulez un nouvel emploi du temps de taux d'apprentissage pour la nouvelle exécution.

Ce qui est chargé depuis le ckpt :

  • ✓ Poids du modèle (encoder + vocab heads + contrast head + decoder, tout ce qui est là)
  • ✓ État de l'optimizer (moments Adam m1/m2 en cours — warm-start du nouvel emploi du temps pour que les premiers centaines d'étapes ne soient pas dominées par le démarrage bruyant de l'estimation de gradient)
  • ✗ Compteur d'étape du scheduler (réinitialisé à 0)
  • ✗ Compteur d'epoch (réinitialisé à 0)
  • ✗ Compteur de batch (réinitialisé à 0)
  • ✗ Identifiant wandb run (nouvel exécution wandb, pas une continuation)

Forme de l'emploi du temps (init/max/final LR, warmup epochs, total epochs) : depuis vos args CLI ou defaults_pretrain.json. Un nouveau NoamLR fresh est construit à partir de ces valeurs et commence à l'étape 0.

--resume (true resume)

Utilisez quand : une exécution précédente a été interrompue (crash, OOM, Ctrl-C) et vous voulez reprendre exactement où elle s'est arrêtée — même dataset, même emploi du temps, même trajectoire d'entraînement.

Ce qui est chargé depuis le ckpt : tout dans le format save_model_for_restart. Les poids du modèle + l'état de l'optimizer + scheduler_step + epoch + batch_idx + wandb_run_id sont tous restaurés. La nouvelle exécution continue à partir de l'étape sauvegardée dans l'emploi du temps sauvegardé (qui est récupéré depuis saved_args du ckpt). La reprise mid-epoch fonctionne aussi — le skip-count du sampler de pretrain_ddp.py reprend à l'index de batch sauvegardé dans l'epoch sauvegardé.

Forme de l'emploi du temps : héritée de saved_args du ckpt. Les remplacements CLI de tout flag de l'emploi du temps (--epochs / --warmup-epochs / --init-lr / --max-lr / --final-lr) sont rejetés avec une erreur hard — pure resume signifie pure resume ; si vous voulez changer l'emploi du temps, supprimez --resume et démarrez une exécution fresh-schedule.

Exigences : le ckpt doit avoir été sauvegardé via save_model_for_restart (c.-à-d., porter les clés optimizer / scheduler_step / epoch / batch_idx). Si l'une d'elles manque, le runner errore avec un message clair et suggère de supprimer --resume.

Le mode par défaut est le bon choix ~90% du temps. Utilisez --resume uniquement quand vous devez vraiment continuer une seule exécution d'entraînement interrompue.

Workflow

Soit $KERMT_REPO le chemin vers votre checkout de repo kermt, et supposez que kermt-setup a déjà construit kermt:latest. Tous les chemins ci-dessous sont sur l'hôte ; l'assistant bind-mount les place aux chemins de conteneur connus.

  1. Pré-vol : assurer conteneur + sonde système.

    "$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')
"

    Surfacez tout gaps à l'utilisateur. Refusez de procéder si ok: false.

  2. Calculer le répertoire d'exécution.

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

  3. Résoudre et valider le checkpoint.

    Résoudre — seulement si --ckpt a été omis. Par défaut au modèle pretrained released hybrid nvidia/NV-KERMT-70M-v2 :

    • Portail de consentement. À moins que --pretrained-release n'ait été passé, demandez à l'utilisateur : « Aucun checkpoint donné — télécharger le modèle released nvidia/NV-KERMT-70M-v2 (NVIDIA Open Model License, https://huggingface.co/nvidia/NV-KERMT-70M-v2) et continuer-pretrain à partir de celui-ci ? [y/N] ». Ne jamais télécharger sans un oui explicite (ou --pretrained-release). Si both --ckpt et --pretrained-release sont donnés, avorter — ils sont en conflit.
    • Lieu de sauvegarde. Par défaut $KERMT_REPO/models/NV-KERMT-70M-v2/ ; honorez --model-dir <dir> si donné. Un bundle déjà complet là-bas est réutilisé.
    • Télécharger (au premier plan ; ~282 MB à la première lecture) :
      "$SKILL_DIR/scripts/kermt_container.sh" run --model-dir <save-dir> -- \
    "python /skill/scripts/fetch_released_model.py --out /model"

      Parsez le JSON ; avorter sur ok: false (surfacer errors). En cas de succès, définissez <user-ckpt> = <save-dir>/kermt_contrastive_v2.0.pt. Les trois fichiers vocab du bundle atterrissent également dans <save-dir>, donc la détection auto de --vocab-dir de l'étape 5 (qui regarde dans le répertoire parent du ckpt) les trouve sans travail supplémentaire.

    Valider le ckpt résolu (ou fourni par l'utilisateur) :

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

    Parsez le JSON. Avorter sur ok: false, affichant l'erreur verbatim. Le message d'erreur redirige l'utilisateur vers kermt-add-cmim-pretrain pour les ckpts encoder-only, ou vers kermt-finetune pour les ckpts finetuned.

  4. Valider les données.

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

    Avorter sur ok: false.

  5. Préparer les données (sauter si --from-prepare est donné). Passer le vocab du ckpt. Regardez dans le répertoire parent du ckpt pour les fichiers conventionnels pretrain_atom_vocab.{json,pkl}, pretrain_bond_vocab.{json,pkl}, et pretrain_smiles_vocab.pkl (la convention de bundling pour les modèles released ; voir references/released-models.md). Si tous les trois sont présents, passez auto via --vocab-dir <ckpt_parent_dir>. Si seulement certains sont présents, passez-les via des flags explicites (--atom-vocab, --bond-vocab, --smiles-vocab). Si aucun n'est présent, demandez au utilisateur pour --vocab-dir — ou refusez de procéder, car la reconstruction d'un fresh vocab depuis le nouveau corpus décallerait silencieusement les vocab heads du ckpt (le vocab du ckpt est l'autorité pour continue-pretrain).

    Notez le pattern de montage à deux couches : passez le répertoire d'hôte à kermt_container.sh --vocab-dir (qui le monte à /vocab à l'intérieur du conteneur), et référencez /vocab depuis la commande prepare_data.py intérieure. Le même pattern s'applique à chaque chemin d'hôte que la commande intérieure doit lire (--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]"

    Les sorties atterrissent à $RUN_DIR/data/prepare_data.json avec vocab_source: "user_provided". L'étape 7 du runner vérifiera que les comptes d'entrées des fichiers vocab correspondent aux tailles des vocab-heads du ckpt et refusera de lancer en cas de décalage.

  6. Estimer le runtime + confirmer avec l'utilisateur.

    • Le temps wall de pretrain dépend de la taille du corpus × epochs × nombre de GPU.
    • Dites à l'utilisateur l'estimation ; demandez « procéder ? » à moins que le flag --yes n'ait été donné (cas agent-non-interactive).
    • Template d'estimation exemple : ~N hours on K GPUs for E epochs over M molecules (~steps/epoch × seconds/step).

  7. Lancer le runner en détaché.

    "$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 ...]"

    Retourne le nom du conteneur + id + chemin du fichier log.

  8. Signaler à l'utilisateur. Sortir un résumé court :

    • Nom du conteneur + id
    • $RUN_DIR/run.json (le manifest avec cmd_replay + image digest)
    • Fichier log : $RUN_DIR/logs/pretrain_ddp.log
    • TensorBoard : $RUN_DIR/logs/tb (ouvrir avec tensorboard --logdir $RUN_DIR/logs/tb)
    • Suggérez d'invoquer kermt-monitor <RUN_DIR> pour vérifier la progression.

Règles strictes

  • Ne jamais télécharger le modèle released sans consentement. Quand --ckpt est omis, téléchargez nvidia/NV-KERMT-70M-v2 seulement après un « oui » explicite de l'utilisateur ou un flag --pretrained-release explicite. --ckpt et --pretrained-release s'excluent mutuellement.
  • Ne jamais modifier le ckpt d'entrée de l'utilisateur. Le runner y crée un symlink dans le save_dir ; le symlink est ce sur lequel pretrain_ddp.py auto-resume. Le fichier source reste intact.
  • Ne jamais remplacer silencieusement l'arch. Si l'utilisateur passe un --hidden-size etc. qui ne correspond pas à la valeur dérivée du ckpt, le runner avorte bruyamment. Les params arch viennent du ckpt, point.
  • Ne jamais bloquer sur le long-running pretrain lui-même. Le runner est invoqué via run_detached ; la skill retourne immédiatement après l'étape 8. Utilisez kermt-monitor pour la progression.
  • Échoez les defaults appliqués à l'utilisateur. Le champ args_applied de run.json enregistre la valeur de chaque flag + source (user / default-config / auto-1gpu / auto-multi-gpu). La skill devrait surfacer un résumé de tout flag non user-spécifié pour que l'utilisateur sache ce qui a été supposé.

Erreurs courantes

  • model_type='finetuned' rejeté → le ckpt est un downstream finetune, pas un pretrain. L'erreur redirige vers le workflow pertinent.
  • grover_base ckpt has no vocab head → ckpt encoder-only (p.ex. le grover_base.pt original-grover). L'erreur redirige vers kermt-add-cmim-pretrain.
  • prepare_data manifest is missing required outputs → l'utilisateur a passé --from-prepare à un répertoire où prepare a été exécuté avec --skip-vocab ou --skip-split. Réexécutez prepare sans ces flags.
  • --gpus all non disponible → installer nvidia-container-toolkit ; vérifier kermt_container.sh check_system.

Reproductibilité

Le champ cmd_replay de run.json est une commande sur une seule ligne qui ré-exécute le pretrain avec les mêmes entrées, hyperparamètres et arch. Pour replayer :

# À l'intérieur du conteneur kermt :
$(jq -r .cmd_replay $RUN_DIR/run.json)

Si ok_to_replay: false dans le manifest (parce que l'arborescence de travail du repo kermt était modifiée au moment du lancement), le replay peut ne pas être bit-exact — épinglez le commit exact via le champ repo.commit et git checkout le d'abord.

Skills similaires

kermt-add-cmim-pretrain
nvidia / skills

Convertir un checkpoint grover_base en modèle hybride cMIM et continuer le pré-entraînement.

3 302
kermt-pretrain-scratch
nvidia / skills

Entraîner un modèle KERMT de zéro sur un corpus chimique personnalisé.

3 302
kermt-finetune
nvidia / skills

Affiner un modèle KERMT pré-entraîné sur un CSV étiqueté fourni par l'utilisateur.

3 302
recipe-recommender
nvidia / skills

Indexer les recettes d'entraînement et recommander la configuration optimale pour LLM.

3 302
mbridge-recipe-recommender
nvidia / skills

Indexer les recettes d'entraînement et recommander la configuration optimale pour LLM.

3 302