Critical supervisor bug: completion check + --resume both used a glob that
only matched '{stage}_step*.pt', so a finished '{stage}_final.pt' (whose
step lives inside the file, not the filename) was invisible. Symptom: after
stage1 hit 31250 and saved final.pt, the supervisor re-launched --resume,
which picked step8000 (oldest sort-order quirk) and started OVERWRITING the
just-finished run from step 8000. Caught mid-overwrite; final.pt survived.
Fix: _latest_step (both stage1 & stage2) now globs step*.pt AND final.pt,
reads each one's 'step' field, returns the true highest. Supervisor's
run_until uses the same helper + reads step from the chosen ckpt to decide
completion. Verified: stage1(31250) now correctly reports 'done' and the
supervisor skips straight to stage2.
Also corrected supervisor stage1 config: bs=4 ga=8 (not bs=8 ga=4) — the
latter OOMs on long-sequence batches (logits 4.6GB alloc). expandable_segments
set to reduce fragmentation.
Full-scale retrain completed: stage1 31250 steps EMA 1.26, stage2 9375
steps EMA 0.75 (beats small-scale avg 0.81).
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.
张宗平