diff --git a/Taskfile.yml b/Taskfile.yml index 60073d115..25d1b4083 100644 --- a/Taskfile.yml +++ b/Taskfile.yml @@ -47,6 +47,39 @@ tasks: cmds: - nice -n 10 uv run python scripts/micro_mint_rig.py --vehicle micro {{.CLI_ARGS}} + rig:gauntlet: + desc: | + pgw#1014 — EVERY mint shape the fleet runs, in one command, with a table. + Each variant is a full production cycle (spawn, export+AOTI, seal, + publish, second-process adopt + parity) on a model that compiles in + under a minute. + + Variants declare the outcome they EXPECT, because some exist to + demonstrate a REFUSAL — so the gauntlet reports agreement, and a flip in + either direction is news. Exit 0 = every variant matched. + + Cardless by default. For the GPU lane (real sm_89, no synthetic probe): + eval "$(./scripts/rig_gpu_env.sh --export-only)" + task rig:gauntlet -- --device cuda --python "$GEN_WORKER_RIG_GPU_PYTHON" + cmds: + - nice -n 10 uv run python scripts/rig_gauntlet.py {{.CLI_ARGS}} + + rig:gpu: + desc: | + pgw#983 — one rig cycle on the REAL card. Builds/refreshes the isolated + cu126 interpreter (see scripts/rig_gpu_env.sh: cu128 has no torch + 2.13.0, cu126 does) and runs with --device cuda, so the cell carries a + MEASURED sm_89 instead of a supplied one. + + The repo's own .venv and the cu130 pin are untouched; this is a second + interpreter, selected explicitly, never a default. + cmds: + - ./scripts/rig_gpu_env.sh > /dev/null + - | + eval "$(./scripts/rig_gpu_env.sh --export-only)" + nice -n 10 "$GEN_WORKER_RIG_GPU_PYTHON" scripts/micro_mint_rig.py \ + --device cuda {{.CLI_ARGS}} + rig:test: desc: | pgw#978 — the rig as a pytest row, including the full cycle (opt-in via diff --git a/changelog.d/pgw1014.md b/changelog.d/pgw1014.md new file mode 100644 index 000000000..26a1177b9 --- /dev/null +++ b/changelog.d/pgw1014.md @@ -0,0 +1,44 @@ +### Added + +- **`task rig:gauntlet` (pgw#1014) — every mint shape the fleet runs, in one + command, with a table.** Each variant is a full production cycle (real child + spawn, `torch.export` + AOTInductor, seal, publish over the real wire, + second-process discover + arm + parity) on a model that compiles in under a + minute. Variants declare the outcome they EXPECT, because some exist to + demonstrate a REFUSAL; the gauntlet reports agreement, and a flip in either + direction is news. Exit 0 = every variant matched. + - `micro` — plain lane, 3 entries, container input + trailing plain input. + - `micro-lora` — `lora_bucket=64`, 5 entries (adapter fork), the + branch-bearing graph and the lifted-arm gates (pgw#1001). + - `micro-4d` — pgw#998's shape: a 4-D latent with BOTH spatial axes dynamic, + so every matmul's M extent is NONLINEAR in the traced symbols. This is + z-image's declaration, and it re-proves pgw#998's fix on every run instead + of trusting a changelog. + - `micro-lora-plain-parent` — EXPECTED RED: a bucket-bearing cell offered to + a plain-lane parent is rejected at discovery with `lane_mismatch`, before + any arm runs. That is A1 step 8's cold-second-pod path, and the reason + attempt 27 must dispatch the adopting pod with the cell's own bucket. + +- **`task rig:gpu` and `scripts/rig_gpu_env.sh` (pgw#983) — GPU-real cycles on + this box.** An isolated second interpreter (`.venv-cu126`, shared across + worktrees) with **torch 2.13.0+cu126**: the fleet's exact torch version, on + an index whose CUDA runs on this box's 12.8 driver. **cu128 was the approved + index and is wrong — it has no torch 2.13.0, only 2.11.0**, and `torch` is an + axis `verify_declared` checks strictly, so a cu128 cell would be further from + the fleet, not closer. The repo's `.venv` and the cu130 pin are untouched; + the GPU interpreter is selected explicitly and is never a default. + + This retires the synthetic-`sm` caveat: cells minted through it carry a + MEASURED `sm_89` — the same compute capability as the fleet's L40S — instead + of a supplied one. Carve-out re-measured on the real card: 3.00 GiB mint + + 1.00 GiB adopt = 4.00 GiB of a 7.63 GiB card, the cap reaches the child + (`vram_cap_bytes = 3.00 GiB`), `CardCensus` reads `basis=sampled`, + `GEN_WORKER_HOST_MOVE_GUARD` untouched, and measured peak is **0.011 GiB**. + +### Fixed + +- The rig's adopt child places the pipeline on the mint's device. A code-only + cell binds constants from RESIDENT weights, so a CPU-resident parent adopting + a CUDA cell fails inside AOTI itself + (`update_constant_buffer_func_ ... API call failed`) — reachable only on a + real card, and the first thing the GPU lane found. diff --git a/examples/micro-diffusion/src/micro_diffusion/aot_declaration_4d.py b/examples/micro-diffusion/src/micro_diffusion/aot_declaration_4d.py new file mode 100644 index 000000000..b3cf44432 --- /dev/null +++ b/examples/micro-diffusion/src/micro_diffusion/aot_declaration_4d.py @@ -0,0 +1,75 @@ +"""The pgw#998 declaration: a 4-D latent grid, both spatial axes dynamic. + +This is the shape z-image declares (`H_lat` and `W_lat` on a 4-D latent under +`dynamic-collapse`) and the shape pgw#998 measured as unlowerable across the +mint's `torch.export.save`/`load` hand-off to its compile child: + + InductorError: LoweringException: RuntimeError: + ('unexpected None!', 512*s18*s57) + +pgw#998's fix (the export handoff carries the ShapeEnv's symbol values) landed +on master. This declaration exists so the gauntlet re-proves that on every run +instead of trusting the changelog — and so that if it ever regresses, it does +so on a 1.1 MB toy in 20 seconds rather than on a rented A100. + +Deliberately kept to ONE target and ONE fork-free class set: the point is the +nonlinear EXTENT, not entry count, and a second target would only add wall +clock to a variant whose whole job is a single lowering question. +""" + +from __future__ import annotations + +from gen_worker import ( + Compile, + Dim, + GraphClass, + Input, + register_export_declaration, +) + +FAMILY = "micro-4d" + +#: The two latent grids, as GRID extents (not token counts). +LATENT_ROWS = (32, 48) +VAE_SCALE = 8 +PIXEL_ROWS = tuple((n * VAE_SCALE, n * VAE_SCALE) for n in LATENT_ROWS) +COND_LEN = 16 +#: Batch is fixed here: the variable under test is the SPATIAL product, and a +#: fork would just double the compile for no extra coverage of it. +ARITY = 1 + + +def build_declaration() -> Compile: + return Compile( + family=FAMILY, + targets=("transformer",), + shapes=PIXEL_ROWS, + text_len=COND_LEN, + dims=( + Dim("N", carried_by=(("t", 0),)), + # BOTH spatial axes dynamic on ONE input — this is the whole + # point. Their product is the M extent of every matmul inside. + Dim("H_lat", carried_by=(("x", 1),), multiple_of=2), + Dim("W_lat", carried_by=(("x", 2),), multiple_of=2), + ), + classes=tuple( + GraphClass(dims={"N": ARITY, "H_lat": n, "W_lat": n}) + for n in LATENT_ROWS), + inputs=( + Input("x", shape=(("config", "in_channels"), "H_lat", "W_lat"), + repeat="N", dtype="float32"), + Input("t", shape=("N",), dtype="float32"), + Input("cond", shape=(COND_LEN, ("config", "cond_dim")), + repeat="N", dtype="float32"), + ), + shape_strategy="dynamic-collapse", + warm_changes_key=False, + ) + + +DECLARATION = build_declaration() + +register_export_declaration(DECLARATION, replace=True) + +__all__ = ["ARITY", "COND_LEN", "DECLARATION", "FAMILY", "LATENT_ROWS", + "PIXEL_ROWS", "VAE_SCALE", "build_declaration"] diff --git a/examples/micro-diffusion/src/micro_diffusion/main_4d.py b/examples/micro-diffusion/src/micro_diffusion/main_4d.py new file mode 100644 index 000000000..f6d43ebd9 --- /dev/null +++ b/examples/micro-diffusion/src/micro_diffusion/main_4d.py @@ -0,0 +1,85 @@ +"""The pgw#998 variant's worker function — same family shape, GRID latents. + +Separate from `main` because it is a different FAMILY (`micro-4d`) with a +different traced signature; sharing a module would put two families' endpoints +in one discovery namespace for no benefit. Everything else — catalog slot with +no code default, declaration registered at import, `ctx.slots` dereferenced +first — is deliberately identical to `main`, so a difference in outcome is a +difference in the SHAPE under test and nothing else. +""" + +from __future__ import annotations + +from typing import List + +import msgspec +import torch + +from gen_worker import Compile, RequestContext, Resources, Slot, endpoint +from gen_worker.families import GenerationDefaults, family + +from .aot_declaration_4d import ( # noqa: F401 — registers at import + ARITY, + COND_LEN, + DECLARATION, + FAMILY, + LATENT_ROWS, + PIXEL_ROWS, +) +from .pipeline import MicroGridPipeline + + +@family(FAMILY) +class Micro4dDefaults(GenerationDefaults, frozen=True): + steps: int = 2 + + +class Micro4dIn(msgspec.Struct): + prompt: str = "" + model: str = "" + + +class Micro4dOut(msgspec.Struct): + checkpoint: str = "" + shape: str = "" + + +@endpoint( + models={"pipeline": Slot(MicroGridPipeline, selected_by="model")}, + compile=Compile( + family=FAMILY, targets=("transformer",), shapes=PIXEL_ROWS, + text_len=COND_LEN), + resources=Resources(gpu=True), +) +class Generate4d: + def setup(self, pipeline: MicroGridPipeline) -> None: + self.pipe = pipeline + + def generate_4d( + self, ctx: RequestContext[Micro4dDefaults], data: Micro4dIn, + ) -> Micro4dOut: + resolved = ctx.slots["pipeline"] + grid = LATENT_ROWS[0] + device = self.pipe.device + config = self.pipe.config + generator = ctx.generator(998) + x: List[torch.Tensor] = [ + torch.randn(config.in_channels, grid, grid, generator=generator, + device=device, dtype=torch.float32) + for _ in range(ARITY) + ] + t = torch.full((ARITY,), 100.0, device=device, dtype=torch.float32) + cond: List[torch.Tensor] = [ + torch.randn(COND_LEN, config.cond_dim, generator=generator, + device=device, dtype=torch.float32) + for _ in range(ARITY) + ] + with torch.no_grad(): + out = self.pipe.transformer(x, t, cond) + return Micro4dOut( + checkpoint=str(resolved.ref.path), + shape=str(tuple(int(n) for n in out.shape))) + + +__all__ = ["DECLARATION", "FAMILY", "Generate4d", "Micro4dDefaults", + "Micro4dIn", "Micro4dOut"] diff --git a/examples/micro-diffusion/src/micro_diffusion/model.py b/examples/micro-diffusion/src/micro_diffusion/model.py index 2f612eae5..a217e0fed 100644 --- a/examples/micro-diffusion/src/micro_diffusion/model.py +++ b/examples/micro-diffusion/src/micro_diffusion/model.py @@ -197,4 +197,36 @@ def forward(self, latent: torch.Tensor) -> torch.Tensor: return torch.tanh(self.proj_out(self.norm(h))) -__all__ = ["MicroConfig", "MicroDecoder", "MicroDenoiser"] +__all__ = ["MicroConfig", "MicroDecoder", "MicroDenoiser", + "MicroGridDenoiser"] + + +class MicroGridDenoiser(MicroDenoiser): + """The pgw#998 shape: a 4-D latent GRID whose H and W are BOTH dynamic. + + Same weights and same blocks as :class:`MicroDenoiser` — only the + interface differs. Each ``x`` element is ``(C, H, W)`` and is flattened to + tokens INSIDE the forward, so every matmul's M extent becomes the PRODUCT + ``H*W``: an extent that is NONLINEAR in the traced symbols. + + That is z-image's declared shape (``H_lat`` and ``W_lat`` on a 4-D latent + under ``dynamic-collapse``), and it is what pgw#998 found unlowerable + across the mint's own `torch.export.save`/`load` hand-off to its compile + child. This class exists so the gauntlet keeps testing that seam on every + run rather than trusting a changelog. + """ + + def forward( # type: ignore[override] + self, + x: List[torch.Tensor], + t: torch.Tensor, + cond: List[torch.Tensor], + ) -> torch.Tensor: + tokens = [] + for grid in x: + channels, height, width = grid.shape + tokens.append( + grid.reshape(channels, height * width).transpose(0, 1)) + out = super().forward(tokens, t, cond) + channels, height, width = x[0].shape + return out.transpose(1, 2).reshape(len(x), channels, height, width) diff --git a/examples/micro-diffusion/src/micro_diffusion/pipeline.py b/examples/micro-diffusion/src/micro_diffusion/pipeline.py index b0392805e..ee14d96e9 100644 --- a/examples/micro-diffusion/src/micro_diffusion/pipeline.py +++ b/examples/micro-diffusion/src/micro_diffusion/pipeline.py @@ -13,7 +13,12 @@ import torch -from .model import MicroConfig, MicroDecoder, MicroDenoiser +from .model import ( + MicroConfig, + MicroDecoder, + MicroDenoiser, + MicroGridDenoiser, +) from .weights import SEED, load_config, load_state, materialize @@ -128,4 +133,20 @@ def __call__( return self.unpatchify(cells, grid) -__all__ = ["MicroConfig", "MicroPipeline"] +__all__ = ["MicroConfig", "MicroGridPipeline", "MicroPipeline"] + + +class MicroGridPipeline(MicroPipeline): + """The pgw#998 vehicle's pipeline: same weights, GRID-shaped transformer. + + Only `.transformer` differs — it is a :class:`MicroGridDenoiser`, so the + traced call takes 4-D latents and every matmul's M extent is the product + of two dynamic symbols. See `aot_declaration_4d`. + """ + + @classmethod + def from_pretrained(cls, path: str, **_kw: Any) -> "MicroGridPipeline": + base = MicroPipeline.from_pretrained(path) + grid = MicroGridDenoiser(base.config) + grid.load_state_dict(base.transformer.state_dict(), strict=True) + return cls(grid.eval(), base.decoder, source=base.source) diff --git a/scripts/rig_gauntlet.py b/scripts/rig_gauntlet.py new file mode 100644 index 000000000..058525717 --- /dev/null +++ b/scripts/rig_gauntlet.py @@ -0,0 +1,236 @@ +#!/usr/bin/env python3 +"""The GAUNTLET — every mint shape the fleet runs, on a model that compiles in +under a minute, in ONE command. + + task rig:gauntlet # every variant, cardless + task rig:gauntlet -- --device cuda + task rig:gauntlet -- --only micro,micro-lora + +Each variant is a FULL production cycle — resolve, handoff, real child spawn, +`torch.export` + AOTInductor, seal, publish over the real wire, and a SECOND +process discovering, arming and comparing every arm against eager. Not a smoke +test: the same machinery a pod runs, at ~15-60 s instead of ~95 minutes. + +WHY A TABLE AND NOT A PASS/FAIL. Some variants exist to demonstrate a REFUSAL +(a bucket-bearing cell offered to a plain-lane parent must be rejected). Those +are not failures, so each variant declares the outcome it EXPECTS and the +gauntlet reports agreement. A variant that flips in EITHER direction is news: +an expected-red going green means someone fixed something and did not say so. + +The exit code is 0 when every variant matched its expectation. +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Dict, List, Optional + +REPO = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(REPO / "tests")) +sys.path.insert(0, str(REPO / "src")) + +RIG = REPO / "scripts" / "micro_mint_rig.py" + +#: The order the table is printed in: cheapest and most fundamental first, so a +#: reader who stops after two rows has still seen the load-bearing ones. +ORDER = ( + "micro", + "micro-lora", + "micro-4d", + "micro-lora-plain-parent", +) + + +def run_one( + name: str, *, device: str, root: Path, python: str, timeout_s: float, +) -> Dict[str, Any]: + out = root / f"{name}.json" + # pgw#1014: DELETE it first. A variant the rig REFUSED (the load gate, a + # cardless `--device cuda`) exits before writing any json, and a stale file + # from a previous run then gets read as THIS run's result — which is how a + # CPU row was very nearly reported as a GPU one, complete with the previous + # cycle's timing and peak. A variant that did not run must look like it did + # not run. + out.unlink(missing_ok=True) + cmd = [ + python, str(RIG), "--vehicle", name, "--device", device, + "--clean", "--root", str(root / name), "--json", str(out), + ] + started = time.monotonic() + try: + proc = subprocess.run( + cmd, capture_output=True, text=True, timeout=timeout_s, + cwd=str(REPO)) + rc, tail = proc.returncode, (proc.stderr or proc.stdout)[-1500:] + except subprocess.TimeoutExpired: + rc, tail = 124, f"timed out after {timeout_s:.0f}s" + elapsed = time.monotonic() - started + + row: Dict[str, Any] = { + "vehicle": name, "rc": rc, "wall_s": round(elapsed, 1), + "green": rc == 0, "tail": tail, + # rc 2 is the rig's REFUSED exit (a precondition, never a verdict on + # the variant). It must not be scored as a red, because a red is a + # claim about the mint path and a refusal is a claim about the box. + "refused": rc == 2, + } + if not out.is_file(): + row["failed_leg"] = "did-not-run" + row["failed_why"] = ( + "REFUSED before any leg (load gate / device precondition) — " + "no result written" if row["refused"] else + "the rig wrote no result") + if out.is_file(): + try: + data = json.loads(out.read_text()) + except ValueError: + return row + row["cycle_s"] = data.get("total_s") + env = data.get("env") or {} + row["device"] = env.get("device") + row["sm"] = env.get("sm") + row["torch"] = env.get("torch") + row["covers"] = env.get("covers") + for leg in data.get("legs") or (): + facts = leg.get("facts") or {} + if leg["name"] == "mint-child": + row["entries"] = len(facts.get("packed_entries") or ()) + row["peak_gb"] = round( + float(facts.get("peak_vram_bytes") or 0) / 2 ** 30, 3) + row["synthetic_sm"] = bool(facts.get("synthetic_runtime")) + if leg["name"] == "adopt": + parity = facts.get("parity_max_abs") or {} + row["parity"] = (max(parity.values()) if parity else None) + row["arm_reason"] = facts.get("arm_reason") or "" + if not leg.get("ok") and "failed_leg" not in row: + row["failed_leg"] = leg["name"] + row["failed_why"] = _why(facts, leg) + return row + + +def _why(facts: Dict[str, Any], leg: Dict[str, Any]) -> str: + """The CLASSIFIED reason a leg failed, never a slice of raw stderr. + + A table cell wide enough for a stack trace tail is a table cell that says + nothing; the whole point of pgw#999 was that these refusals carry a class. + """ + if facts.get("arm_reason"): + return str(facts["arm_reason"]) + for blob in (facts.get("miss_log") or "", facts.get("detail") or "", + facts.get("stderr_tail") or ""): + marker = "rejected by class:" + if marker in blob: + return blob.split(marker, 1)[1].strip().split("\n")[0][:60] + detail = str(facts.get("detail") or leg.get("detail") or "").strip() + return detail.split("\n")[0][:60] or "(unclassified)" + + +def _verdict(row: Dict[str, Any], expect: str) -> str: + if row.get("refused"): + return "NOT-RUN" + actual = "green" if row["green"] else "red" + return "OK" if actual == expect else ("REGRESSED" if expect == "green" + else "NOW-GREEN") + + +def table(rows: List[Dict[str, Any]], vehicles: Dict[str, Any]) -> str: + head = (f"{'variant':<26} {'exp':<5} {'got':<5} {'':<9} {'cycle':>7} " + f"{'ent':>4} {'peak GB':>8} {'parity':>10} note") + lines = [head, "-" * len(head)] + for row in rows: + veh = vehicles[row["vehicle"]] + got = ("n/a" if row.get("refused") + else ("green" if row["green"] else "red")) + verdict = _verdict(row, veh.expect) + parity = row.get("parity") + note = "" + if not row["green"]: + note = f"{row.get('failed_leg', '?')}: {row.get('failed_why', '')}"[:64] + elif veh.expect == "red": + note = "expected a refusal and did NOT get one" + lines.append( + f"{row['vehicle']:<26} {veh.expect:<5} {got:<5} {verdict:<9} " + f"{(row.get('cycle_s') or row['wall_s']):>6.1f}s " + f"{row.get('entries', 0):>4} {row.get('peak_gb', 0):>8.3f} " + f"{(f'{parity:.2e}' if parity else '—'):>10} {note}") + return "\n".join(lines) + + +def main(argv: Optional[List[str]] = None) -> int: + parser = argparse.ArgumentParser(prog="rig_gauntlet") + parser.add_argument("--device", choices=("auto", "cuda", "cpu"), + default="auto") + parser.add_argument("--only", default="", + help="comma-separated subset of the variants") + parser.add_argument("--root", type=Path, default=Path(os.environ.get( + "PGW978_ROOT", str(Path.home() / ".cache" / "gen-worker" / "gauntlet")))) + parser.add_argument("--python", default=os.environ.get( + "GEN_WORKER_RIG_GPU_PYTHON", sys.executable), + help="interpreter to run each cycle with (the GPU lane needs the " + "cu126 venv — see scripts/rig_gpu_env.sh)") + parser.add_argument("--timeout", type=float, default=1800.0) + parser.add_argument("--json", type=Path, default=None) + args = parser.parse_args(list(argv) if argv is not None else None) + + from harness import rig_vehicles + + names = [n.strip() for n in args.only.split(",") if n.strip()] or list(ORDER) + unknown = [n for n in names if n not in rig_vehicles.VEHICLES] + if unknown: + print(f"unknown variant(s): {unknown!r} " + f"(known: {sorted(rig_vehicles.VEHICLES)!r})", file=sys.stderr) + return 2 + + args.root.mkdir(parents=True, exist_ok=True) + print(f"GAUNTLET — {len(names)} variant(s), device={args.device}, " + f"python={Path(args.python).parent.parent.name}\n") + + rows: List[Dict[str, Any]] = [] + for name in names: + veh = rig_vehicles.VEHICLES[name] + print(f" ... {name} ({veh.covers[:70]})", flush=True) + rows.append(run_one(name, device=args.device, root=args.root, + python=args.python, timeout_s=args.timeout)) + + print("\n" + table(rows, rig_vehicles.VEHICLES)) + + env0 = next((r for r in rows if r.get("device")), {}) + print(f"\ncard={env0.get('device', '?')} {env0.get('sm', '')} " + f"torch={env0.get('torch', '?')}") + print(f"covers={env0.get('covers', '?')}") + if any(r.get("synthetic_sm") for r in rows): + print("⚠ synthetic `sm` supplied on at least one variant — those cells " + "are PLUMBING artifacts and must not reach a shared namespace.") + print("⚠ parity is the SDK's cosine ladder plus this rig's max|delta|; the " + "GATE is cosine-based, so it is a directional-degradation " + "instrument, not a max-abs bound.") + + for row in rows: + veh = rig_vehicles.VEHICLES[row["vehicle"]] + if veh.expect == "red" and not row["green"]: + print(f"\n {row['vehicle']}: red AS EXPECTED — {veh.expect_note}") + + if args.json: + args.json.write_text(json.dumps( + {"rows": rows, + "expect": {n: rig_vehicles.VEHICLES[n].expect for n in names}}, + indent=2, sort_keys=True, default=str)) + + mismatched = [r for r in rows + if _verdict(r, rig_vehicles.VEHICLES[r["vehicle"]].expect) != "OK"] + if mismatched: + print(f"\nGAUNTLET: {len(mismatched)} variant(s) did NOT match " + f"expectation: {[r['vehicle'] for r in mismatched]!r}") + return 1 + print(f"\nGAUNTLET: all {len(rows)} variant(s) matched expectation.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/rig_gpu_env.sh b/scripts/rig_gpu_env.sh new file mode 100755 index 000000000..85e739c5f --- /dev/null +++ b/scripts/rig_gpu_env.sh @@ -0,0 +1,84 @@ +#!/usr/bin/env bash +# pgw#983 — build the rig's GPU interpreter environment, reproducibly. +# +# This box's driver is 570.211.01 (CUDA 12.8) and the repo pins +# torch==2.13.0+cu130, which needs a 580-series driver — so the DEFAULT rig +# runs cardless and supplies a synthetic `sm` (pgw#983). This script builds an +# ISOLATED second interpreter that can use the card for real. +# +# WHY cu126 AND NOT cu128. The cu128 index has NO torch 2.13.0 — it stops at +# 2.11.0. `torch` is one of the axes `aot_serve.verify_declared` checks +# STRICTLY, so a cu128 venv would mint cells two minor versions off the fleet: +# further away, not closer. cu126 carries 2.13.0, and CUDA minor-version +# compatibility runs a cu126 build on a 12.8 driver. Only the `cuda` axis +# differs from the fleet (12.6 vs 13.0), and that is stated wherever a +# GPU-minted cell is reported. +# +# WHAT IT DOES NOT DO: no driver change, no system package, nothing outside +# the venv. The repo's own `.venv` is untouched and stays the default. +# +# ./scripts/rig_gpu_env.sh # build/refresh, then print the exports +# eval "$(./scripts/rig_gpu_env.sh --export-only)" +set -euo pipefail + +REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +# The venv is ~6 GB, so every worktree SHARES the canonical checkout's copy +# rather than building its own. `--git-common-dir` names the main .git even +# from inside a linked worktree; its parent is that checkout. +MAIN="$(cd "$(git -C "$REPO" rev-parse --git-common-dir)/.." && pwd)" +VENV="${GEN_WORKER_RIG_GPU_VENV:-$MAIN/.venv-cu126}" +SP="$VENV/lib/python3.12/site-packages" +SHIM="$VENV/cuda-home" +INDEX="https://download.pytorch.org/whl/cu126" + +export_only=0 +[ "${1:-}" = "--export-only" ] && export_only=1 + +build() { + if [ ! -x "$VENV/bin/python" ]; then + echo "rig-gpu: creating $VENV" >&2 + uv venv "$VENV" --python 3.12 >&2 + fi + if ! "$VENV/bin/python" -c "import torch" 2>/dev/null; then + echo "rig-gpu: installing torch 2.13.0+cu126 (~6 GB)" >&2 + UV_HTTP_TIMEOUT=600 uv pip install --python "$VENV" --index-url "$INDEX" \ + "torch==2.13.0" >&2 + # The AOTI CUDA compile needs CUDA HEADERS and ptxas. The torch wheels + # ship runtime libs only, so these two carry the rest. `nvcc` itself is + # NOT required — inductor compiles the wrapper with g++ and the kernels + # through triton/ptxas — but CUDA_HOME must point at a tree that LOOKS + # like a toolkit, which is what the shim below is. + UV_HTTP_TIMEOUT=600 uv pip install --python "$VENV" \ + --index-url https://pypi.org/simple \ + "nvidia-cuda-nvcc-cu12==12.6.*" "nvidia-cuda-cccl-cu12==12.6.*" >&2 + # The SDK's own runtime deps, from PyPI so torch is not re-resolved to a + # different CUDA build. + UV_HTTP_TIMEOUT=600 uv pip install --python "$VENV" \ + --index-url https://pypi.org/simple \ + "grpcio>=1.82.1" "msgspec>=0.18.6" "protobuf>=7.35.0" "requests>=2.32.0" \ + "boto3>=1.41.0" "psutil>=7.0.0" "pyyaml>=6.0.0" "blake3>=1.0.0" \ + "huggingface-hub>=0.26.0" "gguf>=0.10.0" "tomli-w>=1.0.0" \ + "c2pa-python>=0.36" numpy safetensors pillow pytest >&2 + fi + # A CUDA_HOME-shaped tree assembled from the wheels. Rebuilt every run: it + # is only symlinks, and a stale one is a compile error three legs later. + rm -rf "$SHIM" + mkdir -p "$SHIM/include" + ln -sfn "$SP/nvidia/cuda_runtime/lib" "$SHIM/lib64" + ln -sfn "$SP/nvidia/cuda_nvcc/bin" "$SHIM/bin" + ln -sfn "$SP/nvidia/cuda_nvcc/nvvm" "$SHIM/nvvm" + # Headers come from THREE wheels and inductor needs all of them: the runtime + # headers include `crt/host_defines.h` (nvcc wheel) and `nv/target` (cccl + # wheel). Each was a separate compile failure, in that order. + for d in "$SP/nvidia/cuda_runtime/include" "$SP/nvidia/cuda_nvcc/include" \ + "$SP/nvidia/cuda_cccl/include"; do + [ -d "$d" ] || continue + for f in "$d"/*; do ln -sfn "$f" "$SHIM/include/$(basename "$f")"; done + done +} + +[ "$export_only" = 1 ] || build +[ -d "$SHIM/include" ] || build + +echo "export GEN_WORKER_RIG_GPU_PYTHON='$VENV/bin/python'" +echo "export CUDA_HOME='$SHIM'" diff --git a/tests/harness/rig_vehicles.py b/tests/harness/rig_vehicles.py index fd9abd2dc..e31de987c 100644 --- a/tests/harness/rig_vehicles.py +++ b/tests/harness/rig_vehicles.py @@ -59,6 +59,13 @@ class Vehicle: #: What this vehicle proves that the other does not — printed, so a #: reported cycle time is never read against the wrong vehicle. covers: str + #: The outcome this vehicle is EXPECTED to produce. A variant that exists + #: to demonstrate a refusal is not a failing test — but a variant whose + #: outcome flips IS news, in either direction, so the gauntlet compares + #: against this rather than against "green". + expect: str = "green" + #: Why, when `expect` is not "green". + expect_note: str = "" # --------------------------------------------------------------------------- @@ -187,7 +194,13 @@ def _micro_cell(bucket: int = 0) -> Any: # The adopting process rebuilds the pipeline from the SAME generated tree, # which is the whole point of deterministic weights: a second machine with the # seed has the bytes, and the snapshot digest agrees without a download. -pipe = MicroPipeline.from_pretrained(os.environ["PGW978_CHECKPOINT"]) +# The cell is minted for the device the mint ran on, and a code-only cell +# binds its constants from RESIDENT weights — so the adopting pipeline must +# live on that device or the bind fails inside AOTI itself +# (`update_constant_buffer_func_ ... API call failed`). Measured the first +# time this rig ran on a real card; a CPU-only cycle cannot reach it. +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" +pipe = MicroPipeline.from_pretrained(os.environ["PGW978_CHECKPOINT"]).to(DEVICE) config = pipe.config # pgw#999: a BUCKET-bearing cell is keyed on the branch-bearing lane, and @@ -207,12 +220,12 @@ def _feed(arity, tokens): """One legal call of each declared arm, built from a FIXED generator so the eager pass and the served pass see identical bytes.""" gen = torch.Generator().manual_seed(997) - x = [torch.randn(tokens, config.in_channels, generator=gen) + x = [torch.randn(tokens, config.in_channels, generator=gen).to(DEVICE) for _ in range(arity)] - t = torch.full((arity,), 100.0) - cond = [torch.randn(COND_LEN, config.cond_dim, generator=gen) + t = torch.full((arity,), 100.0, device=DEVICE) + cond = [torch.randn(COND_LEN, config.cond_dim, generator=gen).to(DEVICE) for _ in range(arity)] - lat = torch.randn(1, tokens, config.in_channels, generator=gen) + lat = torch.randn(1, tokens, config.in_channels, generator=gen).to(DEVICE) return x, t, cond, lat @@ -229,7 +242,7 @@ def _feed(arity, tokens): # reference captured from the object that is about to be armed is a reference # the arm can contaminate — and a parity number computed against a # contaminated reference is worse than no parity number. -ref = MicroPipeline.from_pretrained(os.environ["PGW978_CHECKPOINT"]) +ref = MicroPipeline.from_pretrained(os.environ["PGW978_CHECKPOINT"]).to(DEVICE) if int(getattr(cfg, "lora_bucket", 0) or 0): from gen_worker import compile_cache as _cc0 _cc0.apply_lora_execution_lane(ref, int(cfg.lora_bucket)) @@ -373,7 +386,160 @@ def _micro_checkpoint_bytes(tree: Path) -> int: ) -VEHICLES: Dict[str, Vehicle] = {v.name: v for v in (TINY, MICRO, MICRO_LORA)} +# --------------------------------------------------------------------------- +# micro-lora-plain-parent — pgw#999's design question, as a standing leg +# --------------------------------------------------------------------------- + +#: A `lora64` cell offered to a parent that boots on the PLAIN lane. This is +#: A1 step 8's "boot a SECOND pod cold and adopt", and it is EXPECTED to be +#: refused: `aot_cells.discover` filters candidates by the adopting pipeline's +#: own lane, so a bucket-less parent computes `lane=plain` and rejects the cell +#: with `lane_mismatch` BEFORE any arm runs — no `arm_aot` gate, no adopt +#: reason. Standing here so the day it silently starts passing is visible. +MICRO_LORA_PLAIN_PARENT = Vehicle( + name="micro-lora-plain-parent", + modules=(RUNTIME_HOOK, "micro_diffusion.main"), + function="generate", + family="micro-diffusion", + ref_path="cozy/micro-diffusion", + syspath=(str(MICRO_SRC),), + build_checkpoint=_micro_checkpoint, + checkpoint_bytes=_micro_checkpoint_bytes, + compile_cell=lambda: _micro_cell(MICRO_LORA_BUCKET), + # The MINT is bucket 64; the ADOPTER is bucket 0. That mismatch is the + # whole point of the leg. + adopt_source=_micro_adopt_source_for(0), + covers=("the pgw#999 design question: a bucket-bearing cell offered to a " + "parent on the plain lane"), + expect="red", + expect_note=("discovery rejects on `lane_mismatch` before arming — the " + "adopting pod must carry the cell's own lora bucket " + "(required line in attempt 27's choreography)"), +) + + + + +# --------------------------------------------------------------------------- +# micro-4d — pgw#998's shape: a NONLINEAR traced extent (z-image's declaration) +# --------------------------------------------------------------------------- + + +def _micro_4d_cell() -> Any: + from gen_worker.registry import CompileCell + + from micro_diffusion.aot_declaration_4d import COND_LEN, PIXEL_ROWS + + return CompileCell( + shapes=PIXEL_ROWS, targets=("transformer",), family="micro-4d", + regional=False, text_len=COND_LEN, dynamic=(), lora_bucket=0, + guidance_scales=(), text_lens=()) + + +_MICRO_4D_ADOPT = ''' +import json, logging, os, sys +for p in %(paths)r: + sys.path.insert(0, p) +logging.basicConfig(level=logging.INFO, stream=sys.stderr) +import harness.rig_runtime # noqa: F401 +import torch +from pathlib import Path +from gen_worker import aot_cells, aot_serve +from gen_worker.models import provision +from gen_worker.registry import CompileCell +from micro_diffusion.aot_declaration_4d import ARITY, COND_LEN, LATENT_ROWS, PIXEL_ROWS +from micro_diffusion.pipeline import MicroGridPipeline + +torch.set_num_threads(2) +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" +cfg = CompileCell( + shapes=PIXEL_ROWS, targets=("transformer",), family="micro-4d", + regional=False, text_len=COND_LEN, dynamic=(), lora_bucket=0, + guidance_scales=(), text_lens=()) +pipe = MicroGridPipeline.from_pretrained(os.environ["PGW978_CHECKPOINT"]).to(DEVICE) +ref = MicroGridPipeline.from_pretrained(os.environ["PGW978_CHECKPOINT"]).to(DEVICE) +config = pipe.config + + +def _feed(grid): + gen = torch.Generator().manual_seed(998) + x = [torch.randn(config.in_channels, grid, grid, generator=gen).to(DEVICE) + for _ in range(ARITY)] + t = torch.full((ARITY,), 100.0, device=DEVICE) + cond = [torch.randn(COND_LEN, config.cond_dim, generator=gen).to(DEVICE) + for _ in range(ARITY)] + return x, t, cond + + +# The row the artifact was NOT seeded on, so the derived range is exercised. +ARMS = [("transformer", LATENT_ROWS[0])] +eager = {} +with torch.no_grad(): + for name, grid in ARMS: + x, t, cond = _feed(grid) + eager[name] = pipe.transformer(x, t, cond).clone() + +cell = aot_cells.discover( + pipe, cfg, base_url=%(base)r, + worker_jwt=lambda: "local-rig-worker-jwt", + cache_dir=Path(%(cache)r)) +out = {"pid": os.getpid(), "ok": cell is not None} +if cell is not None: + meta = aot_serve.unpack_metadata(Path(cell.artifact)) + out.update({ + "cell_key": cell.cell_key, "family": cell.family, "ref": cell.ref, + "snapshot_digest": cell.snapshot_digest, + "artifact_bytes": Path(cell.artifact).stat().st_size, + "entries": sorted((meta.get("entries") or {})), + }) + outcome = provision.arm_aot(pipe, cfg, Path(%(cache)r), Path(cell.artifact), 0) + out["armed"] = bool(outcome) + out["arm_reason"] = str(getattr(outcome, "reason", "") or "") + out["arm_detail"] = str(getattr(outcome, "detail", "") or "")[:400] + if outcome: + deltas = {} + with torch.no_grad(): + for name, grid in ARMS: + x, t, cond = _feed(grid) + deltas[name] = float( + (pipe.transformer(x, t, cond) - eager[name]).abs().max()) + out["parity_max_abs"] = deltas + out["execution_count"] = int(aot_serve.execution_count(pipe)) + out["parity_ok"] = all(v <= 1e-4 for v in deltas.values()) + out["ok"] = bool(out["parity_ok"]) and out["execution_count"] > 0 + else: + out["ok"] = False +print("RIG_ADOPT " + json.dumps(out)) +''' + + +def _micro_4d_adopt_source(base: str, cache: Path) -> str: + return _MICRO_4D_ADOPT % { + "paths": [str(REPO / "tests"), str(REPO / "src"), str(MICRO_SRC)], + "base": base, "cache": str(cache)} + + +MICRO_4D = Vehicle( + name="micro-4d", + modules=(RUNTIME_HOOK, "micro_diffusion.main_4d"), + function="generate-4d", + family="micro-4d", + ref_path="cozy/micro-diffusion", + syspath=(str(MICRO_SRC),), + build_checkpoint=_micro_checkpoint, + checkpoint_bytes=_micro_checkpoint_bytes, + compile_cell=_micro_4d_cell, + adopt_source=_micro_4d_adopt_source, + covers=("pgw#998's shape — a 4-D latent with BOTH spatial axes dynamic, so " + "every matmul's M extent is NONLINEAR in the traced symbols. This " + "is z-image's declaration, and the seam that was unlowerable " + "across the mint's export save/load hand-off"), +) + + +VEHICLES: Dict[str, Vehicle] = { + v.name: v for v in (TINY, MICRO, MICRO_LORA, MICRO_LORA_PLAIN_PARENT, + MICRO_4D)} DEFAULT_VEHICLE = TINY.name @@ -386,4 +552,4 @@ def vehicle(name: str) -> Vehicle: __all__ = ["DEFAULT_VEHICLE", "MICRO", "MICRO_LORA", "MICRO_LORA_BUCKET", - "TINY", "VEHICLES", "Vehicle", "vehicle"] + "MICRO_4D", "MICRO_LORA_PLAIN_PARENT", "TINY", "VEHICLES", "Vehicle", "vehicle"]