kermt-continue-pretrain

kermt-continue-pretrain

ユーザー提供の KERMT チェックポイント（grover_base / cmim / hybrid）から継続事前学習を行います。このスキルはワークフローオーケストレーターとして、入力を検証し、 コーパスを準備してランナーを起動し、実行ディレクトリを返します。

ハードウェア要件

• GPU：CUDA 対応の NVIDIA GPU が 1～N 台必要です。ランナーは torch.cuda.device_count() で自動検出します。--gpus 0,2 で上書きできます。単一 GPU ではランナーは --batch_size 32 --save_interval 500 に切り替えます。複数 GPU の場合は defaults_pretrain.json の値（現在は 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 GB	16–24 GB	32–64
  A100 40 GB, L40, A40	40–48 GB	128
  A100 80 GB, H100, H200	80 GB	256（デフォルト）

  これらは大まかな目安です。--batch-size N を渡すと上書きできます。

• ディスク：コーパスのサイズとエポック数に応じて、数十 GB が必要です（チェックポイント 1 つあたり数百 MB です）。

• ドライバー／CUDA：CUDA 12.6（kermt イメージのベース）に対応するホストであれば使用できます。 kermt-setup が事前に検証します。

入力

必須：

• --csv <path> — 事前学習用 CSV（単一列 smiles）。train/val CSV が別々にある場合は、--val-csv <path> も渡してください。

チェックポイント（省略可。省略時はリリース済みモデルがデフォルトで使用されます）：

• --ckpt <path> — 続きの学習に使用する入力事前学習チェックポイント。vocab heads を含む grover_base、cmim、または hybrid の ckpt である必要があります。バリデーターはそれ以外をすべて拒否し、正しいワークフローに誘導します。省略した場合、スキルはリリース済みの事前学習済み hybrid モデル nvidia/NV-KERMT-70M-v2 のダウンロードを提案し、そのモデルから事前学習を続行します。「Resolve & validate the checkpoint」（ワークフローのステップ 3）を参照してください。リリース済みのバンドルには、ckpt とともに 3 つの vocab ファイルが含まれているため、authoritative-vocab の引き渡し（ステップ 5）は自動的に機能します。
• --pretrained-release — 対話形式のプロンプトを表示せず、リリース済みモデルを使用する明示的なオプトイン（非対話型／エージェント実行用）。--ckpt とは相互排他的です。
• --model-dir <dir> — ダウンロードしたバンドルの保存先（デフォルトは $KERMT_REPO/models/NV-KERMT-70M-v2/）。この場所に完全なバンドルがすでに存在する場合は、再ダウンロードせずに再利用します。

任意：

• --val-csv <path> — 検証用 CSV（別ファイル）。指定しない場合、準備ステップで入力を --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 が train/val の損失を記録します。run 名はプロジェクトと併せて指定した場合にのみ適用されます。デフォルトでは無効です（--resume での ckpt の wandb_run_id 継続処理とは独立しています）。
• --resume — 以下の「モード」セクションを参照してください。
• --gpus 0,2 — 使用する GPU を指定します。デフォルトでは、利用可能なすべての GPU を使用します。
• --from-prepare <dir> — prepare ステップをスキップし、<dir> 内にある既存の prepare_data.json を再利用します。ハイパーパラメーターを変えながら繰り返し実行する場合に便利です。

モード

ランナーには、入力 ckpt を取り込むモードが 2 つあります。どちらを使うかは --resume が設定されているかどうかで決まります。目的に応じて選択してください。

デフォルト（新しいスケジュールでの事前学習継続）

使用する場面：事前学習済み ckpt を使い、新しいコーパスで、異なる目的関数で、または元の計画より多くのエポック数で学習を続けたい場合です。前回の学習のステップカウンターとスケジュール形状は関係なくなり、新しい実行用の学習率スケジュールが必要です。

ckpt から読み込まれるもの：

• ✓ モデルの重み（エンコーダー + vocab heads + contrast head + decoder。存在するものすべて）
• ✓ オプティマイザーの状態（Adam の実行中の m1/m2 モーメント — 新しいスケジュールのウォームスタートに使用され、最初の数百ステップがノイズの多い勾配推定の初期化に支配されるのを防ぎます）
• ✗ スケジューラーのステップカウンター（0 にリセット）
• ✗ エポックカウンター（0 にリセット）
• ✗ バッチカウンター（0 にリセット）
• ✗ wandb の run id（継続ではなく、新しい wandb run を開始）

スケジュール形状（init/max/final LR、ウォームアップエポック数、合計エポック数）：CLI 引数または defaults_pretrain.json から取得します。新しい NoamLR がこれらの値から構築され、ステップ 0 から開始します。

--resume（真の再開）

使用する場面：以前の実行が中断され（クラッシュ、OOM、Ctrl-C）、まったく同じデータセット、同じスケジュール、同じ学習軌道で、停止した箇所から正確に再開したい場合です。

ckpt から読み込まれるもの：save_model_for_restart 形式のすべてです。モデルの重み、オプティマイザーの状態、scheduler_step、epoch、batch_idx、wandb_run_id がすべて復元されます。新しい実行は、保存されたスケジュール内の保存済みステップから続行します（スケジュールは ckpt の saved_args から復元されます）。エポック途中からの再開にも対応しています。pretrain_ddp.py のサンプラーのスキップ数により、保存されたエポック内の保存済みバッチインデックスから再開します。

スケジュールの形状：ckpt の saved_args から引き継がれます。スケジュールフラグ（--epochs / --warmup-epochs / --init-lr / --max-lr / --final-lr）を CLI で上書きすると、ハードエラーで拒否されます。純粋な再開は、あくまで純粋な再開です。スケジュールを変更する場合は、--resume を外して、新しいスケジュールで実行してください。

要件：ckpt は save_model_for_restart を使って保存されている必要があります（つまり、optimizer / scheduler_step / epoch / batch_idx のキーを含む必要があります）。これらのいずれかが欠けている場合、ランナーは明確なメッセージとともにエラーを出し、--resume を外すよう提案します。

デフォルトのモードは、約90%のケースで適切な選択です。--resume を使うのは、単一の中断された学習実行を本当に続行する必要がある場合に限ってください。

ワークフロー

$KERMT_REPO を kermt リポジトリのチェックアウト先のパスとし、kermt-setup によって kermt:latest がすでにビルドされているものとします。以下のパスはすべてホスト上のパスです。ヘルパーによって、既知のコンテナーパスにバインドマウントされます。

1. 事前確認：コンテナーとシステムプローブを確認します。

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

   gaps があればユーザーに提示します。ok: false の場合は続行しません。

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

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

3. チェックポイントを解決し、検証します。

   解決 — --ckpt が省略された場合のみ。 リリース済みの事前学習済みハイブリッドモデル nvidia/NV-KERMT-70M-v2 をデフォルトとして使用します。

   • 同意の確認。 --pretrained-release が指定されていない場合は、ユーザーに次のように尋ねます。 「チェックポイントが指定されていません。リリース済みモデル nvidia/NV-KERMT-70M-v2（NVIDIA Open Model License、https://huggingface.co/nvidia/NV-KERMT-70M-v2）をダウンロードし、そこから継続事前学習しますか？ [y/N]」 明示的な同意（または --pretrained-release）なしに、ダウンロードしてはいけません。--ckpt と --pretrained-release の両方が指定されている場合は、中止します。
   • 保存先。 デフォルトは $KERMT_REPO/models/NV-KERMT-70M-v2/ です。指定されている場合は --model-dir <dir> を使用します。すでに完全なバンドルがある場合は、それを再利用します。
   • ダウンロード（フォアグラウンドで実行。初回取得時は約282 MB）：

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

     JSON を解析し、ok: false の場合は中止して errors を提示します。成功した場合は、<user-ckpt> = <save-dir>/kermt_contrastive_v2.0.pt と設定します。バンドルに含まれる3つの語彙ファイルも <save-dir> に配置されるため、手順5の --vocab-dir 自動検出（ckpt の親ディレクトリを検索）で追加作業なしに検出されます。

   検証：解決済みの ckpt（またはユーザーが指定した ckpt）を検証します。

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

   JSON を解析します。ok: false の場合は、エラーをそのまま表示して中止します。エラーメッセージは、エンコーダー専用の ckpt の場合は kermt-add-cmim-pretrain に、ファインチューニング済みの ckpt の場合は kermt-finetune にユーザーを案内します。

4. データを検証します。

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

   ok: false の場合は中止します。

5. データを準備します（--from-prepare が指定されている場合はスキップします）。 ckpt の語彙を引き継いでください。 ckpt の親ディレクトリで、標準の pretrain_atom_vocab.{json,pkl}、pretrain_bond_vocab.{json,pkl}、pretrain_smiles_vocab.pkl ファイルを探します（リリースモデルのバンドル規約については、agent/README.md の「Released models」セクションを参照してください）。3 つすべてが存在する場合は、--vocab-dir <ckpt_parent_dir> を自動的に渡します。一部のみ存在する場合は、明示的なフラグ（--atom-vocab、--bond-vocab、--smiles-vocab）で渡します。どれも存在しない場合は、ユーザーに --vocab-dir を指定するよう依頼してください。指定がなければ処理を拒否してください。新しいコーパスから新たに語彙を構築すると、ckpt の語彙ヘッドと暗黙的に不一致が生じるためです（continue-pretrain では ckpt の語彙が正となります）。

   2 層マウントパターンに注意してください。ホストディレクトリを kermt_container.sh --vocab-dir に渡すと、コンテナ内の /vocab にマウントされます。内側の prepare_data.py コマンドでは /vocab を参照してください。同じパターンは、内側のコマンドで読み込む必要があるすべてのホストパスにも適用されます（--data <host-csv> → /data/<basename>、--ckpt <host-ckpt> → /ckpt）。

   VOCAB_DIR=$(dirname <user-ckpt>)
   $KERMT_REPO/agent/scripts/kermt_container.sh run \
       --data <user-csv> --vocab-dir $VOCAB_DIR --run-dir $RUN_DIR -- \
       "python agent/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]"
   

   出力は $RUN_DIR/data/prepare_data.json に保存され、vocab_source: "user_provided" となります。runner のステップ 7 では、語彙ファイルのエントリー数が ckpt の語彙ヘッドのサイズと一致することを検証し、不一致の場合は起動を拒否します。

6. 実行時間を見積もり、ユーザーに確認します。

   • 事前学習にかかる時間は、コーパスのサイズ × エポック数 × GPU 数によって異なります。
   • 見積もりをユーザーに伝え、--yes フラグが指定されている場合（エージェントの非対話モード）を除き、「続行しますか？」と尋ねてください。
   • 見積もりの例： ~N hours on K GPUs for E epochs over M molecules (~steps/epoch × seconds/step)。

7. runner をデタッチモードで起動します。

   $KERMT_REPO/agent/scripts/kermt_container.sh run_detached \\
       --name kermt-continue-pretrain-<ts> \\
       --ckpt <user-ckpt> --run-dir $RUN_DIR -- \\
       "python agent/scripts/run_pretrain_local.py \\
            --ckpt /ckpt \\
            --prepare-manifest /runs/data/prepare_data.json \\
            --out /runs \\
            [--epochs N --batch-size N --init-lr F ...]"
   

   コンテナ名、ID、ログファイルのパスが返されます。

8. ユーザーに報告します。 次の内容を簡潔にまとめて出力します。

   • コンテナ名と ID
   • $RUN_DIR/run.json（cmd_replay とイメージダイジェストを含むマニフェスト）
   • ログファイル：$RUN_DIR/logs/pretrain_ddp.log
   • TensorBoard：$RUN_DIR/logs/tb（tensorboard --logdir $RUN_DIR/logs/tb で開きます）
   • 進行状況を確認するには kermt-monitor <RUN_DIR> を実行するよう提案します。

厳守事項

• 同意なしにリリースモデルをダウンロードしてはいけません。 --ckpt が省略された場合、明示的にユーザーから「はい」と回答があった場合、または明示的に --pretrained-release フラグが指定された場合に限り、nvidia/NV-KERMT-70M-v2 をダウンロードします。--ckpt と --pretrained-release は相互に排他的です。
• ユーザーの入力 ckpt を変更してはいけません。 runner は ckpt を save_dir にシンボリックリンクします。pretrain_ddp.py はこのシンボリックリンクから自動再開します。元のファイルは変更されません。
• arch を黙って上書きしてはいけません。 ユーザーが指定した --hidden-size などが ckpt から導出した値と一致しない場合、runner はエラーを明示して中断します。arch パラメーターは ckpt の値を使用します。
• 実行時間の長い事前学習自体を待ってはいけません。 runner は run_detached を介して実行されるため、スキルはステップ 8 の後すぐに戻ります。進行状況の確認には kermt-monitor を使用してください。
• 適用したデフォルト値をユーザーに伝えてください。 run.json の args_applied フィールドには、各フラグの値とその出所（user / default-config / auto-1gpu / auto-multi-gpu）が記録されます。ユーザーが指定していないフラグについて、どのような値を仮定したかをスキルからユーザーに示してください。

よくあるエラー

• model_type='finetuned' rejected → ckpt は下流タスク向けのファインチューニング済みモデルであり、事前学習モデルではありません。このエラーは該当するワークフローに誘導します。
• grover_base ckpt has no vocab head → エンコーダーのみの ckpt です（例：original-grover の grover_base.pt）。このエラーは kermt-add-cmim-pretrain に誘導します。
• prepare_data manifest is missing required outputs → --skip-vocab または --skip-split を指定して prepare を実行したディレクトリに対して、ユーザーが --from-prepare を指定しました。これらのフラグを付けずに prepare を再実行してください。
• --gpus all not available → nvidia-container-toolkit をインストールし、kermt_container.sh check_system を確認してください。

再現性

run.json の cmd_replay フィールドには、同じ入力、ハイパーパラメーター、アーキテクチャで事前学習を再実行する1行のコマンドが記載されています。再実行するには、次のようにします。

# Inside the kermt container:
$(jq -r .cmd_replay $RUN_DIR/run.json)

マニフェストの ok_to_replay: false（起動時に kermt リポジトリの作業ツリーに未コミットの変更があったため）の場合、ビット単位で完全に一致しないことがあります。repo.commit フィールドで正確なコミットを指定し、先に git checkout してください。