kermt-continue-pretrain

kermt-continue-pretrain

从用户提供的 KERMT 检查点（grover_base / cmim / hybrid）继续预训练。此技能是工作流编排器：它会验证输入、 准备语料库、启动运行器，并返回运行目录。

硬件要求

• GPU：1–N 个支持 CUDA 的 NVIDIA GPU。运行器会通过 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 总线顺序与 CUDA 枚举顺序不同）。要指定特定的物理 GPU，请在调用前设置 CUDA_VISIBLE_DEVICES，或运行 python -c "import torch; print([torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())])" 以确认所选设备。

• 显存：默认的 --batch-size 256 面向 A100 级硬件 （80 GB 显存）。在较小的 GPU 上，应调低批次大小以避免 OOM：

  GPU 类型	显存	建议的 --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（每个检查点 都有数百 MB）。

• 驱动程序 / CUDA：任何支持 CUDA 12.6 的主机（kermt 镜像的基础版本）。 kermt-setup 会预先验证此要求。

输入

必填：

• --csv <path> — 预训练 CSV（单列 smiles）。如果有单独的 train/val CSV，也请传入 --val-csv <path>。

检查点（可选 — 如果省略，则默认为已发布的模型）：

• --ckpt <path> — 要继续训练的输入预训练检查点。必须是 grover_base（带有 vocab heads）、cmim 或 hybrid 检查点；验证器会拒绝其他所有类型，并重定向到正确的工作流。如果省略，该技能会提示你下载已发布的预训练 hybrid 模型 nvidia/NV-KERMT-70M-v2，并从该模型继续预训练 — 请参阅“解析并验证检查点”（工作流步骤 3）。已发布的模型包会将三个 vocab 文件与检查点一同提供，因此会自动完成权威 vocab 传递（步骤 5）。
• --pretrained-release — 显式选择使用已发布的模型，不显示交互式提示（适用于非交互式运行或 Agent 运行）。不能与 --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 会记录训练／验证损失；只有同时指定项目时，运行名称才会生效。默认关闭。（与 --resume 下检查点的 wandb_run_id 连续性处理相互独立。）
• --resume — 请参阅下方的“模式”部分。
• --gpus 0,2 — 限定使用部分 GPU。默认使用所有可见 GPU。
• --from-prepare <dir> — 跳过准备步骤，并复用 <dir> 中已有的 prepare_data.json。适用于反复调整超参数。

模式

运行器有两种载入输入检查点的模式，具体取决于是否设置了 --resume。请根据意图选择：

默认模式（使用全新计划继续预训练）

适用场景：你有一个已完成的预训练检查点，并希望继续训练它 — 可以使用新的语料库、不同的目标函数，或仅仅比原计划训练更多轮次。先前训练的步数计数器和计划曲线不再适用；你需要为新运行设置新的学习率计划。

从检查点载入的内容：

• ✓ 模型权重（编码器 + vocab heads + 对比头 + 解码器，包含检查点中已有的所有内容）
• ✓ 优化器状态（Adam 的运行中 m1/m2 矩 — 为新计划进行热启动，避免最初几百步被嘈杂的梯度估计启动过程主导）
• ✗ 调度器步数计数器（重置为 0）
• ✗ 轮次计数器（重置为 0）
• ✗ 批次计数器（重置为 0）
• ✗ wandb 运行 ID（启动新的 wandb 运行，而非延续原运行）

计划曲线（初始／最大／最终 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 中恢复）。也支持在 epoch 中途恢复——pretrain_ddp.py 的采样器跳过计数会从已保存 epoch 中的已保存 batch 索引处继续。

计划配置：继承自 ckpt 的 saved_args。任何计划标志（--epochs / --warmup-epochs / --init-lr / --max-lr / --final-lr）的命令行覆盖都会被硬性报错拒绝——纯粹的恢复就意味着纯粹的恢复；如果你想更改计划，请去掉 --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。模型包中的三个词汇表文件也会放在 <save-dir> 中，因此步骤 5 的 --vocab-dir 自动检测（会在 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”部分）。如果三个文件都存在， 则自动通过 --vocab-dir <ckpt_parent_dir> 传入。如果只有部分文件存在， 则通过显式标志传入（--atom-vocab、--bond-vocab、 --smiles-vocab）。如果一个文件都不存在，则向用户询问 --vocab-dir——或者拒绝继续，因为从新语料库重新构建词汇表会导致词汇表 与 ckpt 的词汇表头不匹配，且不会发出提示（继续预训练时，以 ckpt 的词汇表 为准）。

   请注意两层挂载模式：将宿主目录传给 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"。运行器的步骤 7 会验证词汇表文件的条目数是否 与 ckpt 的词汇表头大小一致；若不匹配，则拒绝启动。

6. 估算运行时间并向用户确认。

   • 预训练耗时取决于语料库大小 × 训练轮数 × GPU 数量。
   • 告知用户预估时间；除非提供了 --yes 标志（Agent 非交互模式），否则询问 “继续吗？”
   • 预估时间示例模板： 在 K 个 GPU 上，对 M 个分子训练 E 轮，约需 N 小时（每轮约 steps/epoch 步 × 每步 seconds/step 秒）。

7. 在后台启动运行器。

   $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。 运行器会将其符号链接到 save_dir；pretrain_ddp.py 会自动从该符号链接恢复训练。源文件保持不变。
• 绝不静默覆盖 arch。 如果用户传入的 --hidden-size 等参数与从 ckpt 得出的值不匹配，运行器会明确报错并中止。架构参数以 ckpt 为准，不可更改。
• 绝不等待耗时较长的预训练任务结束。 运行器通过 run_detached 调用； 该技能会在步骤 8 后立即返回。使用 kermt-monitor 查看进度。
• 向用户说明已应用的默认值。 run.json 的 args_applied 字段会记录每个标志的值及其来源 （user / default-config / auto-1gpu / auto-multi-gpu）。技能应汇总所有非用户指定标志的取值， 以便用户了解采用了哪些默认设置。

常见错误

• 拒绝 model_type='finetuned' → 该 ckpt 是下游微调模型，而非预训练模型。 错误信息会引导至相关工作流。
• grover_base ckpt has no vocab head → 仅包含编码器的 ckpt（例如原始的 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 不可用 → 安装 nvidia-container-toolkit；检查 kermt_container.sh check_system。

可重放性

run.json 中的 cmd_replay 字段是一条单行命令，可使用相同的输入、超参数和架构重新运行预训练。要重放：

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

如果清单中的 ok_to_replay: false（因为启动时 kermt 仓库的工作树有未提交更改），重放结果可能无法做到位级精确——请通过 repo.commit 字段指定确切的提交，并先执行 git checkout。