张宗平
77c2ada9b7
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).
68 lines
2.9 KiB
Bash
Executable File
68 lines
2.9 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 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
|
|
|
|
S1_MAX="${1:-31250}" # 50万 align, bs=4 ga=8 => eff batch 32; 1 ep ≈ 15625, 2 ep = 31250
|
|
S2_MAX="${2:-9375}" # 10万 sft, bs=4 ga=8 => 1 ep ≈ 3125, 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
|
|
# Ask the training module's own _latest_step logic for the true highest-step
|
|
# ckpt (it reads the 'step' field inside both step*.pt and final.pt).
|
|
local cur latest_path
|
|
latest_path=$(.venv/bin/python -c "
|
|
import sys; sys.path.insert(0,'src')
|
|
from tsmm.train.${module##*.} import _latest_step
|
|
import os
|
|
p=_latest_step('$ckpt_dir', '${name}_')
|
|
print(p if p else '')
|
|
" 2>/dev/null)
|
|
if [ -n "$latest_path" ] && [ -f "$latest_path" ]; then
|
|
cur=$(.venv/bin/python -c "
|
|
import torch
|
|
st=torch.load('$latest_path', map_location='cpu', weights_only=False)
|
|
print(int(st.get('step',0)))
|
|
" 2>/dev/null)
|
|
cur=${cur:-0}
|
|
else
|
|
cur=0
|
|
fi
|
|
if [ "$cur" -ge "$max" ]; then
|
|
echo "[sup] $name done: step $cur/$max (from $(basename "$latest_path"))"; 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 4 --grad_accum 8 --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
|