张宗平
5e3030a20e
Addresses both points from review:
1) Resume (so a killed run converges across multiple restarts, like gen_supervisor):
- checkpoints now carry optimizer state + EMA value in a 'train_state' field
- --resume auto-picks the latest step ckpt and continues; --resume_from for explicit
- stage2 resume rebuilds the LoRA structure via enable_lora (same path as fresh)
then overlays saved adapter weights via set_peft_model_state_dict. Using
PeftModel.from_pretrained instead reset requires_grad and shrank the trainable
set 129->33 tensors, corrupting the optimizer state — verified and fixed.
- Smoke-verified: stage1 step4->10, stage2 step4->8, both restore optimizer +
EMA + LoRA with no state mismatch.
2) EMA loss logging (honest answer to 'is the loss swing a real problem?'):
- the old log printed the raw per-micro-batch loss every N steps; on
heterogeneous 6-task data that swings 2-3x purely from sampling (anomaly
JSON is low-loss, forecast numbers high-loss), so it looked 'unstable'
while the model may well have been converging.
- EMA(alpha=0.05) now logged alongside raw + as a tensorboard scalar
(total_loss_ema / sft_loss_ema). Read the EMA, not the raw, to judge
convergence. lr left unchanged (1e-4) pending EMA evidence.
Adds scripts/train_full_supervisor.sh: stage1->stage2 resume loop, same
crash-convergence guarantee as gen_full_supervisor.
26 lines
701 B
Bash
Executable File
26 lines
701 B
Bash
Executable File
#!/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")/.."
|
|
export PYTHONPATH=src PYTHONUNBUFFERED=1
|
|
.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
|
|
# Resume support: pass --resume to continue from the latest step ckpt + lora_adapter.
|