张宗平
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.
58 lines
2.5 KiB
Bash
Executable File
58 lines
2.5 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
# Crash-resilient full-scale training supervisor (stage1 → stage2).
|
|
#
|
|
# Both tsmm.train.stage1 and stage2 now support --resume: each saved checkpoint
|
|
# carries encoder/projector (+LoRA for stage2) + optimizer state + EMA loss, so
|
|
# a kill mid-run keeps all progress and the next --resume continues seamlessly.
|
|
#
|
|
# This supervisor loops "run --resume -> check step" until max_steps is reached.
|
|
# Each surviving iteration advances the step counter by >=1 (ckpt_every), so the
|
|
# run converges across multiple kills.
|
|
#
|
|
# Usage: scripts/train_full_supervisor.sh [stage1_max_steps] [stage2_max_steps]
|
|
set -uo pipefail # no -e: a killed train sub-call must NOT abort the loop
|
|
cd "$(dirname "$0")/.."
|
|
export PYTHONPATH=src PYTHONUNBUFFERED=1
|
|
|
|
S1_MAX="${1:-30000}" # 50万 align, bs=8 ga=4 => 1 epoch ≈ 15625 steps; 2 ep = 31250
|
|
S2_MAX="${2:-10000}" # 10万 sft, bs=4 ga=8 => 1 epoch ≈ 3125 steps; 3 ep = 9375
|
|
|
|
run_until() { # name max ckpt_dir module extra_args...
|
|
local name=$1 max=$2 ckpt_dir=$3 module=$4; shift 4
|
|
local extra=("$@")
|
|
while :; do
|
|
local cur=0
|
|
# read current step from the latest step ckpt's filename
|
|
local latest
|
|
latest=$(ls -1 "$ckpt_dir"/${name}_step*.pt 2>/dev/null | sort -V | tail -1)
|
|
if [ -n "$latest" ]; then
|
|
cur=$(echo "$latest" | grep -oE '[0-9]+' | tail -1)
|
|
elif [ -f "$ckpt_dir/${name}_final.pt" ]; then
|
|
# final carries step inside the file; assume done if present
|
|
cur=$max
|
|
fi
|
|
if [ "$cur" -ge "$max" ]; then
|
|
echo "[sup] $name done: step $cur/$max"; return 0
|
|
fi
|
|
echo "[sup] $name at step $cur/$max — launching --resume $(date '+%H:%M:%S')"
|
|
.venv/bin/python -m "$module" \
|
|
--ckpt_dir "$ckpt_dir" --max_steps "$max" --resume "${extra[@]}"
|
|
# loop: re-check (train may have been killed or completed)
|
|
done
|
|
}
|
|
|
|
echo "[sup] === TRAIN START $(date '+%H:%M:%S') ==="
|
|
echo "[sup] stage1 target=$S1_MAX stage2 target=$S2_MAX"
|
|
|
|
run_until stage1 "$S1_MAX" checkpoints/stage1 tsmm.train.stage1 \
|
|
--data data/align.jsonl --bs 8 --grad_accum 4 --ctx 512 --lr 1e-4 \
|
|
--lambda_contrast 0.1 --perturb_std 0.1 --log_every 50 --ckpt_every 2000 --tensorboard
|
|
|
|
run_until stage2 "$S2_MAX" checkpoints/stage2 tsmm.train.stage2 \
|
|
--data data/sft.jsonl --stage1_ckpt checkpoints/stage1/stage1_final.pt \
|
|
--bs 4 --grad_accum 8 --ctx 1024 --lr 1e-4 --lora_r 16 \
|
|
--log_every 50 --ckpt_every 2000 --tensorboard
|
|
|
|
echo "[sup] === TRAIN ALL DONE $(date '+%H:%M:%S') ==="
|
|
ls -la checkpoints/stage1/stage1_final.pt checkpoints/stage2/stage2_final.pt 2>/dev/null
|