File size: 1,143 Bytes
9477b5c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 | """Tiny smoke test: 5 train steps + 1 eval pass on a tiny subset.
Verifies the pipeline runs end-to-end (forward, loss, backward, opt step,
eval). Should finish in < 2 minutes on a single GH200.
"""
import json
import sys
import time
from pathlib import Path
# Force tiny config overrides
overrides = {
"max_steps": 5,
"batch_size": 2,
"grad_accum": 1,
"data_train_size": 32,
"data_eval_size": 8,
"eval_every": 5,
"save_every": 0,
"log_every": 1,
"K_latents": 2,
"K_curriculum": [[0, 2]],
"output_dir": "/home/ubuntu/work/blt_smoke",
}
base_cfg_path = Path(__file__).resolve().parent.parent / "configs" / "pilot_qwen15b_gsm8k.json"
with open(base_cfg_path) as f:
cfg = json.load(f)
cfg.update(overrides)
tmp_cfg = Path("/tmp/blt_smoke_config.json")
tmp_cfg.write_text(json.dumps(cfg, indent=2))
print(f"Wrote {tmp_cfg}")
import subprocess
t0 = time.time()
r = subprocess.run([
sys.executable, "-u", "-m", "experiments.blt_reasoner.train",
"--config", str(tmp_cfg),
], cwd="/home/ubuntu")
print(f"smoke test exit={r.returncode} elapsed={time.time()-t0:.0f}s")
sys.exit(r.returncode)
|