kermt-pretrain-scratch

Par nvidia · skills

Préentraîne un nouveau modèle KERMT from scratch à partir d'un corpus fourni par l'utilisateur. Construit un nouveau vocabulaire à partir du corpus, instancie l'architecture du modèle depuis les valeurs par défaut, et lance `pretrain_ddp.py` dans le conteneur kermt (en mode détaché pour les longues exécutions). Contrairement à `kermt-continue-pretrain`, aucun checkpoint de départ n'est chargé — le modèle est initialisé aléatoirement.

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

kermt-pretrain-scratch

Préentraîner un nouveau modèle KERMT à partir de zéro sur un corpus fourni par l'utilisateur. Utile quand on veut réentraîner un modèle sur un domaine chimique personnalisé plutôt que d'étendre l'un des checkpoints publics. Considérablement plus cher que kermt-continue-pretrain — pas de démarrage à chaud, donc les courbes de loss doivent descendre de zéro sur de nombreuses epochs.

Chemins des skills et runtime

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

Exigences matérielles

Identiques à kermt-continue-pretrain :

  • GPUs : 1–N avec support CUDA. Le runner auto-détecte via torch.cuda.device_count() ; --gpus 0,2 permet de changer. Fallback monoGPU : --batch_size 32 --save_interval 500. Multi-GPU conserve les valeurs par défaut (--batch_size 256 etc.). Note : --gpus N utilise l'indexage torch.cuda, qui peut différer de l'ordre d'affichage de nvidia-smi sur les hosts multi-GPU (bus PCI vs énumération CUDA). Pour cibler un GPU physique spécifique, définissez CUDA_VISIBLE_DEVICES avant l'invocation, ou lancez python -c "import torch; print([torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())])" pour confirmer le périphérique sélectionné.

  • VRAM : la --batch-size 256 par défaut est calibrée pour du matériel A100 (80 GB VRAM). Sur des GPUs plus petites, réduisez 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 (par défaut)

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

  • Disque : des dizaines de GB pour les shards + vocab + checkpoints, à l'échelle des epochs.

  • Wall time : c'est la grande différence. Préentraîner de zéro sur un corpus de 11M molécules à 100 epochs prend généralement des jours même sur une box multi-GPU. Le skill affiche une estimation avant le lancement ; confirmez avec l'utilisateur.

Quand invoquer

  • L'utilisateur veut entraîner un nouveau modèle sur un corpus personnalisé (p. ex. chimie spécifique au domaine que les ckpts publics ne couvrent pas).
  • L'utilisateur veut reproduire une config de préentraînement de bout en bout sans dépendre d'un ckpt public.

Pour continuer un ckpt public existant, utilisez kermt-continue-pretrain. Pour ajouter un décodeur cMIM à un ckpt grover_base encoder-only, utilisez kermt-add-cmim-pretrain.

Entrées

Obligatoires :

  • --csv <path> — le CSV du corpus de préentraînement avec une colonne smiles. Un fichier par convention ; les corpus multi-fichiers sont reportés. Utilisez --val-csv pour un ensemble de validation séparé.
  • --pretrain-target-mode {vocab|cmim|hybrid} — quel objectif de préentraînement utiliser. Pas de défaut — doit être défini explicitement pour que l'utilisateur fasse un choix éclairé :
    • vocab — prédiction vocab atome + lien style GROVER original (sortie encoder-only, léger).
    • cmim — objectif contrastif + reconstruction SMILES. Nécessite de construire un vocab SMILES à partir du corpus.
    • hybrid — objectifs vocab et contrastif conjointement (la config state-of-the-art du manuscrit KERMT).

Optionnels :

  • --val-csv <path> — CSV de validation séparé. Sans cela, prepare_data auto-divise l'entrée par --val-frac 0.1 (mélange aléatoire avec --seed).
  • Overrides d'hyperparamètres d'entraînement : --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. 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).
  • --wandb-project NAME / --wandb-run-name NAME — logging Weights & Biases optionnel. Quand --wandb-project est défini, rank 0 log train/val losses ; le nom de run n'est honoré que aux côtés d'un project. Désactivé par défaut.
  • --gpus 0,2 — restreindre à un sous-ensemble de GPUs.

Workflow

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

  1. Pré-vol : vérifier conteneur + sonde système (identique à l'étape 1 de kermt-continue-pretrain). Refusez de procéder si check_system signale des lacunes.

  2. Calculer le répertoire de run.

    RUN_DIR=$KERMT_REPO/runs/pretrain-scratch_$(date -u +%Y-%m-%dT%H-%M-%SZ)
  3. Valider le corpus (pas de ckpt à valider, donc c'est la seule vérification d'entrée) :

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

    Abandon sur ok: false.

  4. Préparer les données — pas de pass-through de vocab (on veut un vocab frais depuis le corpus) :

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

    Les sorties arrivent à $RUN_DIR/data/prepare_data.json avec vocab_source: "built_fresh".

  5. Estimer le runtime + avertir fortement. C'est critique pour pretrain-from-scratch :

    • « Préentraîner de zéro est à l'échelle des jours même sur multi-GPU ; les checkpoints KERMT publics ont chacun été entraînés sur des millions de molécules pendant des centaines d'heures-GPU. Si vous voulez surtout tirer parti des connaissances existantes pour une tâche aval, considérez plutôt kermt-continue-pretrain à partir d'un ckpt public, qui converge en heures au lieu de jours. »
    • Montrez taille du corpus × epochs × nombre de GPUs → estimated wall time.
    • Demandez une confirmation explicite sauf si --yes a été donné.
  6. Lancer le runner en détaché.

    "$SKILL_DIR/scripts/kermt_container.sh" run_detached \\
        --name kermt-pretrain-scratch-<ts> \\
        --run-dir $RUN_DIR -- \\
        "python /skill/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 : PAS de flag --ckpt (le runner refuse si les deux --from-scratch et --ckpt sont donnés). Le runner utilise le groupe arch depuis config/defaults_pretrain.json pour dimensionner le modèle.

  7. Rendre compte à l'utilisateur. Incluez toujours tous les éléments suivants — n'omettez pas la ligne TensorBoard quelle que soit la pression sur la longueur de sortie :

    • Nom du conteneur + id
    • $RUN_DIR/run.json (le manifest avec workflow: pretrain-scratch, from_scratch: true, vocab_check: null, arch depuis les defaults, full cmd_replay)
    • Fichier log : $RUN_DIR/logs/pretrain_ddp.log
    • TensorBoard : $RUN_DIR/logs/tb (ouvrir avec tensorboard --logdir $RUN_DIR/logs/tb)
    • Suggérez kermt-monitor <RUN_DIR> pour suivre la progression.

Règles strictes

  • Ne jamais accepter de flag --ckpt. From-scratch est exclusif avec ckpt d'entrée — le runner l'impose ; le skill devrait aussi.
  • Ne jamais mettre --pretrain-target-mode par défaut silencieusement. C'est un choix architectural significatif (vocab = léger, hybrid = SOTA). Invitez l'utilisateur si absent de la CLI.
  • Avertissement fort avant lancement. Le préentraînement from-scratch est le workflow le plus cher. L'utilisateur doit savoir ce à quoi il s'engage.

Erreurs courantes

  • --pretrain-target-mode is required when --from-scratch is set → l'utilisateur a oublié le flag mode. Invitez.
  • --from-scratch is incompatible with --ckpt → l'utilisateur en a fourni les deux ; demandez lequel il entendait.
  • defaults_pretrain.json has no arch group → problème d'état du repo (ne devrait jamais arriver sur un clone frais) ; pointez l'utilisateur vers relancer kermt-setup.

Qu'y a-t-il dans le manifest après un run from-scratch

Mêmes champs de reproductibilité que 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 (non vérifié — vocab construit depuis le corpus est autoritaire pour from-scratch)
  • arch : les valeurs tirées du groupe arch de config/defaults_pretrain.json (avec tous les futurs overrides CLI appliqués).

Répétabilité

Identique à continue-pretrain : cmd_replay est une commande copiable-collable. Si ok_to_replay: false, l'arborescence de travail du repo kermt était sale au moment du lancement — vérifiez repo.commit et git checkout d'abord.

Skills similaires