feat(train): crash-resilient resume + EMA loss logging (stage1 & stage2)

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.
This commit is contained in:
张宗平
2026-07-01 01:06:15 +00:00
parent 1a096ff989
commit 5e3030a20e
5 changed files with 262 additions and 26 deletions
+1
View File
@@ -22,3 +22,4 @@ export PYTHONPATH=src PYTHONUNBUFFERED=1
--log_every 20 \
--ckpt_every 2000 \
--tensorboard
# Resume support: pass --resume to continue from the latest step ckpt + lora_adapter.