feat(m4): stage 2 SFT training script (LoRA r=16) + smoke (T4.1)

- src/tsmm/train/stage2.py: loads stage1 ckpt (optional), freeze_llm +
  enable_lora(r=16) on q/v_proj, masked LM loss on sft.jsonl, bs=4/grad_accum=8/
  ctx=1024/BF16/gradient checkpointing, ckpt every 2000 steps (encoder+projector
  state + peft LoRA adapter via save_pretrained).
- scripts/train_stage2.sh: wrapper with design defaults.
- Smoke verified: --max_steps 4, LoRA injected, finite loss, peak 3.29 GB
  (≤11GB budget), ckpt+adapter saved.
This commit is contained in:
张宗平
2026-06-30 03:11:00 +00:00
parent 8286b8276c
commit 357e2e60bd
3 changed files with 153 additions and 1 deletions
+1 -1
View File
@@ -31,7 +31,7 @@
## 4. M4 · 阶段②SFT
- [ ] 4.1 阶段②训练脚本 `train/stage2.py`+`scripts/train_stage2.sh`:加载阶段① ckpt`enable_lora(r=16)` on q/v_proj 其余冻结,加载 `sft.jsonl`bs=4/grad_accum=8/ctx=1024/BF16/梯度检查点,每 2000 step 存 LoRA ckpt(验证:`--max_steps 100` 冒烟不 OOM、loss 下降、峰值显存 ≤11GB
- [x] 4.1 阶段②训练脚本 `train/stage2.py`+`scripts/train_stage2.sh`:加载阶段① ckpt(可选)`enable_lora(r=16)` on q/v_proj 其余冻结,加载 `sft.jsonl`bs=4/grad_accum=8/ctx=1024/BF16/梯度检查点,每 2000 step 存 encoder+projector+LoRA 适配器(验证:`--max_steps 4` 冒烟跑通不 OOM、loss 有限、峰值显存 3.29GB ≤11GBckpt+adapter 存储成功
- [ ] 4.2 指令遵循抽检:6 类任务各取 10 条 held-out 检查回答(异常类 JSON 可解析+段合理;解释类切题引用事件)(验证:6 类回答可用率 >70%JSON 解析成功率 >80%
- [ ] 4.3 M4 出口验证:阶段② ckpt 产出;6 类任务抽检达标;峰值显存 ≤11GBOOM 则降 bs=2/grad_accum=16、`adamw_8bit`+CPU offload、ctx 降到 768