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).
This commit is contained in:
张宗平
@@ -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 \
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user