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:
Executable
+23
@@ -0,0 +1,23 @@
|
||||
#!/usr/bin/env bash
|
||||
# Stage ② SFT training (T4.1). Loads stage ① ckpt, enables LoRA(r=16) on q/v_proj.
|
||||
set -euo pipefail
|
||||
|
||||
DATA="${1:-data/sft.jsonl}"
|
||||
MAX_STEPS="${2:-10000}"
|
||||
STAGE1_CKPT="${3:-checkpoints/stage1/stage1_final.pt}"
|
||||
CKPT_DIR="${4:-checkpoints/stage2}"
|
||||
|
||||
cd "$(dirname "$0")/.."
|
||||
.venv/bin/python -m tsmm.train.stage2 \
|
||||
--data "$DATA" \
|
||||
--stage1_ckpt "$STAGE1_CKPT" \
|
||||
--ckpt_dir "$CKPT_DIR" \
|
||||
--max_steps "$MAX_STEPS" \
|
||||
--bs 4 \
|
||||
--grad_accum 8 \
|
||||
--ctx 1024 \
|
||||
--lr 1e-4 \
|
||||
--lora_r 16 \
|
||||
--log_every 20 \
|
||||
--ckpt_every 2000 \
|
||||
--tensorboard
|
||||
Reference in New Issue
Block a user