From 77c2ada9b79535bb52962d5c279ec76ff145d3e3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BC=A0=E5=AE=97=E5=B9=B3?= Date: Wed, 1 Jul 2026 10:12:31 +0000 Subject: [PATCH] fix(train): _latest_step must consider final.pt (avoid re-running finished stage) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- scripts/train_full_supervisor.sh | 38 ++++++++++++++++++++------------ src/tsmm/train/stage1.py | 26 +++++++++++++++++++--- src/tsmm/train/stage2.py | 24 ++++++++++++++++++-- 3 files changed, 69 insertions(+), 19 deletions(-) diff --git a/scripts/train_full_supervisor.sh b/scripts/train_full_supervisor.sh index db7c67a..7379688 100755 --- a/scripts/train_full_supervisor.sh +++ b/scripts/train_full_supervisor.sh @@ -12,27 +12,37 @@ # 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 +export PYTHONPATH=src PYTHONUNBUFFERED=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True -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 +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 - 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 + # 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"; return 0 + 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" \ @@ -45,7 +55,7 @@ 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 \ + --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 \ diff --git a/src/tsmm/train/stage1.py b/src/tsmm/train/stage1.py index 476c95a..a212458 100644 --- a/src/tsmm/train/stage1.py +++ b/src/tsmm/train/stage1.py @@ -55,9 +55,29 @@ class EMA: def _latest_step(ckpt_dir: str, prefix: str) -> Optional[str]: - """Return the highest-numbered step checkpoint path in ``ckpt_dir``, or None.""" - steps = sorted(glob.glob(os.path.join(ckpt_dir, f"{prefix}step*.pt"))) - return steps[-1] if steps else None + """Return the path of the checkpoint with the highest recorded step. + + Compares stage1_step{N}.pt files AND stage1_final.pt (whose step lives inside + the file, not the filename) by reading each one's 'step' field. Without this + a finished final.pt (step=max) would be invisible to --resume and the + supervisor would restart from an older step ckpt, overwriting progress. + """ + best_path = None + best_step = -1 + candidates = list(glob.glob(os.path.join(ckpt_dir, f"{prefix}step*.pt"))) + final = os.path.join(ckpt_dir, f"{prefix}final.pt") + if os.path.exists(final): + candidates.append(final) + for path in candidates: + try: + st = torch.load(path, map_location="cpu", weights_only=False) + s = int(st.get("step", _step_number(path))) + except Exception: + s = _step_number(path) + if s > best_step: + best_step = s + best_path = path + return best_path def _step_number(path: str) -> int: diff --git a/src/tsmm/train/stage2.py b/src/tsmm/train/stage2.py index 6502ce9..9597ba6 100644 --- a/src/tsmm/train/stage2.py +++ b/src/tsmm/train/stage2.py @@ -43,9 +43,29 @@ class EMA: def _latest_step(ckpt_dir: str, prefix: str): + """Path of the checkpoint with the highest recorded step. + + Compares stage2_step{N}.pt AND stage2_final.pt by reading each 'step' field + (final.pt's step is inside the file). Prevents the supervisor from ignoring + a finished final and resuming from an older step ckpt. + """ import glob - steps = sorted(glob.glob(os.path.join(ckpt_dir, f"{prefix}step*.pt"))) - return steps[-1] if steps else None + best_path = None + best_step = -1 + candidates = list(glob.glob(os.path.join(ckpt_dir, f"{prefix}step*.pt"))) + final = os.path.join(ckpt_dir, f"{prefix}final.pt") + if os.path.exists(final): + candidates.append(final) + for path in candidates: + try: + st = torch.load(path, map_location="cpu", weights_only=False) + s = int(st.get("step", _step_number(path))) + except Exception: + s = _step_number(path) + if s > best_step: + best_step = s + best_path = path + return best_path def _step_number(path: str) -> int: