fix(train): _latest_step must consider final.pt (avoid re-running finished stage)

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

@@ -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
+S1_MAX="${1:-31250}"   # 50万 align, bs=4 ga=8 => eff batch 32; 1 ep ≈ 15625, 2 ep = 31250
-S2_MAX="${2:-10000}"   # 10万 sft,   bs=4 ga=8 => 1 epoch ≈ 3125 steps;  3 ep = 9375
+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
+    # Ask the training module's own _latest_step logic for the true highest-step
-    # read current step from the latest step ckpt's filename
+    # ckpt (it reads the 'step' field inside both step*.pt and final.pt).
-    local latest
+    local cur latest_path
-    latest=$(ls -1 "$ckpt_dir"/${name}_step*.pt 2>/dev/null | sort -V | tail -1)
+    latest_path=$(.venv/bin/python -c "
-    if [ -n "$latest" ]; then
+import sys; sys.path.insert(0,'src')
-      cur=$(echo "$latest" | grep -oE '[0-9]+' | tail -1)
+from tsmm.train.${module##*.} import _latest_step
-    elif [ -f "$ckpt_dir/${name}_final.pt" ]; then
+import os
-      # final carries step inside the file; assume done if present
+p=_latest_step('$ckpt_dir', '${name}_')
-      cur=$max
+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 \

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."""
+    """Return the path of the checkpoint with the highest recorded step.
-    steps = sorted(glob.glob(os.path.join(ckpt_dir, f"{prefix}step*.pt")))
+
-    return steps[-1] if steps else None
+    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:

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")))
+    best_path = None
-    return steps[-1] if steps else 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: