kermt-finetune

kermt-finetune

在用户提供的带标签 CSV 上微调预训练的 KERMT 编码器。此技能负责协调工作流：验证 ckpt、验证数据、准备数据、在分离模式下启动运行程序，并返回运行目录和容器名称。

硬件要求

• GPU：默认使用 1 张（单 GPU）；传入 --gpus 0（或其他 ID）来选择 GPU。在多 GPU 主机上进行更快的训练时，传入 --num-gpus N（N>1），即可使用 N 张 GPU 运行数据并行 DDP；此时 --batch-size 表示每张 GPU 的批大小（全局有效批大小 = batch_size × N）。
• 显存：默认配置 batch_size 32 至少需要 8 GB。显存较小时也可以运行，但需要使用更小的批大小；传入 --batch-size N 进行设置。
• 磁盘：每次运行需要几 GB（检查点 + 特征 + 日志）。
• 驱动程序 / CUDA：主机需支持 CUDA 12.6（kermt 镜像的基础版本）。kermt-setup 会提前验证这一点。

输入

必填：

• --csv <path> — 带标签的 CSV。第一列为 smiles；其余每一列 都是一个目标。

检查点（可选 — 如未指定，则默认使用已发布的模型）：

• --ckpt <path> — 输入预训练检查点（grover_base / cmim / hybrid）。 验证器会拒绝已微调的检查点，并提示改用 kermt-infer。如未指定，该技能会询问是否下载已发布的预训练混合模型 nvidia/NV-KERMT-70M-v2，并基于该模型进行微调 — 请参阅“解析并验证检查点”（工作流步骤 3）。
• --pretrained-release — 明确选择使用已发布的模型，无需交互式提示（适用于非交互式 / Agent 运行）。不能与 --ckpt 同时使用。
• --model-dir <dir> — 下载的模型包保存位置（默认 $KERMT_REPO/models/NV-KERMT-70M-v2/）。如果该目录中已有完整的模型包， 则会复用，而不会重新下载。

可选：

• --dataset-type {regression | classification | multiclass} — 默认 regression（来自 defaults_finetune.json）。决定损失函数、指标默认值 和输出头初始化方式。对于分类任务，请传入 --dataset-type classification。

• --targets COL [COL ...] — 显式指定目标列名称。如果未指定，验证器会自动检测数值型的非 smiles 列，并在继续之前向用户确认。

• --val-csv <path> 和 --test-csv <path> — 用户提供的验证集和测试集 划分。要么同时传入两者，要么都不传（该技能会使用已配置的 --split-type 自动划分）。

• --split-type {random | scaffold_balanced | index_predetermined} — 默认值来自 defaults_finetune.json，为 scaffold_balanced。

  • random 和 scaffold_balanced：从训练 CSV 内部构建验证集 / 测试集划分。无需传入 --val-csv / --test-csv。
  • index_predetermined：必须通过 --val-csv + --test-csv 传入预先划分的 CSV（另外还需提供每折索引文件 — 请参阅 kermt/util/utils.split_data）。当数据集自带标准划分时使用此选项（例如 tests/data/Biogen_for_grover/scaffold/ balance/<endpoint>/{train,val,test}.csv）。

• --metric NAME — mae（回归默认值）、auc（分类默认值）， 或 kermt.util.metrics.get_metric_func 接受的任意名称。

• --epochs N / --batch-size N / --init-lr F / --max-lr F / --final-lr F / --warmup-epochs F / --weight-decay F / --dropout F / --bond-drop-rate F / --dist-coff F / --early-stop-epoch N / --seed N — 训练超参数覆盖项。未指定的参数将从 agent/config/defaults_finetune.json 中填充。

• --ffn-hidden-size N / --ffn-num-layers N — 共享 FFN 主干维度。

• --ffn-num-task-specific-layers N / --ffn-task-specific-hidden-size H — 每个目标的 FFN 输出头（默认值为 0，即关闭；适用于目标异构的多目标 微调）。当 N > 0 时，两者必须同时设置。

• --ensemble-size N / --num-folds N — 多模型 / k 折交叉验证。两者默认值均为 1。

• --gpus 0 — 单进程微调所用的单个 GPU ID（默认值为 0）。当 --num-gpus > 1 时会被忽略。

• --num-gpus N — 数据并行 DDP 微调所用的 GPU 数量。默认值为 1 （单进程，行为不变）。N>1 时会以 WORLD_SIZE=N 运行 main.py finetune（每个 GPU 一个进程）；--batch-size 为每 GPU 的批次大小。

• --from-prepare <dir> — 跳过准备步骤，并复用 <dir> 中已有的 prepare_data.json。适用于调整超参数时反复尝试。

工作流

假设 $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); CUDA via container toolkit')
   "
   

   如果 ok: false，则拒绝继续。

2. 计算运行目录。

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

3. 解析并验证检查点。 解析 — 仅在省略了 --ckpt 时执行。 默认使用已发布的 预训练混合模型 nvidia/NV-KERMT-70M-v2：

   • 同意确认。 除非传入了 --pretrained-release，否则询问用户： "未提供 checkpoint — 是否下载已发布的模型 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。

   验证解析得到的（或用户提供的）ckpt：

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

   解析 JSON。若 ok: false 则中止。验证器会拒绝已 微调过的 ckpt（has_task_ffn: true），并重定向至 kermt-infer。

4. 验证数据。

   $KERMT_REPO/agent/scripts/kermt_container.sh run --data <user-csv> -- \
       "python agent/scripts/check_data.py --mode finetune --csv /data/<basename> [--targets COL1 COL2 ...]"
   

   如果用户未给出 --targets，则展示 JSON 中的 auto_detected_targets 并请用户确认后再继续。若 ok: false 则中止。

5. 准备数据（如果给出了 --from-prepare 则跳过）。

   预检：检查同级的 val.csv / test.csv。 在调用 prepare_data 之前，检查 <user-csv> 的父目录。如果训练 CSV 旁边 存在形似规范命名的同级 val.csv（或 val_*.csv — 常见变体 包括 val_T.csv、val_clean.csv）AND 匹配的 test.csv / test_*.csv，则该数据集自带 预定义的划分。在这种情况下，设置 --split-type index_predetermined AND 传入 --val-csv / --test-csv — 否则所配置的 split_type（默认为 scaffold_balanced）会从头重新划分训练 CSV， 并静默丢弃用户的 val/test 文件。若有疑问 — 或当同级文件 使用非规范后缀（_T、_v2 等）— 向用户展示该情况并询问 他们想要哪种方式。

   引用目标名称。 如果任何 --targets 列名 包含 shell 元字符（>, &, |, (, ), $ 等）， 在 CLI 上传递时对每个名称使用单引号，以免 shell 吞掉名称的一部分。示例：--targets 'Log_Caco2_Papp_A>B' 'logD'。CSV 表头本身由下游 训练器直接读取，不受影响，但 prepare_data.json 清单的 targets[] 字段会捕获 shell 所传递的内容 — 未加引号的 元字符会在那里被截断。

   挂载说明： kermt_container.sh --data <host-csv> 会将 <host-csv> 的父目录挂载到 /data。因此 --val-csv 和 --test-csv 必须引用该同一父目录中的文件。如果 val/test 位于单独的目录（例如同级的 splits/ 文件夹），请对目录 而非文件使用 --data <dir> 来挂载这三者的父目录。

   $KERMT_REPO/agent/scripts/kermt_container.sh run --data <user-csv> --run-dir $RUN_DIR -- \
       "python agent/scripts/prepare_data.py --mode finetune \\
            --csv /data/<basename> --out /runs/data \\
            --split-type <split_type> \\
            [--val-csv /data/<val-basename> --test-csv /data/<test-basename>] \\
            [--val-frac 0.1 --test-frac 0.1 --seed 0] \\
            --targets <COL1> [COL2 ...]"
   

   输出落在 $RUN_DIR/data/prepare_data.json。对于 scaffold_balanced 和 index_predetermined，prep 会生成单个 clean_full_csv + .npz； 运行器会将它们传递给 main.py finetune，后者会使用用户提供的 seed 在内部调用 split_data。

6. 估算运行时间 + 回显所应用的默认值。

   • 微调的挂钟时间在 1 块 GPU 上通常为几分钟到几小时。
   • 展示一份摘要，列出每个由默认值填充 与用户提供的标志，以便用户了解做了哪些假设。运行器 会将其记录在 args_applied 中。
   • 示例消息： "Filling from defaults_finetune.json: epochs=30, batch_size=32, split_type=scaffold_balanced. Override any of these with --<flag>."

7. 目标确认关卡（硬性要求）。 在启动运行器之前，无论目标列表是如何确定的（CLI --targets、第 4 步中的自动检测，或用户的自然语言请求，例如 “在 Caco2 和 HLM 上进行微调”），都要向用户回显最终目标列表并明确给出数量： "Will finetune on N target(s): COL1, COL2, ..."。如果用户请求指定的子集与此列表不符（例如，用户通过自然语言请求了 2 个任务，但列表中仍有 4 个），则视为存在差异，并再次提示用户，同时列出差异——绝不可在目标集合有误时默默继续。除非用户指定了 --yes，否则必须等待用户明确确认后才能启动。

8. 以分离模式启动运行器。（与预训练技能保持一致。）

   $KERMT_REPO/agent/scripts/kermt_container.sh run_detached \\
       --name kermt-finetune-<ts> \\
       --ckpt <user-ckpt> --run-dir $RUN_DIR -- \\
       "python agent/scripts/run_finetune_local.py \\
            --ckpt /ckpt \\
            --prepare-manifest /runs/data/prepare_data.json \\
            --dataset-type <type> \\
            --out /runs \\
            [--gpus 0] \\
            [--num-gpus N] \\
            [--epochs N --batch-size N --init-lr F ...] \\
            [--ffn-num-task-specific-layers N --ffn-task-specific-hidden-size H]"
   

   返回容器名称、ID 和日志文件路径。

9. 向用户报告。 输出简短摘要：

   • 容器名称和 ID
   • $RUN_DIR/run.json（包含 cmd_replay 和映像摘要的清单）
   • 日志文件：$RUN_DIR/logs/finetune.log
   • TensorBoard：$RUN_DIR/logs/tb（使用 tensorboard --logdir $RUN_DIR/logs/tb 打开）
   • 最终检查点保存于 $RUN_DIR/ckpt/fold_0/model_0/model.pt （best-val）和 last_checkpoint.pt（同级文件，用作自动恢复目标）。 留出测试集的预测结果和指标保存于 $RUN_DIR/ckpt/fold_0/test_result.csv。路径会随 --num-folds / --ensemble-size 而变化。
   • 查看进度：kermt-monitor <RUN_DIR>（单次查看）或 docker logs -f <container-name>（持续输出）。
   • 等待运行完成（适用于较短的测试运行）： docker wait <container-name> —— 完成时输出退出代码。

硬性规则

• 未经同意，绝不下载已发布模型。 未提供 --ckpt 时，只有在用户明确回答“是”或明确指定 --pretrained-release 标志后，才下载 nvidia/NV-KERMT-70M-v2。--ckpt 和 --pretrained-release 不能同时使用。
• 绝不修改用户的输入 ckpt。 运行器通过 --checkpoint_path 传入其路径；task/train.py 以只读方式将其加载到模型中，并附加一个新的 FFN 头。源文件保持不变。
• 架构取自 ckpt，而非 CLI/默认值。 运行器从 ckpt 的 saved_args 中提取 hidden_size、depth、num_attn_head、activation、embedding_output_type、 self_attention（以及适用时的 attn_hidden / attn_out）。此运行器没有 --hidden-size 标志。
• 绝不等待长时间运行的微调任务完成。 此技能通过 run_detached 启动任务，并在第 9 步后立即返回。请使用 kermt-monitor。
• 向用户回显已应用的默认值。 run.json 的 args_applied 字段会记录每个标志的取值及来源（user / default-config）。请用一行摘要列出所有采用默认值的标志，让用户了解系统采用了哪些假设。

常见错误

• finetune_init requires a pretrain ckpt (grover_base / cmim / hybrid) → 传入的 ckpt 已经过微调（包含任务 FFN 头）。请选择预训练 ckpt；如果要使用现有的微调模型进行预测，请使用 kermt-infer。如需在同一数据集上恢复微调，请绕过此技能，直接调用 python main.py finetune --checkpoint_path <ckpt> ...——Agent 技能不支持恢复训练，因为无法通过机器验证已保存的任务标识与新训练数据是否一致。
• prepare_data manifest reports ok=False → 检查 errors 中失败的步骤（通常是 clean_smiles 或 save_features）。修复后重新运行。
• ffn_num_task_specific_layers=N>0 but ffn_task_specific_hidden_size is unset → MTL 头需要显式指定隐藏层大小。请传入 --ffn-task-specific-hidden-size H。
• finetune is single-GPU（来自 --gpus 0,1）→ --gpus 为单进程微调选择单个设备。要使用多 GPU，请改用 --num-gpus N（DDP）。

可复现性

run.json 的 cmd_replay 字段是一条单行命令，可使用相同的输入、超参数和架构重新运行微调任务。要在 kermt 容器内重放：

$(jq -r .cmd_replay $RUN_DIR/run.json)

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