kermt-add-cmim-pretrain

Par nvidia · skills

Convertit un checkpoint grover_base (encoder seul ou encoder + têtes de vocabulaire) en checkpoint hybride en ajoutant un décodeur cMIM et un latent_dist initialisés aléatoirement, puis reprend le pré-entraînement sur le corpus de l'utilisateur en mode hybride (vocab + contrast). Équivaut à kermt-continue-pretrain avec une étape ponctuelle de conversion de checkpoint en amont.

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

kermt-add-cmim-pretrain

Convertir un checkpoint grover_base (ancien GROVER d'origine grover.encoders.* ou moderne kermt.encoders.*, avec ou sans vocab heads) en checkpoint hybride complet (cMIM + vocab), puis continuer le pretraining sur le corpus de l'utilisateur en mode hybride.

C'est un simple wrapper : upgrade_to_hybrid.py produit un nouveau ckpt classé comme model_type: hybrid via check_checkpoint.py, et le reste du workflow est identique à kermt-continue-pretrain.

État : expérimental. Ce workflow fonctionne de bout en bout mais n'a pas été comparé au training hybride from-scratch du manuscrit (qui produit le checkpoint publié). À utiliser comme alternative expérimentale à kermt-pretrain-scratch quand on veut étendre un checkpoint grover_base existant plutôt que redémarrer depuis une init aléatoire. Validez la performance downstream sur votre propre benchmark avant de vous fier au ckpt upgradé pour la production.

Chemins skill et runtime

Définissez SKILL_DIR au chemin absolu du répertoire skill installé. Exportez KERMT_REPO comme chemin absolu vers le checkout KERMT utilisé pour l'exécution du modèle. Le helper de conteneur bundlé monte ce checkout à /workspace et ce skill à /skill (lecture seule). Les commandes dans le conteneur utilisent /skill/scripts/ ; les valeurs par défaut sont bundlées dans config/.

Prérequis matériels

Identiques à kermt-continue-pretrain (le décodeur cMIM ajoute des paramètres mais sans augmenter sensiblement les besoins ; la marge VRAM devrait être suffisante). L'étape upgrade elle-même est rapide (~5 s) et CPU uniquement — seul le continue-pretrain subséquent consomme GPU.

Quand invoquer

  • L'utilisateur dispose d'un checkpoint grover_base (encoder uniquement ou avec vocab heads) et veut l'étendre en pretraining hybride (vocab + contrastive cMIM).
  • Utile pour ajouter l'objectif contrastif de reconstruction SMILES à un encodeur préentraîné sans redémarrer le pretraining from scratch (ce que ferait kermt-pretrain-scratch sur une échelle de jours).

Pour continuer un ckpt hybride ou cmim existant : utilisez kermt-continue-pretrain directement. Pour entraîner un modèle neuf sur un corpus personnalisé : utilisez kermt-pretrain-scratch.

Entrées

Requis :

  • --ckpt <path> — ckpt grover_base à upgrader. Validé via check_checkpoint.py --mode upgrade_to_hybrid ; rejeté si le ckpt a déjà une contrast head ou FFN de tâche.
  • --csv <path> — CSV de corpus pretraining. Même structure que l'entrée --csv de kermt-continue-pretrain.

Optionnel (identique à kermt-continue-pretrain) :

  • --val-csv <path> — CSV de validation séparé. Sans cela, prepare_data auto-divise par --val-frac 0.1.
  • Surcharges d'hyperparamètres d'entraînement (--epochs N, --batch-size N, triple lr, --warmup-epochs F, etc.).
  • --vocab-loss-weight F / --latent-dim N / --contrastive-temperature F.
  • --wandb-project NAME / --wandb-run-name NAME — logging optionnel Weights & Biases (nom de run honoré seulement avec un projet). Désactivé par défaut.
  • --gpus 0,2.

Workflow

Soit $KERMT_REPO le chemin vers votre checkout de repo kermt.

  1. Pré-vol : check_system (identique à l'étape 1 de kermt-continue-pretrain).

  2. Calculer le répertoire run :

    RUN_DIR=$KERMT_REPO/runs/add-cmim-pretrain_$(date -u +%Y-%m-%dT%H-%M-%SZ)
  3. Valider le ckpt d'entrée avec check_checkpoint --mode upgrade_to_hybrid. Arrêtez sur ok: false. Le validateur rejette les ckpts qui ont déjà une contrast head (suggère kermt-continue-pretrain) ou des FFN task heads (le ckpt a été finetué ; suggère d'utiliser le checkpoint pretraining original).

  4. Valider le corpus via check_data --mode pretrain. Arrêtez sur ok: false.

  5. Préparer les données avec --mode pretrainsans --vocab-dir. L'upgrade construit des vocab heads frais dimensionnés au vocab du corpus, donc on veut que prepare_data produise un vocab neuf depuis le corpus plutôt que de passer le vieux vocab du ckpt (qui peut même ne pas exister pour les ckpts grover_base legacy encoder-only) :

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

    Le manifest de sortie a vocab_source: "built_fresh" et inclut un smiles_vocab (construit depuis le corpus, nécessaire pour le nouveau décodeur).

  6. Upgrader le 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"

    Présentez le résumé JSON à l'utilisateur — en particulier warnings[], qui inclut des notes sur la dérive d'arch encoder (par ex. le GROVER legacy avait deux clés act_func_* supplémentaires que KERMTEmbedding moderne n'a pas) et la note de restriction argparse --backbone de pretrain_ddp.py si le backbone du ckpt upgradé est autre chose que gtrans.

  7. Estimer le runtime + confirmer avec l'utilisateur. Même heuristique que kermt-continue-pretrain (taille corpus × epochs × nombre GPU → temps mur).

  8. Lancer le runner détaché.

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

    Le runner voit le ckpt upgradé comme model_type: hybrid, donc il auto-dispatche --pretrain_mode hybrid --vocab_loss_weight 1.0 avec smiles_vocab raccordé.

  9. Rapporter à l'utilisateur le chemin du ckpt upgradé + le même pointeur run.json / chemin log / pattern URL tensorboard que kermt-continue-pretrain.

Règles strictes

  • Ne jamais modifier le ckpt d'entrée de l'utilisateur. L'upgrade écrit un nouveau fichier à <run_dir>/upgraded.pt ; le ckpt source reste inchangé.
  • Les vocab heads sont toujours neufs. Même si le grover_base d'entrée a des vocab heads, ils sont jetés et reconstruits dimensionnés au vocab du nouveau corpus. Le continue-pretraining du ckpt upgradé entraînera ces nouvelles heads aux côtés du décodeur.
  • Ne pas relâcher auto-matiquement les choix --backbone. Si l'avertissement upgrade se déclenche parce que le backbone du ckpt d'entrée n'est pas gtrans (par ex. legacy dualtrans), présentez l'avertissement et demandez à l'utilisateur. Ne modifiez PAS silencieusement parsing.py pour ajouter le backbone legacy à la liste des choix.

Erreurs courantes

  • check_checkpoint a rejeté le ckpt avec model_type=hybrid ou cmim → le ckpt de l'utilisateur a déjà une contrast head. Redirigez vers kermt-continue-pretrain.
  • check_checkpoint a rejeté le ckpt avec task_ffn=true → le ckpt a été finetué. Le workflow upgrade supporte seulement les checkpoints pretraining.
  • prepare manifest missing smiles_vocab → prepare_data a été invoqué avec --skip-vocab ou équivalent qui a omis le smiles vocab. Relancez prepare sans ces flags.
  • unexpected key(s) in encoder load warning → les architectures GROVER legacy ont sauvegardé quelques poids act_func_* que KERMTEmbedding moderne n'utilise pas. Bénin ; le reste de l'encodeur a chargé correctement.

Ce qui est dans run.json après une exécution réussie

Mêmes champs de reproducibilité que kermt-continue-pretrain, plus le summary.json de l'étape upgrade est capturé sous le chemin inputs.upgrade_summary pour que la provenance du ckpt upgradé soit auditable.

Rejeu

Identique à kermt-continue-pretrain : cmd_replay reconstruit l'invocation run_pretrain_local.py --ckpt <upgraded.pt> .... Pour refaire le flux add-cmim complet end-to-end, l'utilisateur a aussi besoin du ckpt grover_base d'entrée et du corpus — tous deux sont capturés dans les manifests prepare_data et upgrade par chemin absolu.

Skills similaires