kermt-add-cmim-pretrain

kermt-add-cmim-pretrain

将 grover_base 检查点（旧版原始 GROVER grover.encoders.* 或新版 kermt.encoders.*，可带或不带词表头）转换为完整的混合（cMIM + 词表）检查点，然后以混合模式在用户的语料上继续预训练。

这是一个轻量封装：upgrade_to_hybrid.py 会生成新的检查点，经 check_checkpoint.py 检查后，其类型为 model_type: hybrid；其余工作流与 kermt-continue-pretrain 完全相同。

状态：实验性。 此工作流从头到尾均可正常运行，但尚未与论文中从头开始的混合模型训练（该训练会生成已发布的检查点）进行基准测试。当你想扩展现有的 grover_base 检查点，而不是从随机初始化重新开始时，可将其作为 kermt-pretrain-scratch 的实验性替代方案。在依赖升级后的检查点进行生产工作之前，请使用你自己的基准测试验证下游性能。

硬件要求

与 kermt-continue-pretrain 相同（cMIM 解码器会增加参数，但增加幅度不大；显存余量应该足够）。升级步骤本身速度很快（约 5 s），且仅使用 CPU——只有后续的继续预训练会占用 GPU。

何时调用

• 用户有一个 grover_base 检查点（仅编码器或带词表头），并希望将其扩展为混合（词表 + cMIM 对比式）预训练模型。
• 适用于向预训练编码器添加 SMILES 重建对比目标，而无需从头重新开始预训练（kermt-pretrain-scratch 需要数天时间）。

如需继续训练现有的混合或 cmim 检查点，请直接使用 kermt-continue-pretrain。如需在自定义语料上训练新模型，请使用 kermt-pretrain-scratch。

输入

必需：

• --ckpt <path> — 要升级的 grover_base 检查点。通过 check_checkpoint.py --mode upgrade_to_hybrid 验证；如果检查点 已有对比头或任务 FFN，则会被拒绝。
• --csv <path> — 预训练语料 CSV。格式与 kermt-continue-pretrain 的 --csv 输入相同。

可选（与 kermt-continue-pretrain 相同）：

• --val-csv <path> — 单独的验证 CSV。如果未提供，prepare_data 会按 --val-frac 0.1 自动拆分。
• 训练超参数覆盖项（--epochs N、--batch-size N、lr 三元组、 --warmup-epochs F 等）。
• --vocab-loss-weight F / --latent-dim N / --contrastive-temperature F。
• --wandb-project NAME / --wandb-run-name NAME — 可选的 Weights & Biases 日志记录（仅在同时指定项目时才会使用运行名称）。默认关闭。
• --gpus 0,2。

工作流

将 $KERMT_REPO 设为 kermt 仓库检出目录的路径。

1. 预检：check_system（与 kermt-continue-pretrain 的步骤 1 相同）。

2. 计算运行目录：

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

3. 使用 check_checkpoint --mode upgrade_to_hybrid 验证输入 ckpt。 若 ok: false，则中止。验证器会拒绝已包含对比头的 ckpt（建议使用 kermt-continue-pretrain）或包含任务 FFN 头的 ckpt（该 ckpt 已经过微调；建议使用原始预训练检查点）。

4. 通过 check_data --mode pretrain 验证语料。 若 ok: false，则中止。

5. 使用 --mode pretrain 准备数据——不要使用 --vocab-dir。 升级过程会构建与语料词汇表大小相符的新词汇表头，因此我们希望 prepare_data 从语料构建新词汇表，而不是沿用 ckpt 的旧词汇表（仅编码器的旧版 grover_base ckpt 甚至可能没有旧词汇表）：

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

   输出清单包含 vocab_source: "built_fresh" 和 smiles_vocab（从语料构建，供新解码器使用）。

6. 升级 ckpt。

   $KERMT_REPO/agent/scripts/kermt_container.sh run --ckpt <user-ckpt> --run-dir $RUN_DIR -- \
       "python agent/scripts/upgrade_to_hybrid.py \\
            --ckpt /ckpt \\
            --prepare-manifest /runs/data/prepare_data.json \\
            --out /runs/upgraded.pt"
   

   向用户展示 JSON 摘要，尤其是 warnings[]，其中包括编码器架构偏差说明（例如，旧版 GROVER 有两个额外的 act_func_* 键，而现代 KERMTEmbedding 没有），以及当升级后 ckpt 的 backbone 不是 gtrans 时，pretrain_ddp.py 的 --backbone argparse 限制说明。

7. 估算运行时间并向用户确认。 使用与 kermt-continue-pretrain 相同的估算方法（语料大小 × 轮数 × GPU 数量 → 实际运行时间）。

8. 以分离模式启动运行器。

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

   运行器会将升级后的 ckpt 识别为 model_type: hybrid，因此会自动分派 --pretrain_mode hybrid --vocab_loss_weight 1.0，并传入 smiles_vocab。

9. 向用户报告升级后的 ckpt 路径，以及与 kermt-continue-pretrain 相同格式的 run.json 指针、日志路径和 tensorboard URL。

硬性规则

• 绝不修改用户的输入 ckpt。 升级会在 <run_dir>/upgraded.pt 写入新文件；源 ckpt 保持不变。
• 词汇表头始终重新构建。 即使输入的 grover_base 有词汇表头，也会丢弃并按新语料的词汇表大小重新构建。继续预训练升级后的 ckpt 时，这些新词汇表头会与解码器一同训练。
• 不要自动放宽 --backbone 的可选项。 如果输入 ckpt 的 backbone 不是 gtrans（例如旧版 dualtrans）而触发升级警告，请向用户说明警告并征求意见。不要悄悄修改 parsing.py，将旧版 backbone 加入可选项列表。

常见错误

• check_checkpoint rejected the ckpt，且 model_type=hybrid 或 cmim → 用户的 ckpt 已经有对比头。请引导用户使用 kermt-continue-pretrain。
• check_checkpoint rejected the ckpt，且 task_ffn=true → 该 ckpt 已经过微调。升级工作流仅支持预训练检查点。
• prepare manifest missing smiles_vocab → 调用 prepare_data 时使用了 --skip-vocab，或采用了其他省略 smiles 词汇表的方式。请在不使用这些标志的情况下重新运行 prepare。
• unexpected key(s) in encoder load 警告 → 旧版 GROVER 架构保存了几个现代 KERMTEmbedding 不使用的 act_func_* 权重。这不会造成影响；编码器其余部分已正确加载。

成功运行后 run.json 中的内容

包含与 kermt-continue-pretrain 相同的可复现性字段，此外，升级步骤的 summary.json 会通过 inputs.upgrade_summary 路径记录，以便审计升级后 ckpt 的来源。

可重放性

与 kermt-continue-pretrain 相同：cmd_replay 会重建 run_pretrain_local.py --ckpt <upgraded.pt> ... 调用。若要端到端重新执行完整的 add-cmim 流程，用户还需要输入 grover_base ckpt 和语料——二者的绝对路径均记录在 prepare_data 和升级清单中。