kermt-pretrain-scratch

kermt-pretrain-scratch

从头开始，在用户提供的语料库上预训练一个全新的 KERMT 模型。适用于希望在自定义化学领域重新训练模型，而不是在某个已发布检查点的基础上继续训练的情况。开销显著高于 kermt-continue-pretrain——没有预热启动，因此损失曲线需要经过多个周期才能从头下降。

硬件要求

与 kermt-continue-pretrain 相同：

• GPU：1–N 个支持 CUDA 的 GPU。运行器会通过 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 总线顺序与 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，具体用量随训练周期数增加。

• 训练时间：这是最大的区别。使用 1,100 万分子语料库从头开始预训练 100 个周期，通常需要数天，即使使用多 GPU 服务器也一样。 该工具会在启动前给出耗时估算；请与用户确认。

何时调用

• 用户希望在自定义语料库上训练新模型（例如，针对已发布检查点未涵盖的特定化学领域）。
• 用户希望端到端复现预训练配置，而不依赖已发布的检查点。

若要继续训练现有的已发布检查点，请使用 kermt-continue-pretrain。若要为仅含编码器的 grover_base 检查点添加 cMIM 解码器，请使用 kermt-add-cmim-pretrain。

输入

必需：

• --csv <path> — 预训练语料库 CSV，包含 smiles 列。按惯例使用单个文件；暂不支持多文件语料库。若有单独的验证集，请使用 --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 小时。如果你的主要 目的是利用已有知识完成下游任务，请考虑改用已发布检查点运行 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）。如果命令行未指定，请询问用户。
• 启动前发出强烈警告。 从零开始预训练是开销最高的工作流。用户需要 知道自己将要投入什么。

常见错误

• --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 组读取的值 （并应用任何未来的命令行覆盖项）。

可重放性

与 continue-pretrain 相同：cmd_replay 是可复制粘贴的命令。如果 ok_to_replay: false，则启动时 kermt 仓库工作树处于脏状态——请先检查 repo.commit 并执行 git checkout。