"""Resumable build and four-checkpoint evaluation, one GPU owner at a time.""" import argparse import fcntl import importlib.metadata import json import os import signal import subprocess import sys import time from common import ROOT, RUN, SOURCE, BASELINE, FAST, read_json, write_json, records, sha256, stamp from serve import server PYTHON = str(ROOT / "venv/bin/python") def stage(name, command, outputs, env=None): markers = RUN / "stages" marker = markers / (name + ".json") signature = {"command": command, "env": env or {}} signature["script_hashes"] = {part: sha256(ROOT / part) for part in command if part.endswith(".py") and (ROOT / part).is_file()} if marker.exists(): old = read_json(marker) assert old["signature"] == signature, f"Stage settings changed: {name}" assert all(p.exists() for p in outputs), f"Stage output missing: {name}" print("Already complete:", name, flush=True) return write_json(RUN / "status.json", {"stage": name, "state": "running", "started": stamp()}) print("Starting:", name, flush=True) log = RUN / "logs" / (name + ".log") log.parent.mkdir(parents=True, exist_ok=True) started = time.monotonic() with log.open("a") as out: process = subprocess.Popen(command, cwd=ROOT, env=dict(os.environ, OMP_NUM_THREADS="8", **(env or {})), stdout=out, stderr=subprocess.STDOUT, start_new_session=True) try: code = process.wait() if code: raise subprocess.CalledProcessError(code, command) except BaseException: try: os.killpg(process.pid, signal.SIGTERM) except ProcessLookupError: pass try: process.wait(timeout=45) except subprocess.TimeoutExpired: try: os.killpg(process.pid, signal.SIGKILL) except ProcessLookupError: pass process.wait(timeout=15) raise assert all(p.exists() for p in outputs), f"Missing outputs after {name}" write_json(marker, {"signature": signature, "completed": stamp(), "seconds": time.monotonic()-started, "outputs": {str(p.relative_to(ROOT)): p.stat().st_size for p in outputs}}) def build(): cal = RUN / "calibration" stage("baseline", [PYTHON, "swift15/checkpoint.py", "baseline"], [BASELINE / "hyperqwen-build.json"]) stage("generation", [PYTHON, "-u", "swift15/generate.py"], [cal / "generation-summary.json"]) assert len(records(cal / "gen.jsonl")) == len(records(cal / "prompts.jsonl")), "Incomplete generation" stage("vocabulary", [PYTHON, "swift15/quantize.py", "vocabulary"], [cal / "draft_vocab_ids.json"]) stage("capture", [PYTHON, "-u", "drafter/capture.py"], [cal / "hidden.npy", cal / "seqs.json"], {"MODEL": str(BASELINE), "CALIBRATION_DATA": str(cal)}) stage("lm-head", [PYTHON, "-u", "swift15/quantize.py", "lm-head"], [cal / "lm-head-report.json"]) stage("draft-head", [PYTHON, "prepare/build_draft_vocab.py", str(FAST), "--ids", str(cal / "draft_vocab_ids.json")], [FAST / "mtp_draft_vocab_ids.pt"]) stage("mtp-hessians", [PYTHON, "-u", "drafter/train_mtp.py", "--model", str(BASELINE), "--original-model", str(SOURCE), "--data", str(cal), "--out", str(cal / "mtp-replay"), "--eval-only", "1", "--dump-hessians", str(cal / "mtp_hessians.pt"), "--max-seqs", "400", "--val-frac", ".4", "--depths", "2", "--micro-tokens", "1", "--head-chunk", "128", "--draft-ids", str(cal / "draft_vocab_ids.json")], [cal / "mtp_hessians.pt"]) stage("mtp-quant", [PYTHON, "-u", "swift15/quantize.py", "mtp"], [cal / "mtp-report.json"]) stage("release", [PYTHON, "swift15/release.py"], [FAST / "hyperqwen-build.json", FAST / "hyperqwen-audit.json"]) def evaluate(full): evaluation_manifest = read_json(RUN / "evaluation/manifest.json") quality_concurrency = int(evaluation_manifest.get( "quality_concurrency", os.environ.get("SWIFT15_EVAL_CONCURRENCY", "8"))) if quality_concurrency < 1: raise ValueError("Quality concurrency must be positive") reference = ROOT / "models/Qwen3.8-27B-W4A16-AutoRound-fast-eval" stage("reference-embedding-repair", [PYTHON, "swift15/reference.py"], [reference / "embedding-repair.json"]) models = {"qwen-fast": reference, "swift10": ROOT / "models/Swift-Qwen3.8-27b-W4A16-AWQ", "swift15-baseline": BASELINE, "swift15-fast": FAST} for tag, model in models.items(): finished = RUN / "stages" / (tag + ("-full" if full else "-pilot") + ".json") if finished.exists(): continue with server(model, tag): stage(tag + "-speed", [PYTHON, "swift15/measure.py", tag, "--kind", "speed"], [RUN / "results" / tag / "speed-c1.json"]) stage(tag + "-ppl", [PYTHON, "swift15/measure.py", tag, "--kind", "ppl"], [RUN / "results" / tag / "perplexity.json"]) stage(tag + "-pilot-tasks", [PYTHON, "-u", "swift15/evaluate.py", tag + "-pilot", "--pilot"], [RUN / "results" / (tag + "-pilot") / "summary.json"]) if full: stage(tag + "-full-tasks", [PYTHON, "-u", "swift15/evaluate.py", tag + "-full", "--concurrency", str(quality_concurrency)], [RUN / "results" / (tag + "-full") / "summary.json"]) write_json(finished, {"completed": stamp()}) subprocess.run([PYTHON, "swift15/report.py"], cwd=ROOT, check=True) if os.environ.get("SWIFT15_SINGLE_USER_ONLY") == "1": return stage("runtime-compat", [PYTHON, "swift15/runtime_compat.py"], [RUN / "runtime-compatibility.json"]) for tag, int8 in [("swift15-fast-batch", False), ("swift15-fast-batch-int8", True)]: finished = RUN / "stages" / (tag + ".json") if finished.exists(): continue with server(FAST, tag, batch=True, int8=int8): stage(tag + "-speed", [PYTHON, "swift15/measure.py", tag, "--kind", "speed", "--concurrency", "8"], [RUN / "results" / tag / "speed-c8.json"]) stage(tag + "-ppl", [PYTHON, "swift15/measure.py", tag, "--kind", "ppl"], [RUN / "results" / tag / "perplexity.json"]) stage(tag + "-tasks", [PYTHON, "-u", "swift15/evaluate.py", tag, "--pilot", "--concurrency", "8"], [RUN / "results" / tag / "summary.json"]) write_json(finished, {"completed": stamp(), "activation_dtype": "INT8 MLP only" if int8 else "BF16"}) subprocess.run([PYTHON, "swift15/report.py"], cwd=ROOT, check=True) if (RUN / "calibration/manifest.json").exists(): subprocess.run([PYTHON, "swift15/release.py"], cwd=ROOT, check=True) def main(): ap = argparse.ArgumentParser() ap.add_argument("mode", choices=["build", "pilot", "full"]) ap.add_argument("--evaluation-only", action="store_true", help="Evaluate existing checkpoints in a separately prepared run directory") ap.add_argument("--wait-for-pid", type=int, help="Wait for an already-running calibration process before acquiring the GPU") args = ap.parse_args() def stop(signum, frame): raise KeyboardInterrupt("Pipeline terminated") signal.signal(signal.SIGTERM, stop) RUN.mkdir(parents=True, exist_ok=True) with (RUN / "pipeline.lock").open("w") as lock: fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) if args.wait_for_pid: print("Waiting for existing calibration process", args.wait_for_pid, flush=True) while os.path.exists(f"/proc/{args.wait_for_pid}"): time.sleep(5) hardware = subprocess.check_output(["nvidia-smi", "--query-gpu=name,uuid,memory.total,power.limit,driver_version", "--format=csv"], text=True) versions = {name: importlib.metadata.version(name) for name in ["torch", "vllm", "transformers", "safetensors", "compressed-tensors"]} commit = subprocess.check_output(["git", "rev-parse", "HEAD"],cwd=ROOT,text=True).strip() write_json(RUN / "environment.json", {"started": stamp(), "hardware": hardware, "versions": versions, "hyperqwen_commit":commit,"python":sys.version, "patches_sha256":{p.name:sha256(p) for p in (ROOT/"patches").glob("*.patch")}}) try: if not args.evaluation_only: build() if args.mode != "build": evaluate(full=args.mode == "full") truncated = {} for path in (RUN / "results").glob("*/summary.json"): result = read_json(path) if result.get("truncated_task_ids"): truncated[path.parent.name] = result["truncated_task_ids"] write_json(RUN / "status.json", {"state": "complete", "mode": args.mode, "completed": stamp(), "truncated_counted_as_wrong": truncated}) except BaseException as error: previous = read_json(RUN / "status.json") if (RUN / "status.json").exists() else {} write_json(RUN / "status.json", dict(previous, state="failed", error=repr(error), failed=stamp())) raise if __name__ == "__main__": main()