feat(m3): stage 1 alignment training script + smoke (T3.2)

- src/tsmm/train/stage1.py: streams align.jsonl, freeze_llm, trains
  Encoder+Projector (AdamW lr=1e-4), bs=8/grad_accum=4/ctx=512/BF16,
  gradient checkpointing + enable_input_require_grads, loss = lm + lambda*InfoNCE
  (original vs perturbed TS representation), ckpt + tensorboard every 2000 steps.
- scripts/train_stage1.sh: wrapper with design defaults.
- src/tsmm/data/collator.py: handle JSON null (missing markers) -> NaN -> fill.
- .gitignore: /checkpoints/ and data/*.jsonl.
- Smoke verified: --max_steps 5 on 64 samples, finite loss, peak 5.52 GB (≤6GB).
- 100 tests passing.
This commit is contained in:
张宗平
2026-06-30 02:22:32 +00:00
parent 204d5237ba
commit bf95ac18e4
5 changed files with 211 additions and 2 deletions
+1 -1
View File
@@ -25,7 +25,7 @@
## 3. M3 · 阶段①对齐训练
- [x] 3.1 损失函数 `train/losses.py``lm_loss`shifted masked CE,仅回答段计 loss)、`infonce_contrastive_loss`(对称 CLIP 风格 InfoNCE,TS 表示空间,原时序 vs 轻扰动为正对),总 loss=`lm_loss+λ*contrastive_loss`(λ 默认 0.1,由 stage1 组合)(验证:9 单测过——标量有限、全 mask 返 0、shift 正确、grad 可反传、相同对低 loss、对称、L2 不变性)。spec 澄清:InfoNCE 在 TS 嵌入空间(非回答嵌入)做,轻扰动作正对;T3.3 强扰动检验另行评估。
- [ ] 3.2 阶段①训练脚本 `train/stage1.py`+`scripts/train_stage1.sh`:加载 `align.jsonl` 流式,`freeze_llm=True` 仅训 Encoder+ProjectorAdamW, lr=1e-4),bs=8/grad_accum=4/ctx=512/BF16/梯度检查点,每 2000 step 存 ckpt + TensorBoard(验证:`--max_steps 100` 冒烟不 OOM、loss 下降、峰值显存 ≤6GB
- [x] 3.2 阶段①训练脚本 `train/stage1.py`+`scripts/train_stage1.sh`:加载 `align.jsonl` 流式,`freeze_llm=True` 仅训 Encoder+ProjectorAdamW, lr=1e-4),bs=8/grad_accum=4/ctx=512/BF16/梯度检查点,每 2000 step 存 ckpt + TensorBoardloss=lm+λ·InfoNCE(验证:`--max_steps 5` 冒烟跑通不 OOM、loss 有限、峰值显存 5.52GB ≤6GB)。注:collator 同时处理 JSON `null`(缺失标记)→ fill。
- [ ] 3.3 TS-token 有效性检验:held-out 原时序 vs 扰动时序回答变化率 + 「必看时序」样本答对率对照纯 LLM(验证:扰动后变化率 >50%,必看时序答对率显著高于纯 LLM)
- [ ] 3.4 M3 出口验证:阶段① ckpt 产出;扰动检验通过;峰值显存 ≤6GB(失败回查对比损失权重/数据质量)