159 lines
6.7 KiB
Python
159 lines
6.7 KiB
Python
#!/usr/bin/env python
|
|
"""Start a training run on the GPU host and hand the GPU back when it ends.
|
|
|
|
Detached by design: a training run is a day long, and tying it to an SSH channel means
|
|
it dies with the connection -- or worse, finishes while the client stays blocked on a
|
|
pipe that never closes. Both happened while this was being built. The run becomes a
|
|
transient systemd unit on the host; this process only starts it, waits, and restores the
|
|
LLM services afterwards.
|
|
|
|
gpu_train.py --epochs 48 --steps 1500 # start and watch
|
|
gpu_train.py --status # what is it doing now
|
|
gpu_train.py --stop # stop it and restore the LLMs
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
import time
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
from suvi import gpubox
|
|
|
|
REMOTE = "/var/home/htpc/suvi"
|
|
LOG = f"{REMOTE}/train.log"
|
|
|
|
|
|
def command(args):
|
|
return (
|
|
"bash -c 'cd /work && python -u train.py "
|
|
f"--shards /work/shards --manifest /work/shards/manifest.json "
|
|
f"--out /work/runs/{args.name} "
|
|
f"--epochs {args.epochs} --batch {args.batch} --lr {args.lr} "
|
|
f"--base {args.base} --depth {args.depth} --workers {args.workers} "
|
|
f"--val-samples {args.val_samples} --eval-every {args.eval_every}"
|
|
+ (f" --steps-per-epoch {args.steps}" if args.steps else "")
|
|
+ (f" --max-steps {args.max_steps}" if args.max_steps else "")
|
|
+ (f" --overfit {args.overfit}" if args.overfit else "")
|
|
+ (f" --resume /work/runs/{args.name}/last.pt" if args.resume else "")
|
|
+ "'"
|
|
)
|
|
|
|
|
|
def main(argv=None):
|
|
parser = argparse.ArgumentParser(description=__doc__,
|
|
formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
parser.add_argument("--name", default="v4")
|
|
parser.add_argument("--epochs", type=int, default=20)
|
|
parser.add_argument("--steps", type=int, default=None,
|
|
help="steps per epoch (default: the whole train split)")
|
|
parser.add_argument("--val-samples", type=int, default=64)
|
|
parser.add_argument("--eval-every", type=int, default=250)
|
|
parser.add_argument("--max-steps", type=int, default=None,
|
|
help="stop after this many steps (the 300-step probe gate)")
|
|
parser.add_argument("--overfit", type=int, default=None,
|
|
help="run the overfit gate on this many frozen samples")
|
|
parser.add_argument("--batch", type=int, default=8)
|
|
parser.add_argument("--lr", type=float, default=1e-4)
|
|
parser.add_argument("--base", type=int, default=32)
|
|
parser.add_argument("--depth", type=int, default=3)
|
|
parser.add_argument("--workers", type=int, default=6)
|
|
parser.add_argument("--resume", action="store_true")
|
|
parser.add_argument("--status", action="store_true")
|
|
parser.add_argument("--since", type=int, default=None,
|
|
help="print history rows with step > SINCE, plus health")
|
|
parser.add_argument("--stop", action="store_true")
|
|
parser.add_argument("--watch", type=int, default=0,
|
|
help="seconds to wait for completion (0 = start and return)")
|
|
args = parser.parse_args(argv)
|
|
|
|
if args.status:
|
|
print(gpubox.tail(LOG, 30))
|
|
print("containers:", gpubox.run("podman ps --format '{{.Names}}'", check=False).strip())
|
|
print(f"GTT {gpubox.gtt_used() / 1024**3:.1f} GiB used")
|
|
return 0
|
|
|
|
if args.since is not None:
|
|
return _report_since(args)
|
|
|
|
if args.stop:
|
|
stopped = gpubox.stop_containers()
|
|
print(f"stopped {stopped or 'nothing'}")
|
|
gpubox.run("systemctl --user start llama-swap.service", check=False)
|
|
return 0
|
|
|
|
state = gpubox.stop_llms()
|
|
try:
|
|
name, log = gpubox.launch(command(args), mounts=[(REMOTE, "/work")], log=LOG)
|
|
print(f"launched {name}; log at {log}", flush=True)
|
|
if not args.watch:
|
|
print("Running detached. gpu_train.py --status to check, --stop to end.")
|
|
print("NOTE: the LLM services stay stopped until --stop or the run finishes.")
|
|
_paused_by_us(state)
|
|
return 0
|
|
deadline = time.time() + args.watch
|
|
while time.time() < deadline and gpubox.running(name):
|
|
time.sleep(30)
|
|
print(gpubox.tail(log, 40))
|
|
finally:
|
|
if args.watch:
|
|
gpubox.stop_containers()
|
|
print("restore failures:", gpubox.start_llms(state) or "none")
|
|
return 0
|
|
|
|
|
|
def _report_since(args):
|
|
"""New history rows past a step, plus the health a loss curve does not show.
|
|
|
|
A run can go non-finite, stall with the process still alive, or lose its
|
|
container entirely; each of those wastes hours if nobody looks. Returns 1 when
|
|
something is wrong, so a supervising loop can react without parsing prose.
|
|
"""
|
|
import json
|
|
|
|
raw = gpubox.run(f"cat {REMOTE}/runs/{args.name}/history.jsonl 2>/dev/null",
|
|
check=False)
|
|
rows = [json.loads(line) for line in raw.splitlines() if line.strip()]
|
|
fresh = [row for row in rows if row.get("step", 0) > args.since]
|
|
for row in fresh:
|
|
if row.get("untrained"):
|
|
print(f"step {row['step']:>6} val {row['val_loss']:.5f} "
|
|
f"({row['val_psnr']:.2f} dB) [untrained]")
|
|
else:
|
|
print(f"step {row['step']:>6} train {row['train_loss']:.5f} "
|
|
f"val {row['val_loss']:.5f} ({row['val_psnr']:.2f} dB) "
|
|
f"grad {row.get('grad_mean', 0):.3f}/{row.get('grad_max', 0):.2f}")
|
|
|
|
problems = []
|
|
for row in fresh:
|
|
for key in ("train_loss", "val_loss", "val_psnr"):
|
|
if key in row and not (row[key] == row[key] and abs(row[key]) < 1e30):
|
|
problems.append(f"step {row['step']}: {key} is not finite")
|
|
containers = gpubox.run("podman ps --format '{{.Names}}'", check=False).strip()
|
|
if "suvi" not in containers:
|
|
problems.append("no suvi container is running")
|
|
disk = gpubox.run(f"df --output=pcent {REMOTE} | tail -1", check=False).strip()
|
|
if disk.rstrip("%").strip().isdigit() and int(disk.rstrip("%").strip()) > 95:
|
|
problems.append(f"host disk at {disk}")
|
|
for problem in problems:
|
|
print(f"PROBLEM: {problem}")
|
|
return 1 if problems else 0
|
|
|
|
|
|
def _paused_by_us(state):
|
|
"""Record what was stopped, so a later --stop can put it back.
|
|
|
|
Written to disk rather than held in memory because the run outlives this process by
|
|
a day, and a restore that only works while the launcher is alive is not a restore.
|
|
"""
|
|
import json
|
|
|
|
path = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".gpu_paused.json")
|
|
with open(path, "w") as handle:
|
|
json.dump(state, handle)
|
|
print(f" paused state recorded in {path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|