mcpField
K
kermt-pretrain-scratch

kermt-pretrain-scratch

github.com/nvidia-bionemoデータ・データベース502209 回インストール2 日前に更新
リポジトリ

kermt-pretrain-scratch について

kermt-pretrain-scratch は、独自の化学コーパスで新しい KERMT モデルを学習するためのワークフロースキルです。データを検証し、コーパスから語彙を作成して、vocab、cmim、hybrid の目的関数で学習を開始します。

主な機能

  1. 01KERMT モデルのスクラッチ事前学習
  2. 02SMILES コーパス CSV の検証
  3. 03コーパスからの語彙作成
  4. 04vocab、cmim、hybrid 目的関数の選択
  5. 05学習パラメーターと GPU の設定

活用シーン

  1. 01独自の化学領域向けモデルの学習
  2. 02既存チェックポイントに依存しない事前学習設定の再現

インストール

コマンドをコピーしてターミナルで実行します。

$npx skills add https://github.com/nvidia-bionemo/bionemo-agent-toolkit --skill kermt-pretrain-scratch

インストール後、Claude Code で/kermt-pretrain-scratchを入力して読み込まれたか確認します。

同梱ファイルevals/evals.jsonskill-card.md

詳細

kermt-pretrain-scratch

ユーザーが指定したコーパスを使って、まったく新しい KERMT モデルをスクラッチから事前学習します。リリース済みチェックポイントを拡張するのではなく、独自の化学分野に合わせてモデルを再学習したい場合に便利です。kermt-continue-pretrainより大幅にコストがかかります。ウォームスタートがないため、損失曲線がスクラッチから多くのエポックをかけて下降する必要があります。

ハードウェア要件

kermt-continue-pretrainと同じです。

  • GPU:CUDA 対応の GPU が 1 台以上必要です。実行ツールはtorch.cuda.device_count()で自動検出します。--gpus 0,2で上書きできます。単一 GPU のフォールバック設定は--batch_size 32 --save_interval 500です。複数 GPU の場合は既定値(--batch_size 256など)を使用します。注:--gpus Nはtorch.cudaのインデックスを使用します。複数 GPU ホストでは、nvidia-smiの表示順(PCI バス順)と異なる場合があります。特定の物理 GPU を指定するには、実行前にCUDA_VISIBLE_DEVICESを設定するか、python -c "import torch; print([torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())])"を実行して、選択するデバイスを確認してください。

  • VRAM:既定の--batch-size 256は A100 クラスのハードウェア(VRAM 80 GB)向けです。小型の GPU では、OOM を避けるために値を下げてください。

    GPU クラスVRAM推奨--batch-size
    L4, T4, V100 16 GB16–24 GB32–64
    A100 40 GB, L40, A4040–48 GB128
    A100 80 GB, H100, H20080 GB256(既定値)

    これらはおおよその開始値です。--batch-size Nを指定して上書きしてください。

  • ディスク:シャード、語彙、チェックポイント用に、エポック数に応じて数十 GB 必要です。

  • 所要時間:ここが大きく異なる点です。1,100 万分子のコーパスを使ったスクラッチからの事前学習は、通常 100 エポックで複数 GPU 搭載のマシンでも数日かかります。このスキルは起動前に所要時間の見積もりを表示します。ユーザーに確認してください。

呼び出すタイミング

  • ユーザーが独自のコーパスで新しいモデルを学習したい場合(例:リリース済みチェックポイントではカバーされていない、化学の特定分野)。
  • リリース済みチェックポイントに依存せず、事前学習の設定を最初から最後まで再現したい場合。

既存のリリース済みチェックポイントを継続して学習する場合は、kermt-continue-pretrainを使用してください。エンコーダーのみの grover_base チェックポイントに cMIM デコーダーを追加する場合は、kermt-add-cmim-pretrainを使用してください。

入力

必須:

  • --csv <path> — smiles列を含む事前学習用コーパスの CSV。慣例上、ファイルは 1 つです。複数ファイルのコーパスは未対応です。別の検証セットを使う場合は--val-csvを指定してください。
  • --pretrain-target-mode {vocab|cmim|hybrid} — 使用する事前学習の目的関数。既定値はありません。ユーザーが情報に基づいて選択できるよう、明示的に設定する必要があります。
    • vocab — GROVER 方式の元の原子+結合の語彙予測(エンコーダーのみの出力、軽量)。
    • cmim — 対照学習+SMILES 再構成の目的関数。コーパスから SMILES 語彙を構築する必要があります。
    • hybrid — 語彙と対照学習の両方の目的関数を組み合わせて使用します(KERMT 論文の最先端の設定)。

任意:

  • --val-csv <path> — 別の検証用 CSV。指定しない場合、prepare_data は--val-frac 0.1を使って入力を自動分割します(--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。 指定されていない値はagent/config/defaults_pretrain.jsonから設定されます。
  • --vocab-loss-weight F(hybrid のみ)/ --latent-dim N / --contrastive-temperature F(cmim と hybrid のみ)。
  • --wandb-project NAME / --wandb-run-name NAME — Weights & Biases による任意のログ記録。--wandb-projectを設定すると、rank 0 が学習損失と検証損失を記録します。実行名はプロジェクトと併せて指定した場合のみ有効です。既定では無効です。
  • --gpus 0,2 — 使用する GPU を指定します。

ワークフロー

$KERMT_REPO を、kermt リポジトリのチェックアウト先のパスとします。

  1. 事前確認:コンテナとシステムのプローブを確認する(kermt-continue-pretrain のステップ 1 と同じです)。check_system で不足が報告された場合は、処理を続行しないでください。

  2. 実行ディレクトリを計算する。

    RUN_DIR=$KERMT_REPO/runs/pretrain-scratch_$(date -u +%Y-%m-%dT%H-%M-%SZ)
    
  3. コーパスを検証する(検証対象の ckpt がないため、これが唯一の入力チェックです)。

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

    ok: false の場合は中止してください。

  4. データを準備する — 語彙を引き継ぎません(コーパスから新しい語彙を作成します)。

    $KERMT_REPO/agent/scripts/kermt_container.sh run --data <user-csv> --run-dir $RUN_DIR -- \
        "python agent/scripts/prepare_data.py --mode pretrain \\
             --csv /data/<basename> --out /runs/data \\
             [--val-csv /data/<val-basename>] [--val-frac 0.1] [--seed 0]"
    

    出力は $RUN_DIR/data/prepare_data.json に保存され、vocab_source: "built_fresh" となります。

  5. 実行時間を見積もり、強く警告する。 スクラッチからの事前学習では、これは非常に重要です。

    • 「スクラッチからの事前学習には、マルチ GPU 環境でも数日かかります。リリース済みの KERMT チェックポイントはそれぞれ、数百万個の分子を使い、数百 GPU 時間をかけて学習されました。主に下流タスクで既存の知識を活用したい場合は、代わりにリリース済みの ckpt を使う kermt-continue-pretrain を検討してください。こちらは数日ではなく数時間で収束します。」
    • コーパスのサイズ × エポック数 × GPU 数を示し、推定実行時間を提示します。
    • --yes が指定されていない場合は、明示的な確認を求めます。
  6. ランナーをデタッチして起動する。

    $KERMT_REPO/agent/scripts/kermt_container.sh run_detached \\
        --name kermt-pretrain-scratch-<ts> \\
        --run-dir $RUN_DIR -- \\
        "python agent/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 ...]"
    

    注:--ckpt フラグは指定しません(--from-scratch と --ckpt の両方が指定されると、ランナーは実行を拒否します)。ランナーは agent/config/defaults_pretrain.json の arch グループを使ってモデルのサイズを決定します。

  7. ユーザーに報告する。 以下の項目を必ずすべて含めてください — 出力長に制限がある場合も、TensorBoard の行を省略しないでください。

    • コンテナ名と ID
    • $RUN_DIR/run.json(workflow: pretrain-scratch、from_scratch: true、vocab_check: null、デフォルトの arch、完全な cmd_replay を含むマニフェスト)
    • ログファイル:$RUN_DIR/logs/pretrain_ddp.log
    • TensorBoard:$RUN_DIR/logs/tb(tensorboard --logdir $RUN_DIR/logs/tb で開きます)
    • 進捗確認には kermt-monitor <RUN_DIR> を提案します。

厳守事項

  • --ckpt フラグを決して受け入れないでください。 スクラッチからの実行では入力 ckpt を併用できません。ランナーはこれを強制しますが、このスキルでも同様にしてください。
  • --pretrain-target-mode を黙ってデフォルト設定しないでください。 これは重要なアーキテクチャ上の選択です(vocab は軽量、hybrid は SOTA)。CLI で指定されていない場合は、ユーザーに確認してください。
  • 起動前に強く警告してください。 スクラッチからの事前学習は、最もコストの高いワークフローです。ユーザーが何に同意するのかを把握できるようにする必要があります。

よくあるエラー

  • --pretrain-target-mode is required when --from-scratch is set → ユーザーがモードフラグを忘れています。確認してください。
  • --from-scratch is incompatible with --ckpt → ユーザーが両方を指定しています。どちらを意図していたか尋ねてください。
  • defaults_pretrain.json has no arch group → リポジトリの状態に問題があります(新規クローンで起きることはありません)。kermt-setup を再度実行するようユーザーに案内してください。

スクラッチからの実行後にマニフェストに記録される内容

continue-pretrain と同じ再現性フィールド(repo.commit、kermt_image、cmd_replay、args_applied)に加えて、以下が記録されます。

  • workflow:"pretrain-scratch"
  • from_scratch:true
  • inputs.ckpt:null
  • ckpt_symlink:null
  • vocab_check:null(検証されていません — コーパスから構築した語彙がスクラッチからの実行で正式なものとして扱われます)
  • arch:agent/config/defaults_pretrain.json の arch グループから取得した値(今後追加される CLI オーバーライドがあれば、それも適用されます)。

再実行可能性

continue-pretrain と同じく、cmd_replay はコピーして貼り付けられるコマンドです。ok_to_replay: false の場合、起動時に kermt リポジトリの作業ツリーに未コミットの変更がありました。まず repo.commit を確認して git checkout してください。

コメント