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).

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:

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: