kermt-finetune

Affinez un encodeur KERMT pré-entraîné sur un CSV étiqueté fourni par l'utilisateur. La skill est l'orchestrateur de workflow : valider le checkpoint, valider les données, préparer les données, lancer le runner en détaché, retourner un répertoire de run + nom de conteneur.

Chemins de la skill et du runtime

Définissez SKILL_DIR comme 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'helper de 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/. Consultez Released models pour les exigences de 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 script scripts/fetch_released_model.py fourni télécharge le bundle de modèle via HTTPS dans le répertoire hôte sélectionné par l'utilisateur. Les modèles publics fonctionnent sans identifiants ; si HF_TOKEN est défini, l'helper de 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 de run choisi.

Exigences matérielles

• GPUs : 1 par défaut (single-GPU) ; passez --gpus 0 (ou n'importe quel id) pour en sélectionner un. Pour un entraînement plus rapide sur un hôte multi-GPU, passez --num-gpus N (N>1) pour exécuter DDP data-parallel sur N GPUs — --batch-size est alors par GPU (batch global effectif = batch_size × N).
• VRAM : ≥ 8 GB pour la configuration batch_size 32 par défaut. Une VRAM plus faible fonctionne avec des batch sizes plus petits — passez --batch-size N pour remplacer.
• Disque : quelques GB par run (checkpoint + features + logs).
• Driver / CUDA : tout hôte prenant en charge CUDA 12.6 (la base de l'image kermt). kermt-setup valide ceci à l'avance.

Entrées

Obligatoires :

• --csv <path> — CSV étiqueté. La première colonne est smiles ; chaque autre colonne est une cible.

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

• --ckpt <path> — checkpoint pré-entraîné d'entrée (grover_base / cmim / hybrid). Le validateur refuse les ckpts déjà affinés avec une redirection vers kermt-infer. S'il est omis, la skill propose de télécharger le modèle hybrid pré-entraîné released nvidia/NV-KERMT-70M-v2 et d'affiner à partir de celui-ci — voir « Resolve & validate the checkpoint » (étape du workflow 3).
• --pretrained-release — opt-in explicite pour utiliser le modèle released sans l'invite interactive (pour les runs non-interactifs / 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 complet déjà là est réutilisé, non re-téléchargé.

Optionnels :

• --dataset-type {regression | classification | multiclass} — par défaut regression (de defaults_finetune.json). Conduit la loss, les valeurs par défaut de métrique et l'initialisation de head. Pour les tâches de classification, passez --dataset-type classification.

• --targets COL [COL ...] — noms de colonnes cibles explicites. S'il est omis, le validateur détecte automatiquement les colonnes numériques non-smiles et la skill confirme avec l'utilisateur avant de continuer.

• --val-csv <path> et --test-csv <path> — splits val + test fournis par l'utilisateur. Passez les deux ou ne passez rien (la skill auto-split en utilisant le --split-type configuré).

• --split-type {random | scaffold_balanced | index_predetermined} — par défaut scaffold_balanced de defaults_finetune.json.

  • random et scaffold_balanced : construisez le split val/test en interne à partir du train CSV. Pas de --val-csv / --test-csv nécessaire.
  • index_predetermined : requiert les CSVs pré-split passés via --val-csv + --test-csv (et, séparément, les fichiers d'index par fold — voir kermt/util/utils.split_data). Utilisez ceci quand le dataset livre son propre split canonique (par exemple tests/data/Biogen_for_grover/scaffold/balance/<endpoint>/{train,val,test}.csv).

• --metric NAME — mae (régression par défaut), auc (classification par défaut), ou n'importe quel nom que kermt.util.metrics.get_metric_func accepte.

• --epochs N / --batch-size N / --init-lr F / --max-lr F / --final-lr F / --warmup-epochs F / --weight-decay F / --dropout F / --bond-drop-rate F / --dist-coff F / --early-stop-epoch N / --seed N — remplacements d'hyper-paramètre d'entraînement. Tout ce qui n'est pas donné est rempli de config/defaults_finetune.json.

• --ffn-hidden-size N / --ffn-num-layers N — dimensions du trunk FFN partagé.

• --ffn-num-task-specific-layers N / --ffn-task-specific-hidden-size H — heads FFN par cible (par défaut 0 = désactivé ; utile pour les affinages multi-cibles hétérogènes). Les deux doivent être définis ensemble quand N > 0.

• --ensemble-size N / --num-folds N — multi-modèles / k-fold CV. Par défaut 1 chacun.

• --gpus 0 — id GPU unique pour l'affinage single-process (par défaut 0). Ignoré quand --num-gpus > 1.

• --num-gpus N — nombre de GPUs pour l'affinage data-parallel DDP. Par défaut 1 (single-process, inchangé). N>1 exécute main.py finetune avec WORLD_SIZE=N (un process par GPU) ; --batch-size est par GPU.

• --from-prepare <dir> — ignorez l'étape prepare et réutilisez un prepare_data.json existant dans <dir>. Utile lors de l'itération sur les hyper-paramètres.

Workflow

Soit $KERMT_REPO le chemin vers votre checkout du repo kermt, et supposons que kermt-setup ait construit kermt:latest. Tous les chemins ci-dessous sont sur l'hôte ; l'helper les bind-monte à des chemins de conteneur connus.

1. Pré-vol : assurer le 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); CUDA via container toolkit')
   "

   Refusez de continuer si ok: false.

2. Calculez le répertoire de run.

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

3. Resolve & validez le checkpoint.

   Resolve — uniquement si --ckpt a été omis. Par défaut le modèle hybrid pré-entraîné released nvidia/NV-KERMT-70M-v2 :

   • Portail de consentement. Sauf si --pretrained-release a été passé, demandez à l'utilisateur : « Aucun checkpoint fourni — 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 affiner à partir de celui-ci ? [y/N] ». Ne téléchargez jamais sans un oui explicite (ou --pretrained-release). Si les deux --ckpt et --pretrained-release sont donnés, abandon — ils entrent en conflit.
   • Emplacement de sauvegarde. Par défaut $KERMT_REPO/models/NV-KERMT-70M-v2/ ; honorez --model-dir <dir> s'il est donné. Un bundle complet existant est réutilisé.
   • Télécharger (premier plan ; ~282 MB au premier téléchargement) :

     "$SKILL_DIR/scripts/kermt_container.sh" run --model-dir <save-dir> -- \
         "python /skill/scripts/fetch_released_model.py --out /model"

     Analysez le JSON ; abandonnez sur ok: false (afficher errors). Au succès, définissez <user-ckpt> = <save-dir>/kermt_contrastive_v2.0.pt.

   Validez le ckpt résolut (ou fourni par l'utilisateur) :

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

   Analysez le JSON. Abandonnez sur ok: false. Le validateur rejette les ckpts déjà affinés (has_task_ffn: true) avec une redirection vers kermt-infer.

4. Validez les données.

   "$SKILL_DIR/scripts/kermt_container.sh" run --data <user-csv> -- \
       "python /skill/scripts/check_data.py --mode finetune --csv /data/<basename> [--targets COL1 COL2 ...]"

   Si --targets n'a pas été donné par l'utilisateur, afficher auto_detected_targets du JSON et demandez à l'utilisateur de confirmer avant de continuer. Abandonnez sur ok: false.

5. Préparez les données (ignorer si --from-prepare donné).

   Pré-vol : vérifier les fichiers val.csv / test.csv voisins. Avant d'invoquer preparedata, inspectez le répertoire parent de <user-csv>. Si un val.csv voisin canonique (ou `val.csv— les variantes communes incluentval_T.csv,valclean.csv) ET untest.csv/test.csvcorrespondant existent à côté du train CSV, le dataset livre son propre split pré-défini. **Dans ce cas, définissez--split-type index_predeterminedET passez--val-csv/--test-csv** — sinon lesplit_typeconfiguré (par défautscaffold_balanced) re-splittera le train CSV à partir de zéro et supprimera silencieusement les fichiers val/test de l'utilisateur. En cas de doute — ou quand les fichiers voisins utilisent des suffixes non-canoniques (_T,_v2`, etc.) — afficher la situation à l'utilisateur et demandez lequel il veut.

   Quotage des noms de cibles. Si l'un des noms de colonnes --targets contient des métacaractères shell (>, &, |, (, ), $, etc.), mettez en single-quote chacun lors du passage sur la CLI pour empêcher le shell de manger une partie du nom. Exemple : --targets 'Log_Caco2_Papp_A>B' 'logD'. L'en-tête CSV lui-même est lu directement par le trainer en aval et n'est pas affecté, mais le champ targets[] du manifest prepare_data.json capture ce que le shell livre — les métacaractères non-quotés sont tronqués là.

   Note de montage : kermt_container.sh --data <host-csv> monte le répertoire parent de <host-csv> à /data. --val-csv et --test-csv doivent donc référencer des fichiers dans ce même répertoire parent. Si val/test vivent dans un répertoire séparé (par exemple un dossier splits/ voisin), montez le parent de tous les trois en utilisant --data <dir> sur un répertoire plutôt qu'un fichier.

   "$SKILL_DIR/scripts/kermt_container.sh" run --data <user-csv> --run-dir $RUN_DIR -- \
       "python /skill/scripts/prepare_data.py --mode finetune \\
            --csv /data/<basename> --out /runs/data \\
            --split-type <split_type> \\
            [--val-csv /data/<val-basename> --test-csv /data/<test-basename>] \\
            [--val-frac 0.1 --test-frac 0.1 --seed 0] \\
            --targets <COL1> [COL2 ...]"

   Les sorties se trouvent à $RUN_DIR/data/prepare_data.json. Pour scaffold_balanced et index_predetermined, la préparation émet un seul clean_full_csv + .npz ; le runner les passe à main.py finetune qui appelle split_data en interne avec la seed fournie par l'utilisateur.

6. Estimez le runtime + affichez les défauts appliqués.

   • Le temps mur de l'affinage est généralement de minutes à heures sur 1 GPU.
   • Affichez un résumé de chaque flag qui a été rempli à partir des défauts vs fourni par l'utilisateur, afin que l'utilisateur sache ce qui a été supposé. Le runner enregistre ceci dans args_applied.
   • Message exemple : « Remplissage à partir de defaults_finetune.json : epochs=30, batch_size=32, split_type=scaffold_balanced. Remplacez n'importe lequel de ceux-ci avec --<flag>. »

7. Portail de confirmation des cibles (exigence stricte). Avant de lancer le runner, quel que soit le mode de détermination de la liste de cibles (CLI --targets, auto-détection à l'étape 4, ou une demande en langage naturel de l'utilisateur comme « affinez sur Caco2 et HLM »), affichez la liste des cibles finales à l'utilisateur avec un décompte explicite : « Affinez sur N cible(s) : COL1, COL2, ... ». Si la demande de l'utilisateur a spécifié un sous-ensemble qui ne correspond pas à cette liste (par exemple, ils ont demandé 2 tâches en langage naturel mais la liste en a toujours 4), traitez-le comme une divergence et re-demandez avec la différence — n'avancez jamais silencieusement sur le mauvais ensemble de cibles. Attendez une confirmation explicite avant de lancer à moins que --yes n'ait été donné.

8. Lancez le runner en détaché. (Cohérent avec les skills de pré-entraînement.)

   "$SKILL_DIR/scripts/kermt_container.sh" run_detached \\
       --name kermt-finetune-<ts> \\
       --ckpt <user-ckpt> --run-dir $RUN_DIR -- \\
       "python /skill/scripts/run_finetune_local.py \\
            --ckpt /ckpt \\
            --prepare-manifest /runs/data/prepare_data.json \\
            --dataset-type <type> \\
            --out /runs \\
            [--gpus 0] \\
            [--num-gpus N] \\
            [--epochs N --batch-size N --init-lr F ...] \\
            [--ffn-num-task-specific-layers N --ffn-task-specific-hidden-size H]"

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

9. Rapportez à l'utilisateur. Sortez un court résumé :

   • Nom du conteneur + id
   • $RUN_DIR/run.json (manifest avec cmd_replay + image digest)
   • Fichier log : $RUN_DIR/logs/finetune.log
   • TensorBoard : $RUN_DIR/logs/tb (ouvrir avec tensorboard --logdir $RUN_DIR/logs/tb)
   • Les checkpoints finaux se trouvent à $RUN_DIR/ckpt/fold_0/model_0/model.pt (meilleure-val) et last_checkpoint.pt (voisin, cible auto-resume). Les prédictions et métriques de test retenues se trouvent à $RUN_DIR/ckpt/fold_0/test_result.csv. Les chemins varient avec --num-folds / --ensemble-size.
   • Pour suivre la progression : kermt-monitor <RUN_DIR> (one-shot) ou docker logs -f <container-name> (streaming).
   • Pour bloquer jusqu'à la fin du run (utile pour les courts runs de test) : docker wait <container-name> — affiche le code de sortie à la fin.

Règles strictes

• Ne téléchargez jamais le modèle released sans consentement. Quand --ckpt est omis, téléchargez nvidia/NV-KERMT-70M-v2 uniquement après un « oui » explicite de l'utilisateur ou un flag --pretrained-release explicite. --ckpt et --pretrained-release s'excluent mutuellement.
• Ne modifiez jamais le ckpt d'entrée de l'utilisateur. Le runner passe son chemin via --checkpoint_path ; task/train.py le charge en lecture seule dans le modèle et attache un nouveau head FFN. Le fichier source reste intact.
• L'arch provient du ckpt, pas de CLI/défauts. Le runner extrait hidden_size, depth, num_attn_head, activation, embedding_output_type, self_attention (+ attn_hidden / attn_out quand applicable) de saved_args du ckpt. Il n'y a pas de flag --hidden-size sur ce runner.
• Ne bloquez jamais sur l'affinage long. La skill lance via run_detached et revient immédiatement après l'étape 9. Utilisez kermt-monitor.
• Affichez les défauts appliqués à l'utilisateur. Le champ args_applied de run.json enregistre la valeur de chaque flag + source (utilisateur / default-config). Affichez un résumé d'une ligne de chaque flag rempli à partir des défauts afin que l'utilisateur sache ce qui a été supposé.

Erreurs courantes

• finetune_init requires a pretrain ckpt (grover_base / cmim / hybrid) → le ckpt que vous avez passé est déjà affiné (a des heads task FFN). Choisissez un ckpt pré-entraîné à la place, ou utilisez kermt-infer si vous voulez exécuter des prédictions avec le modèle affiné existant. Pour reprendre un affinage sur le MÊME dataset, contournez la skill et appelez python main.py finetune --checkpoint_path <ckpt> ... directement — la skill agent ne supporte pas le resume car l'identité de task sauvegardée ne peut pas être vérifiée automatiquement par rapport aux nouvelles données d'entraînement.
• prepare_data manifest reports ok=False → vérifiez errors pour l'étape échouée (généralement clean_smiles ou save_features). Corrigez et re-lancez.
• ffn_num_task_specific_layers=N>0 but ffn_task_specific_hidden_size is unset → les heads MTL ont besoin d'une taille cachée explicite. Passez --ffn-task-specific-hidden-size H.
• finetune is single-GPU (de --gpus 0,1) → --gpus sélectionne un device pour l'affinage single-process. Pour multi-GPU, utilisez --num-gpus N (DDP) à la place.

Rejoabilité

Le champ cmd_replay de run.json est une commande single-line qui re-lance l'affinage avec les mêmes entrées, hyper-paramètres et arch. Pour rejouer à 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'arbre de travail du repo kermt était sale au moment du lancement), la rejoabilité peut ne pas être bit-exacte — épinglez le commit exact via le champ repo.commit et git checkout le d'abord.