From 1b78deb107c16800b486c84a5c723aadbb1fb8b1 Mon Sep 17 00:00:00 2001 From: Eric Sutphen Date: Thu, 3 Sep 2026 20:47:07 -0400 Subject: [PATCH 1/6] issue #76 (inferswarm): execution harness for the Gemma heterogeneous numerical qualification Implements the minimum harness required to execute the frozen issue #74 methodology (inferswarm@f394dc9): case-driven single-arm reference runner (RTX 3090), three-stage chain runner (2x RTX 3060 node-01 + remote last stage on node-03 via a per-case R4 wire service), RowPruningSink capture (#71-compatible, host-side final-row pruning of the full BF16 logits matrix), o_proj checkpoint wrappers (layer-0 in/out, global-layer-15 attention projection), pure host-float64 reducer matching the frozen REDUCER.md identity, and torch-free unit/source-contract tests. No execution/model math is added or changed: all model execution flows through the accepted R6 GemmaDenseStage replay-prefill semantics. --- benchmarks/inferswarm_76/__init__.py | 258 +++++++++++++++ benchmarks/inferswarm_76/capture.py | 209 ++++++++++++ benchmarks/inferswarm_76/chain_runner.py | 256 +++++++++++++++ .../inferswarm_76/last_stage_service.py | 309 ++++++++++++++++++ benchmarks/inferswarm_76/reducer.py | 150 +++++++++ benchmarks/inferswarm_76/reference_runner.py | 225 +++++++++++++ benchmarks/inferswarm_76/stage_entry.py | 190 +++++++++++ tests/research/test_inferswarm_76_harness.py | 196 +++++++++++ 8 files changed, 1793 insertions(+) create mode 100644 benchmarks/inferswarm_76/__init__.py create mode 100644 benchmarks/inferswarm_76/capture.py create mode 100644 benchmarks/inferswarm_76/chain_runner.py create mode 100644 benchmarks/inferswarm_76/last_stage_service.py create mode 100644 benchmarks/inferswarm_76/reducer.py create mode 100644 benchmarks/inferswarm_76/reference_runner.py create mode 100644 benchmarks/inferswarm_76/stage_entry.py create mode 100644 tests/research/test_inferswarm_76_harness.py diff --git a/benchmarks/inferswarm_76/__init__.py b/benchmarks/inferswarm_76/__init__.py new file mode 100644 index 000000000..16c4e0981 --- /dev/null +++ b/benchmarks/inferswarm_76/__init__.py @@ -0,0 +1,258 @@ +"""Issue #76 (InferSwarm) execution harness: shared case/capture utilities. + +This package executes the frozen issue #74 numerical-equivalence methodology +(inferswarm @ f394dc9, docs/qualification/gemma4-12b-it-v1/) on the frozen +physical topology: + +- matched single-GPU FreeToken reference: inferswarm04 RTX 3090 (24 GiB); +- frozen three-stage RTX 3060 distributed candidate: node inferswarm01 + (2x RTX 3060, stages 1-2) + inferswarm03 (RTX 3060, last stage via the + accepted R4 wire service). + +It adds NO new execution math. All model execution flows through the accepted +R6 runtime (benchmarks.inferswarm_r6.stage_runtime.GemmaDenseStage) with the +accepted replay-prefill greedy semantics. The only additions are: + +1. case-driven corpora loading (exact frozen token IDs, never retokenized); +2. capture of the full 15-envelope checkpoint set at positions 0/1/3/7 by + ARMING the existing #71 capture sink seams plus thin out-of-tree wrappers + for the three checkpoints the R6/#71 seams do not emit + (layer-0 o_proj input/output, global-layer-15 attention o_proj output, + full final-row BF16 logits); +3. reference top-1 margin diagnostics per frozen selection input + ``matched-reference-top1-margin``; +4. host-float64 reduction exactly per frozen REDUCER.md. + +Execution-branch discipline (issue #76): this file and everything under +benchmarks/inferswarm_76/ plus its tests freeze as the physical implementation +producer BEFORE the first model execution. After that freeze no execution or +model math may change during the campaign. +""" + +from __future__ import annotations + +import hashlib +import json +import math +from pathlib import Path +from typing import Any, Iterable, Sequence + +# Frozen contract identity (mirrors inferswarm@f394dc9 methodology.json). +CONTRACT_ID = "inferswarm.gemma4-heterogeneous-numerical-equivalence/1" +METHODOLOGY_COMMIT = "f394dc9fbf9979574324f2d037580659f1d63b39" +FREETOKEN_BASE = "d4d16089165917704a87f4e2f0c4a09969646f95" + +EXPECTED_CHECKPOINT_SHA256 = ( + "5a84cb313260ac447237b890387116dfa8682e49a6b44bc585ae8353abbff18d" +) +MODEL_REVISION = "707f0a3b8a3c7ad586ed01e27eafbad8a27dd0f7" + +GENERATED_TOKENS = 8 +EXACT_TOKEN_POSITIONS = tuple(range(8)) +CAPTURE_POSITIONS = (0, 1, 3, 7) +RUNTIME_CAPACITY_TOKENS = 64 # frozen single replay chunk bound + +FAMILIES = ( + "local-bf16-backend-operation-output", + "hidden-residual-stream", + "final-normalized-hidden-state", + "bf16-logits", + "fp32-consumer-logits", +) +METRICS = ( + "max-absolute-difference", + "rms-difference", + "p99-absolute-error", +) +ENVELOPES = tuple(f"{family}:{metric}" for family in FAMILIES for metric in METRICS) + +# checkpoint_id -> (family, semantic_dtype) exactly per frozen +# manifests/checkpoint-family-map.json at the methodology commit. +CHECKPOINT_FAMILY_MAP = { + "embedding-output": ("local-bf16-backend-operation-output", "bfloat16"), + "layer-0-o-proj-input": ("local-bf16-backend-operation-output", "bfloat16"), + "layer-0-o-proj-output": ("local-bf16-backend-operation-output", "bfloat16"), + "global-layer-15-attention-o-proj-output": ( + "local-bf16-backend-operation-output", + "bfloat16", + ), + "post-global-layer-15-residual": ("hidden-residual-stream", "bfloat16"), + "post-global-layer-31-residual": ("hidden-residual-stream", "bfloat16"), + "post-global-layer-47-residual": ("hidden-residual-stream", "bfloat16"), + "final-normalized-hidden-state": ("final-normalized-hidden-state", "bfloat16"), + "full-final-row-bf16-logits": ("bf16-logits", "bfloat16"), + "full-final-row-fp32-consumer-logits": ("fp32-consumer-logits", "float32"), +} + +# Capture names emitted by the R6/#71 seams vs. the envelope checkpoint IDs. +# The #71 seams emit per stage; the harness maps stage-local records to the +# frozen global checkpoint IDs (single arm: all checkpoints on one device). +SEAM_TO_CHECKPOINT = { + "single": { + "embedding_output": "embedding-output", + "layer0_o_proj_input": "layer-0-o-proj-input", + "layer0_o_proj_output": "layer-0-o-proj-output", + "layer15_attn_o_proj_output": "global-layer-15-attention-o-proj-output", + "after_layer_15": "post-global-layer-15-residual", + "after_layer_31": "post-global-layer-31-residual", + "after_layer_47": "post-global-layer-47-residual", + "final_norm": "final-normalized-hidden-state", + "full_final_row_bf16_logits": "full-final-row-bf16-logits", + "final_row_fp32": "full-final-row-fp32-consumer-logits", + }, +} +# Chain arm: which stage-role emits which frozen checkpoints. +CHAIN_STAGE_CHECKPOINTS = { + "first": { + "embedding_output": "embedding-output", + "layer0_o_proj_input": "layer-0-o-proj-input", + "layer0_o_proj_output": "layer-0-o-proj-output", + "layer15_attn_o_proj_output": "global-layer-15-attention-o-proj-output", + "after_layer_15": "post-global-layer-15-residual", + "boundary_send_hidden": "post-global-layer-15-residual-boundary", + }, + "middle": { + "boundary_recv_hidden": "boundary1-recv", + "after_layer_31": "post-global-layer-31-residual", + "boundary_send_hidden": "post-global-layer-31-residual-boundary", + }, + "last": { + "boundary_recv_hidden": "boundary2-recv", + "after_layer_47": "post-global-layer-47-residual", + "final_norm": "final-normalized-hidden-state", + "full_final_row_bf16_logits": "full-final-row-bf16-logits", + "final_row_fp32": "full-final-row-fp32-consumer-logits", + }, +} +# Only checkpoints that feed envelopes (boundary probes are exact-layer only). +ENVELOPE_CHECKPOINT_IDS = frozenset(CHECKPOINT_FAMILY_MAP) + + +def canonical_json_bytes(value: Any) -> bytes: + """Byte-identical twin of the frozen inferswarm canonical JSON.""" + return ( + json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")) + + "\n" + ).encode() + + +def sha256_bytes(value: bytes) -> str: + return hashlib.sha256(value).hexdigest() + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + for chunk in iter(lambda: stream.read(1 << 20), b""): + digest.update(chunk) + return digest.hexdigest() + + +def verify_case_identity(case: dict[str, Any]) -> dict[str, Any]: + """Fail closed unless a corpus row still hashes to its frozen identity.""" + identity = { + "content_class": case["content_class"], + "length_regime": list(case["length_regime"]), + "prompt_text": case["prompt_text"], + "token_ids": list(case["token_ids"]), + } + prompt_sha = sha256_bytes(case["prompt_text"].encode("utf-8")) + ids_sha = sha256_bytes(canonical_json_bytes(list(case["token_ids"]))) + case_sha = sha256_bytes(canonical_json_bytes(identity)) + if prompt_sha != case["prompt_sha256"]: + raise ValueError(f"{case['case_id']}: prompt hash mismatch") + if ids_sha != case["token_ids_sha256"]: + raise ValueError(f"{case['case_id']}: token-ids hash mismatch") + if case_sha != case["case_sha256"]: + raise ValueError(f"{case['case_id']}: case hash mismatch") + if len(case["token_ids"]) != case["token_count"]: + raise ValueError(f"{case['case_id']}: token count mismatch") + return dict(case) + + +def load_corpus(path: Path, *, expected_count: int | None = None) -> list[dict]: + corpus = json.loads(Path(path).read_text()) + cases = [verify_case_identity(row) for row in corpus["cases"]] + if expected_count is not None and len(cases) != expected_count: + raise ValueError( + f"corpus {path}: expected {expected_count} cases, found {len(cases)}" + ) + return cases + + +def nearest_rank_higher(values: Sequence[float], percentile: float = 0.99) -> float: + """Frozen reducer tail rule: ascending sort, one-based ceil(0.99*N).""" + if not values: + raise ValueError("percentile domain must not be empty") + if not 0.0 < percentile <= 1.0: + raise ValueError("percentile out of range") + finite = [float(v) for v in values] + if not all(math.isfinite(v) and v >= 0.0 for v in finite): + raise ValueError("absolute-error inputs must be finite and nonnegative") + ordered = sorted(finite) + return ordered[math.ceil(percentile * len(ordered)) - 1] + + +def tensor_metrics( + reference: Sequence[float], candidate: Sequence[float] +) -> dict[str, float]: + """Frozen host-float64 tensor metrics over the complete domain.""" + if len(reference) != len(candidate): + raise ValueError("reference/candidate domain size mismatch") + if not reference: + raise ValueError("empty comparison domain") + errors = [] + for r, c in zip(reference, candidate): + e = abs(float(r) - float(c)) + if not math.isfinite(e): + raise ValueError("nonfinite absolute error in comparison domain") + errors.append(e) + square_sum = math.fsum(e * e for e in errors) + return { + "max-absolute-difference": max(errors), + "rms-difference": math.sqrt(square_sum / len(errors)), + "p99-absolute-error": nearest_rank_higher(errors, 0.99), + } + + +def conservative_case_family( + checkpoint_metrics: Iterable[dict[str, float]], +) -> dict[str, float]: + """Per-metric maximum across all declared checkpoints and positions.""" + result: dict[str, float] = {} + for row in checkpoint_metrics: + for metric in METRICS: + value = float(row[metric]) + if metric not in result or value > result[metric]: + result[metric] = value + missing = [m for m in METRICS if m not in result] + if missing: + raise ValueError(f"case family reduction missing metrics: {missing}") + return result + + +def envelopes_from_case_metrics( + per_checkpoint: dict[str, dict[str, float]], +) -> dict[str, str]: + """15 frozen envelope strings (exact hex binary64) for one case.""" + by_family: dict[str, list[dict[str, float]]] = {f: [] for f in FAMILIES} + for checkpoint_id, metrics in per_checkpoint.items(): + family = CHECKPOINT_FAMILY_MAP[checkpoint_id][0] + by_family[family].append(metrics) + envelopes: dict[str, str] = {} + for family, rows in by_family.items(): + if not rows: + raise ValueError(f"no checkpoint evidence for family {family}") + family_max = conservative_case_family(rows) + for metric in METRICS: + envelopes[f"{family}:{metric}"] = family_max[metric].hex() + if len(envelopes) != 15: + raise ValueError(f"expected 15 envelopes, built {len(envelopes)}") + return envelopes + + +def hex_to_float(value: str) -> float: + out = float.fromhex(value) + if not math.isfinite(out): + raise ValueError(f"nonfinite hex float {value!r}") + return out diff --git a/benchmarks/inferswarm_76/capture.py b/benchmarks/inferswarm_76/capture.py new file mode 100644 index 000000000..d59d49d4a --- /dev/null +++ b/benchmarks/inferswarm_76/capture.py @@ -0,0 +1,209 @@ +"""#76 capture arming: full 15-envelope checkpoint set from the R6 runtime. + +Arms the accepted #71 ``CaptureSink`` seam (``runtime._capture_sink``) and +adds the three checkpoints the frozen checkpoint-family map requires that the +R6/#71 seams do not emit: + +- ``layer0_o_proj_input`` (input rows of layer-0 attention o_proj GEMM) +- ``layer0_o_proj_output`` +- ``layer15_attn_o_proj_output`` (representative global-attention projection) +- ``full_final_row_bf16_logits`` (complete final-row BF16 logits; the #71 + seam emits the full [seq, vocab] matrix which the harness reduces to the + final row OFF-device, on host, without touching device math) + +All wrappers are instance-level (bound-method rebind), installed by the +harness process only, and add zero device-side math: each wrapper calls the +original method, then hands the EXACT native tensor to the sink for a host +copy. Un-instrumented execution is unchanged (wrappers absent). + +Checkpoint identity vs. the stage's own layer numbering: the wrappers key off +GLOBAL layer ids (``block.global_layer_ids``), so the same wrapper works for +the single arm (stage owns [0,48)) and chain stages owning [0,16)/[16,32). +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + + +def arm_full_capture(runtime, sink) -> None: + """Install the #76 checkpoint wrappers on one runtime instance.""" + import torch # noqa: F401 (presence check; wrappers run under it) + + block = runtime.block + global_ids = list(block.global_layer_ids) + + # --- layer-0 attention o_proj in/out (layer 0 must be owned) ---------- + if 0 in global_ids: + local0 = global_ids.index(0) + attn0 = block.layers[local0].self_attn + o_proj = attn0.o_proj + orig_forward = o_proj.forward + + def o_proj_forward(x, *args, **kwargs): + runtime._emit( + "layer0_o_proj_input", + x, + global_layer=0, + extra={"op": "o_proj_gemm_input"}, + ) + out = orig_forward(x, *args, **kwargs) + runtime._emit( + "layer0_o_proj_output", + out, + global_layer=0, + extra={"op": "o_proj_gemm_output"}, + ) + return out + + o_proj.forward = o_proj_forward + + # --- representative global-attention o_proj output (layer 15) --------- + if 15 in global_ids: + local15 = global_ids.index(15) + attn15 = block.layers[local15].self_attn + o_proj15 = attn15.o_proj + orig15 = o_proj15.forward + + def o_proj15_forward(x, *args, **kwargs): + out = orig15(x, *args, **kwargs) + runtime._emit( + "layer15_attn_o_proj_output", + out, + global_layer=15, + extra={"op": "global_attention_o_proj_output"}, + ) + return out + + o_proj15.forward = o_proj15_forward + + # --- full final-row BF16 logits (host-side row slice) ------------------ + # The runtime emits the full [seq, vocab] BF16 matrix at "bf16_logits". + # The harness post-processes that record to the final row on host; no + # device math is added here. + runtime._capture_sink = sink + + +def final_row_from_bf16_record(record_tensor): + """Return the exact final row of a captured full BF16 logits matrix.""" + if record_tensor.dim() != 2: + raise ValueError("bf16_logits capture must be a 2-D [seq, vocab] tensor") + return record_tensor[record_tensor.shape[0] - 1] + + +def reduce_capture_records( + records: list[dict[str, Any]], + *, + mapping: dict[str, str], +) -> dict[tuple[int, str], Any]: + """Group sink records by (capture_position, frozen_checkpoint_id). + + Returns raw host tensors keyed for the reducer. Records whose seam name + is not in the mapping are ignored (diagnostic-only seams). + """ + out: dict[tuple[int, str], Any] = {} + for record in records: + meta = record["meta"] + name = meta.get("checkpoint") + if name not in mapping: + continue + checkpoint_id = mapping[name] + position = int(meta["step"]) + out[(position, checkpoint_id)] = record["tensor"] + return out + + +class RowPruningSink: + """#76 per-case capture sink (#71 CaptureSink-compatible surface). + + Differences from the R6/#71 sink, all host-side only: + + - ``bf16_logits`` (full [seq, vocab] BF16 matrix) is reduced to its + FINAL ROW immediately after the host copy; the retained artifact is + the exact frozen ``full-final-row-bf16-logits`` checkpoint domain. + - tensors are NOT kept after ``save``; only hashes/metadata plus the + pruning-carried final row flow into the persisted bundle. + """ + + def __init__(self, *, role: str, gpu_uuid: str | None = None): + import socket + import time as _time + + self.role = role + self.host = socket.gethostname() + self.gpu_uuid = gpu_uuid + self.records: list[dict[str, Any]] = [] + self._now = _time.time + + def emit( + self, + *, + checkpoint: str, + step: int | None, + global_layer: int | None, + position_range: list[int] | None, + source_device: str, + tensor, + extra: dict[str, Any] | None = None, + ) -> None: + import hashlib + + import torch + + host_copy = tensor.detach().cpu() + if checkpoint == "bf16_logits": + host_copy = final_row_from_bf16_record(host_copy) + checkpoint = "full_final_row_bf16_logits" + raw = host_copy.detach().contiguous() + meta = { + "schema": "inferswarm.issue76.row-pruned-capture/1", + "checkpoint": checkpoint, + "step": step, + "global_layer": global_layer, + "position_range": position_range, + "source_device": source_device, + "role": self.role, + "host": self.host, + "gpu_uuid": self.gpu_uuid, + "captured_at": self._now(), + "shape": list(raw.shape), + "dtype": str(raw.dtype).replace("torch.", ""), + "byte_count": raw.numel() * raw.element_size(), + "sha256": hashlib.sha256( + raw.view(torch.uint8).numpy().tobytes() + ).hexdigest(), + "nan_count": ( + int(torch.isnan(raw).sum().item()) + if raw.dtype.is_floating_point + else 0 + ), + "inf_count": ( + int(torch.isinf(raw).sum().item()) + if raw.dtype.is_floating_point + else 0 + ), + } + if extra: + meta["extra"] = extra + self.records.append({"meta": meta, "tensor": raw}) + + def save(self, out_dir: str | Path, tag: str) -> dict[str, Any]: + import torch + + out = Path(out_dir) + out.mkdir(parents=True, exist_ok=True) + metas = [r["meta"] for r in self.records] + tensors = [r["tensor"] for r in self.records] + torch.save( + {"records": metas, "tensors": tensors}, + out / f"capture-{tag}.pt", + ) + self.records = [] + return { + "schema": "inferswarm.issue76.capture-manifest/1", + "out_dir": str(out), + "tag": tag, + "record_count": len(metas), + "record_sha256": [m["sha256"] for m in metas], + } diff --git a/benchmarks/inferswarm_76/chain_runner.py b/benchmarks/inferswarm_76/chain_runner.py new file mode 100644 index 000000000..39aa8b626 --- /dev/null +++ b/benchmarks/inferswarm_76/chain_runner.py @@ -0,0 +1,256 @@ +"""#76 chain-arm case runner: three-stage RTX 3060 distributed candidate. + +Drives the accepted R6 chain topology (stage 1 = node-01 GPU-0 layers [0,16), +stage 2 = node-01 GPU-1 layers [16,32), stage 3 = node-03 last stage via the +#76 R4 wire service) through the frozen replay-prefill greedy semantics, +per case: + + for step in 0..7: + RESET all stages; replay = prompt + generated[0:step] + stage1 PREFILL(token_ids) -> hidden -> stage2 PREFILL(hidden) + -> stage3 PREFILL(hidden) -> token (capture_step at 0/1/3/7) + +Capture handling mirrors the single arm: stages 1-2 use local +``RowPruningSink`` instances; the remote stage saves via CASE_SAVE. +Margins/NaN/Inf for the final row come from the last-stage response. + +Run from a #76 worktree on inferswarm01 with the last-stage service already +listening on inferswarm03. +""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +import time +from pathlib import Path + +from benchmarks.inferswarm_76 import ( + CAPTURE_POSITIONS, + GENERATED_TOKENS, + load_corpus, + verify_case_identity, +) + + +def _producer(repo: Path) -> dict: + sha = subprocess.check_output( + ["git", "-c", f"safe.directory={repo}", "-C", str(repo), + "rev-parse", "HEAD"], text=True).strip() + status = subprocess.check_output( + ["git", "-c", f"safe.directory={repo}", "-C", str(repo), + "status", "--porcelain"], text=True) + return {"commit": sha, "dirty": bool(status)} + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", required=True) + parser.add_argument("--plan", required=True) + parser.add_argument("--corpus", required=True) + parser.add_argument("--case-ids", default=None) + parser.add_argument("--last-stage-host", default="10.0.0.219") + parser.add_argument("--last-stage-port", type=int, default=18485) + parser.add_argument("--out-dir", required=True) + parser.add_argument("--tag", required=True) + parser.add_argument("--attempt-id", required=True) + args = parser.parse_args(argv) + + import multiprocessing + + repo = Path(__file__).resolve().parents[2] + producer = _producer(repo) + if producer["dirty"]: + print(json.dumps({"status": "BLOCKED_DIRTY_SOURCE"})) + return 2 + + raw = json.loads(Path(args.corpus).read_text()) + if isinstance(raw, list): + cases = [verify_case_identity(row) for row in raw] + else: + cases = load_corpus(args.corpus) + if args.case_ids: + wanted = set(args.case_ids.split(",")) + cases = [c for c in cases if c["case_id"] in wanted] + missing = wanted - {c["case_id"] for c in cases} + if missing: + raise SystemExit(f"unknown case ids: {sorted(missing)}") + + from benchmarks.inferswarm_76.stage_entry import I76StageClient + from benchmarks.inferswarm_r6.wire_client import RemoteLastStageClient + + plan = json.loads(Path(args.plan).read_text()) + shared = plan.get("declared_shared_state") + + context = multiprocessing.get_context("spawn") + stages = [] + try: + for index, block in enumerate(plan["blocks"][:-1]): + stages.append( + I76StageClient( + context, + role="first" if index == 0 else "middle", + adapter_data={ + **block, + "declared_shared_state": shared if index == 0 else None, + "runtime_capacity_tokens": 64, + }, + model_path=args.model, + gpu_index=index, + ) + ) + stages.append( + RemoteLastStageClient( + host=args.last_stage_host, + port=args.last_stage_port, + experiment_id=plan["digest"], + connect_timeout=600.0, + ) + ) + for stage in stages[:-1]: + ready = stage.recv() + if ready.get("op") == "ERROR": + raise RuntimeError(f"stage failed: {ready}") + + results = [] + out_root = Path(args.out_dir) + out_root.mkdir(parents=True, exist_ok=True) + + for case in cases: + case_dir = out_root / case["case_id"] + case_dir.mkdir(parents=True, exist_ok=True) + t0 = time.perf_counter() + + # arm per-case capture on stages 1-2 (fresh sink per case) + gpu_uuids = [ + subprocess.check_output( + ["nvidia-smi", "-i", str(i), "--query-gpu=uuid", + "--format=csv,noheader"], text=True).strip() + for i in (0, 1) + ] + for stage_index, (stage, uuid) in enumerate( + zip(stages[:-1], gpu_uuids) + ): + stage.request({ + "op": "CASE_ARM", + "out_dir": str(case_dir), + "tag": args.tag, + "gpu_uuid": uuid, + "after_layers": [15] if stage_index == 0 else [31], + }) + stages[-1].request({"op": "CASE_BEGIN", "case_id": case["case_id"]}) + + prompt = list(case["token_ids"]) + generated: list[int] = [] + margins: list[dict] = [] + nan_inf_total = 0 + for step in range(GENERATED_TOKENS): + for stage in stages: + stage.request({"op": "RESET"}) + replay = prompt + generated + capture_now = step in CAPTURE_POSITIONS + hidden = None + response: dict = {} + for index, stage in enumerate(stages): + if index == 0: + response = stage.request({ + "op": "PREFILL", + "token_ids": replay, + "position": 0, + **({"capture_step": step} if capture_now else {}), + }) + hidden = response.get("hidden") + else: + response = stage.request({ + "op": "PREFILL", + "hidden": hidden, + "position": 0, + **({"capture_step": step} if capture_now else {}), + }) + hidden = response.get("hidden") + token = response["token_id"] + if response.get("top1_index") is None: + raise RuntimeError( + "last-stage response missing margin diagnostics") + margin = { + "step": step, + "top1_index": response["top1_index"], + "top1_value_hex": response["top1_value_hex"], + "top2_index": response["top2_index"], + "top2_value_hex": response["top2_value_hex"], + "margin_hex": response["margin_hex"], + "nan_inf_count": response["nan_inf_count"], + } + margins.append(margin) + nan_inf_total += int(response["nan_inf_count"]) + generated.append(int(token)) + + # persist per-case captures from stages 1-2 and the remote stage + manifests = {} + for stage_index, stage in enumerate(stages[:-1]): + ack = stage.request({"op": "SAVE_CAPTURE", + "suffix": args.tag}) + manifests[f"stage{stage_index + 1}"] = ack["manifest"] + save_ack = stages[-1].request( + {"op": "CASE_SAVE", "tag": args.tag}) + manifests["stage3"] = save_ack["manifest"] + + summary = { + "schema": "inferswarm.issue76.chain-case-run/1", + "attempt_id": args.attempt_id, + "case_id": case["case_id"], + "case_sha256": case["case_sha256"], + "prompt_sha256": case["prompt_sha256"], + "token_ids_sha256": case["token_ids_sha256"], + "generated_token_ids": generated, + "step_margins": margins, + "min_top1_margin_hex": min( + float.fromhex(m["margin_hex"]) for m in margins).hex(), + "nan_inf_count": nan_inf_total, + "producer": producer, + "role": "chain", + "capture_positions": list(CAPTURE_POSITIONS), + "capture_manifests": manifests, + "wall_seconds": time.perf_counter() - t0, + } + path = case_dir / f"chain-case-{args.tag}.json" + if path.exists(): + raise SystemExit(f"refusing to overwrite {path}") + path.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n") + results.append(summary) + print(json.dumps({ + "case_id": case["case_id"], "status": "CASE_COMMITTED", + "tokens": generated, "nan_inf": nan_inf_total}), flush=True) + + index = { + "schema": "inferswarm.issue76.chain-run-index/1", + "attempt_id": args.attempt_id, + "producer": producer, + "tag": args.tag, + "case_count": len(results), + "cases": [ + {"case_id": r["case_id"], "case_sha256": r["case_sha256"], + "generated_token_ids": r["generated_token_ids"], + "nan_inf_count": r["nan_inf_count"], + "min_top1_margin_hex": r["min_top1_margin_hex"]} + for r in results + ], + } + index_path = out_root / f"index-{args.tag}.json" + if index_path.exists(): + raise SystemExit(f"refusing to overwrite {index_path}") + index_path.write_text(json.dumps(index, indent=2, sort_keys=True) + "\n") + print(json.dumps({"status": "RUN_COMPLETE", "cases": len(results)})) + return 0 + finally: + for stage in stages: + try: + stage.shutdown() + except Exception: # noqa: BLE001 + pass + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/inferswarm_76/last_stage_service.py b/benchmarks/inferswarm_76/last_stage_service.py new file mode 100644 index 000000000..f13f62082 --- /dev/null +++ b/benchmarks/inferswarm_76/last_stage_service.py @@ -0,0 +1,309 @@ +"""#76 remote last-stage service (inferswarm03) over the accepted R4 wire. + +Identical framing, boundary contract, and execution semantics to the accepted +R6 ``last_stage_service``. Differences (all evidence-side, no execution math): + +- per-CASE capture sinks (``RowPruningSink``) instead of one sink for the + whole connection — the 584-case campaign cannot hold all captures in host + RAM; each case's bundle is persisted at ``CASE_SAVE``; +- the o_proj/full-row capture wrappers from ``benchmarks.inferswarm_76``; +- ``--allow-producer`` records the #76 frozen producer explicitly (the plan + file carries the historical R6 producer; #76 is an authorized execution + campaign under a NEW frozen producer derived from the same base). + +The service speaks two extra request ops beyond the R6 protocol: + CASE_BEGIN {case_id} -> swaps in a fresh sink for the case + CASE_SAVE {tag} -> persists the case bundle, returns manifest +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import socket +import subprocess +import time +from pathlib import Path + +from freetoken.research.r4_wire import ( + WireError, + encode_frame, + payload_checksum, + recv_frame, + send_exact, + validate_request, +) + +WIRE_PROTOCOL_ID = "inferswarm.r4.boundary-wire/1" +MAX_TOKEN_COUNT = 64 # matches strategy PREFILL_CHUNK (single-chunk replays) +ROW_WIDTH = 3840 +BOUNDARY_CONTRACT = { + "dtype": "bfloat16", + "layout": "plane-major-contiguous", + "planes": 1, + "row_width": ROW_WIDTH, + "element_bytes": 2, + "max_token_count": MAX_TOKEN_COUNT, +} + + +def serve( + *, + listen_host: str, + listen_port: int, + participant_plan: str, + model_path: str, + gpu_uuid: str, + diagnostic: bool, + ready_file: str | None = None, + allow_producer: str | None = None, + out_dir: str, +) -> None: + os.environ["CUDA_VISIBLE_DEVICES"] = gpu_uuid + import torch + + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.layers.rotary import set_rope_device + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + set_rope_device(torch.device("cuda:0")) + + repo_root = Path(__file__).resolve().parents[2] + running_sha = subprocess.check_output( + ["git", "-c", f"safe.directory={repo_root}", "-C", str(repo_root), + "rev-parse", "HEAD"], text=True, + ).strip() + plan = json.loads(Path(participant_plan).read_text()) + plan_producer = plan.get("provenance", {}).get("r6", {}).get("producer_sha") + if plan_producer and running_sha != plan_producer: + if allow_producer and allow_producer == running_sha: + producer_check = { + "mode": "EXPLICIT_OVERRIDE_ISSUE76_EXECUTION", + "plan_frozen_producer": plan_producer, + "running_producer": running_sha, + "reason": "issue #76 authorized execution campaign under a " + "new frozen producer derived from the same base", + } + else: + raise RuntimeError( + f"last-stage running producer {running_sha!r} != plan's frozen " + f"producer {plan_producer!r}; pass --allow-producer " + f"{running_sha!r} for the #76 execution campaign" + ) + else: + producer_check = { + "mode": "PLAN_FROZEN", + "plan_frozen_producer": plan_producer, + "running_producer": running_sha, + } + + from benchmarks.inferswarm_76.capture import RowPruningSink, arm_full_capture + from benchmarks.inferswarm_r6.stage_runtime import GemmaDenseStage + + block = plan["blocks"][-1] + runtime = GemmaDenseStage( + role="last", + model_path=model_path, + adapter_data={ + **block, + "runtime_capacity_tokens": int( + plan.get("runtime_capacity_tokens", 256) + ), + "declared_shared_state": plan.get("declared_shared_state"), + }, + ) + + sink: RowPruningSink | None = None + case_dir_root = Path(out_dir) + case_dir_root.mkdir(parents=True, exist_ok=True) + current_case_dir: Path | None = None + current_case_id: str | None = None + + def _arm(case_id: str) -> None: + nonlocal sink, current_case_dir, current_case_id + current_case_id = case_id + current_case_dir = case_dir_root / case_id + current_case_dir.mkdir(parents=True, exist_ok=True) + sink = RowPruningSink(role="last", gpu_uuid=gpu_uuid) + runtime._capture_sink = sink + runtime._capture_after_layers = frozenset() + # stage owns globals [32,48): after_layer_47 is emitted explicitly by + # the prefill path; the last stage owns no layer 0/15. + + # The last stage owns global layers [32,48); layer-0/15 wrappers do not + # apply. arm_full_capture no-ops for absent layers. + + buffer_bytes = MAX_TOKEN_COUNT * ROW_WIDTH * 2 + host_u8 = torch.empty(buffer_bytes, dtype=torch.uint8) + stats = {"boundaries_served": 0, "activation_bytes_rx": 0, + "result_bytes_tx": 0, "cases_saved": 0} + + server = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + server.bind((listen_host, listen_port)) + server.listen(1) + if ready_file: + Path(ready_file).write_text(json.dumps({ + "plan_digest": plan.get("digest"), + "stage": "last", + "gpu_uuid": gpu_uuid, + "listen": [listen_host, listen_port], + "pid": os.getpid(), + "producer_freetoken_sha": running_sha, + "producer_check": producer_check, + "runtime": runtime.report("P4_ready_for_resident_execution"), + })) + experiment_id = plan.get("digest") or "r6-dense-chain" + try: + conn, _addr = server.accept() + conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + session_id = None + while True: + identity = {"protocol": WIRE_PROTOCOL_ID, "experiment_id": experiment_id} + header, payload = recv_frame(conn, identity) + kind = header.get("kind") + if kind == "hello": + session_id = header["session_id"] + runtime.reset_session_state() + response = { + "kind": "response", "protocol": WIRE_PROTOCOL_ID, + "experiment_id": experiment_id, "session_id": session_id, + "op": "HELLO_ACK", "runtime_ready": True, + } + send_exact(conn, encode_frame(response)) + continue + if kind != "request" or session_id is None: + raise WireError(f"unexpected frame kind {kind!r} before session") + if header.get("session_id") != session_id: + raise WireError("session identity mismatch") + op = header["op"] + if op == "OPEN_SESSION": + runtime.reset_session_state() + response = { + "kind": "response", "protocol": WIRE_PROTOCOL_ID, + "experiment_id": experiment_id, "session_id": session_id, + "op": "SESSION_ACK", + } + send_exact(conn, encode_frame(response)) + continue + if op == "CASE_BEGIN": + _arm(header["case_id"]) + runtime.reset_session_state() + response = { + "kind": "response", "protocol": WIRE_PROTOCOL_ID, + "experiment_id": experiment_id, "session_id": session_id, + "op": "CASE_ACK", "case_id": current_case_id, + } + send_exact(conn, encode_frame(response)) + continue + if op == "CASE_SAVE": + if sink is None or current_case_dir is None: + raise WireError("CASE_SAVE before CASE_BEGIN") + manifest = sink.save(str(current_case_dir), header["tag"]) + stats["cases_saved"] += 1 + response = { + "kind": "response", "protocol": WIRE_PROTOCOL_ID, + "experiment_id": experiment_id, "session_id": session_id, + "op": "SAVE_ACK", "case_id": current_case_id, + "manifest": manifest, + } + send_exact(conn, encode_frame(response)) + continue + if op != "BOUNDARY": + raise WireError(f"unsupported request op {op!r}") + token_count = int(header["token_count"]) + validate_request( + header, + contract=BOUNDARY_CONTRACT, + checksum=payload_checksum(payload) if diagnostic else None, + payload=payload, + ) + hidden = ( + torch.frombuffer(bytearray(payload), dtype=torch.uint8) + .view(torch.bfloat16) + .reshape(token_count, ROW_WIDTH) + .to(device="cuda:0", non_blocking=False) + ) + if header.get("capture_step") is not None: + runtime._capture_step = int(header["capture_step"]) + if header["operation"] == "prefill": + token, logits = runtime.prefill(None, hidden, int(header["position"])) + else: + if token_count != 1: + raise WireError("decode boundary requires exactly one token") + token, logits = runtime.decode(hidden, int(header["position"])) + runtime._capture_step = None + row = logits + nan_inf = int(torch.isnan(row).sum().item() + + torch.isinf(row).sum().item()) + top2 = torch.topk(row, 2) + stats["boundaries_served"] += 1 + stats["activation_bytes_rx"] += len(payload) + response = { + "kind": "response", "protocol": WIRE_PROTOCOL_ID, + "experiment_id": experiment_id, "session_id": session_id, + "op": "TOKEN_RESULT", + "token_id": int(token), + "top1_index": int(top2.indices[0].item()), + "top1_value_hex": float(top2.values[0].item()).hex(), + "top2_index": int(top2.indices[1].item()), + "top2_value_hex": float(top2.values[1].item()).hex(), + "margin_hex": float( + top2.values[0].item() - top2.values[1].item()).hex(), + "nan_inf_count": nan_inf, + "consumer_sha256": payload_checksum(payload), + "compute_ns": 0, + } + frame = encode_frame(response) + send_exact(conn, frame) + stats["result_bytes_tx"] += len(frame) + del hidden + finally: + try: + Path(os.environ.get( + "I76_LAST_STAGE_FINAL_REPORT", + "/tmp/i76-last-stage.json")).write_text(json.dumps({ + "schema": "inferswarm.issue76.last-stage-final-report/1", + "plan_digest": plan.get("digest"), + "producer_freetoken_sha": running_sha, + "producer_check": producer_check, + "stats": stats, + "runtime": runtime.report("P5_post_run"), + })) + except Exception: # noqa: BLE001 + pass + server.close() + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--listen-host", default="0.0.0.0") + parser.add_argument("--listen-port", type=int, default=18485) + parser.add_argument("--participant-plan", required=True) + parser.add_argument("--model", required=True) + parser.add_argument("--gpu-uuid", required=True) + parser.add_argument("--diagnostic", action="store_true") + parser.add_argument("--ready-file") + parser.add_argument("--allow-producer", default=None) + parser.add_argument("--out-dir", required=True, + help="root dir for per-case capture bundles") + args = parser.parse_args(argv) + serve( + listen_host=args.listen_host, + listen_port=args.listen_port, + participant_plan=args.participant_plan, + model_path=args.model, + gpu_uuid=args.gpu_uuid, + diagnostic=args.diagnostic, + ready_file=args.ready_file, + allow_producer=args.allow_producer, + out_dir=args.out_dir, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/inferswarm_76/reducer.py b/benchmarks/inferswarm_76/reducer.py new file mode 100644 index 000000000..6d2416a98 --- /dev/null +++ b/benchmarks/inferswarm_76/reducer.py @@ -0,0 +1,150 @@ +"""#76 case reducer: 15 frozen envelopes from paired reference/candidate runs. + +Pure host-float64 reduction exactly per frozen REDUCER.md +(inferswarm@f394dc9): full-domain absolute differences, fsum-of-squares RMS, +nearest-rank/higher p99, per-case per-family maximum across all declared +checkpoints and replay positions. + +Inputs are the per-case capture bundles (.pt, RowPruningSink format) from the +single arm (reference) and the chain arm (candidate). The chain bundles live +across stage subdirs; this module resolves the frozen checkpoint IDs from the +union of stage records using the stage-role mappings. +""" + +from __future__ import annotations + +import json +import math +from pathlib import Path +from typing import Any + +from benchmarks.inferswarm_76 import ( + CAPTURE_POSITIONS, + CHECKPOINT_FAMILY_MAP, + ENVELOPE_CHECKPOINT_IDS, + SEAM_TO_CHECKPOINT, + conservative_case_family, + envelopes_from_case_metrics, + tensor_metrics, +) + +# Chain-stage seam -> frozen checkpoint id (envelope-relevant only). +CHAIN_SEAM_MAP = { + "first": { + "embedding_output": "embedding-output", + "layer0_o_proj_input": "layer-0-o-proj-input", + "layer0_o_proj_output": "layer-0-o-proj-output", + "layer15_attn_o_proj_output": "global-layer-15-attention-o-proj-output", + "after_layer_15": "post-global-layer-15-residual", + }, + "middle": { + "after_layer_31": "post-global-layer-31-residual", + }, + "last": { + "after_layer_47": "post-global-layer-47-residual", + "final_norm": "final-normalized-hidden-state", + "full_final_row_bf16_logits": "full-final-row-bf16-logits", + "final_row_fp32": "full-final-row-fp32-consumer-logits", + }, +} +SINGLE_SEAM_MAP = SEAM_TO_CHECKPOINT["single"] + + +def _load_bundle(path: Path) -> list[dict[str, Any]]: + import torch + + bundle = torch.load(path, map_location="cpu", weights_only=False) + return [ + {"meta": meta, "tensor": tensor} + for meta, tensor in zip(bundle["records"], bundle["tensors"]) + ] + + +def _grouped( + records: list[dict[str, Any]], seam_map: dict[str, str] +) -> dict[tuple[int, str], Any]: + out: dict[tuple[int, str], Any] = {} + for record in records: + meta = record["meta"] + name = meta.get("checkpoint") + if name not in seam_map: + continue # boundary probes: exact-layer evidence, not envelopes + checkpoint_id = seam_map[name] + position = int(meta["step"]) + key = (position, checkpoint_id) + if key in out: + raise ValueError(f"duplicate capture {key}") + out[key] = record + return out + + +def _require_dtype(record: dict[str, Any], checkpoint_id: str) -> None: + expected = CHECKPOINT_FAMILY_MAP[checkpoint_id][1] + observed = record["meta"]["dtype"] + if observed != expected: + raise ValueError( + f"{checkpoint_id}: semantic dtype {observed} != frozen {expected}" + ) + + +def reduce_case( + *, + case_id: str, + reference_bundle: Path, + chain_bundles: dict[str, Path], +) -> dict[str, Any]: + """Compare one case; returns envelope hex strings + integrity detail.""" + ref_records = _load_bundle(reference_bundle) + ref = _grouped(ref_records, SINGLE_SEAM_MAP) + + cand: dict[tuple[int, str], Any] = {} + stage_of: dict[tuple[int, str], str] = {} + for stage_role, bundle_path in chain_bundles.items(): + records = _load_bundle(bundle_path) + grouped = _grouped(records, CHAIN_SEAM_MAP[stage_role]) + for key, record in grouped.items(): + if key in cand: + raise ValueError(f"duplicate chain capture {key} ({stage_role})") + cand[key] = record + stage_of[key] = stage_role + + per_checkpoint: dict[str, dict[str, float]] = {} + integrity = {"missing_reference": [], "missing_candidate": [], + "dtype_mismatches": [], "nan_inf": []} + for checkpoint_id in sorted(ENVELOPE_CHECKPOINT_IDS): + for position in CAPTURE_POSITIONS: + key = (position, checkpoint_id) + if key not in ref: + integrity["missing_reference"].append(list(key)) + continue + if key not in cand: + integrity["missing_candidate"].append(list(key)) + continue + _require_dtype(ref[key], checkpoint_id) + _require_dtype(cand[key], checkpoint_id) + for side, record in (("ref", ref[key]), ("cand", cand[key])): + meta = record["meta"] + if meta["nan_count"] or meta["inf_count"]: + integrity["nan_inf"].append( + {"side": side, "key": list(key), + "nan": meta["nan_count"], "inf": meta["inf_count"]}) + r = ref[key]["tensor"] + metrics = tensor_metrics( + [float(v) for v in r.flatten().tolist()], + [float(v) for v in cand[key]["tensor"].flatten().tolist()], + ) + per_checkpoint[f"{checkpoint_id}@{position}"] = metrics + if any(integrity[k] for k in + ("missing_reference", "missing_candidate", "dtype_mismatches")): + raise ValueError(f"{case_id}: incomplete capture set: {integrity}") + if integrity["nan_inf"]: + raise ValueError(f"{case_id}: NaN/Inf at declared finite checkpoint: " + f"{integrity['nan_inf'][:3]}") + + envelopes = envelopes_from_case_metrics(per_checkpoint) + return { + "case_id": case_id, + "envelopes": envelopes, + "integrity": integrity, + "checkpoint_count": len(per_checkpoint), + } diff --git a/benchmarks/inferswarm_76/reference_runner.py b/benchmarks/inferswarm_76/reference_runner.py new file mode 100644 index 000000000..6bfba3053 --- /dev/null +++ b/benchmarks/inferswarm_76/reference_runner.py @@ -0,0 +1,225 @@ +"""#76 single-arm case runner: matched FreeToken reference on the RTX 3090. + +Case-driven replay-prefill greedy generation with the full 15-envelope +checkpoint capture at positions 0/1/3/7 and per-step top-1 margin +diagnostics. One process may run MANY cases sequentially (model load is the +dominant cost); Phase-A fresh-process realizations are created by launching +this runner again. + +Per case it writes: + //capture-.pt raw host tensors (#71 bundle format) + //case-.json case summary (tokens, margins, + checkpoint hashes, NaN/Inf counts, + producer, gpu identity) + +The case summary is the atomic unit of committed evidence; the .pt bundle is +retained immutably next to it. Nothing is overwritten: a repeat run must use +a different --tag. +""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +import time +from pathlib import Path + +from benchmarks.inferswarm_76 import ( + CAPTURE_POSITIONS, + GENERATED_TOKENS, + RUNTIME_CAPACITY_TOKENS, + load_corpus, + verify_case_identity, +) + + +def _producer(repo: Path) -> dict: + sha = subprocess.check_output( + ["git", "-c", f"safe.directory={repo}", "-C", str(repo), + "rev-parse", "HEAD"], text=True).strip() + status = subprocess.check_output( + ["git", "-c", f"safe.directory={repo}", "-C", str(repo), + "status", "--porcelain"], text=True) + return {"commit": sha, "dirty": bool(status)} + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", required=True) + parser.add_argument("--corpus", required=True, + help="JSON file: corpus manifest OR a list of case rows") + parser.add_argument("--case-ids", default=None, + help="comma-separated subset of case ids to run") + parser.add_argument("--out-dir", required=True) + parser.add_argument("--tag", required=True) + parser.add_argument("--gpu", default="0") + parser.add_argument("--attempt-id", required=True) + args = parser.parse_args(argv) + + import os + + os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu + repo = Path(__file__).resolve().parents[2] + + producer = _producer(repo) + if producer["dirty"]: + print(json.dumps({"status": "BLOCKED_DIRTY_SOURCE"})) + return 2 + + raw = json.loads(Path(args.corpus).read_text()) + if isinstance(raw, list): + cases = [verify_case_identity(row) for row in raw] + else: + cases = load_corpus(args.corpus) + if args.case_ids: + wanted = set(args.case_ids.split(",")) + cases = [c for c in cases if c["case_id"] in wanted] + missing = wanted - {c["case_id"] for c in cases} + if missing: + raise SystemExit(f"unknown case ids: {sorted(missing)}") + + import torch + + from benchmarks.inferswarm_76.capture import RowPruningSink, arm_full_capture + from benchmarks.inferswarm_r6.stage_runtime import GemmaDenseStage + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.layers.rotary import set_rope_device + from freetoken.research.r6_dense_census import ( + DenseBlockSpec, + checkpoint_census, + freeze_dense_block_plan, + ) + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + set_rope_device(torch.device("cuda:0")) + + gpu_uuid = subprocess.check_output( + ["nvidia-smi", "-i", "0", "--query-gpu=uuid", "--format=csv,noheader"], + text=True).strip() + + model = Path(args.model).resolve() + census = checkpoint_census(model, text_prefix="model.language_model") + shared = { + "id": "tied-embedding-lm-head", + "kind": "tied-weight-shared-state", + "tensor_keys": ["model.language_model.embed_tokens.weight"], + "bytes": census["bytes_by_owner_category"]["embedding/input"], + "materialization_policy": "single-cuda-tensor-used-for-input-and-output", + } + plan = freeze_dense_block_plan( + census, [DenseBlockSpec(0, 48, True, True)], declared_shared_state=shared + ) + runtime = GemmaDenseStage( + role="single", + model_path=str(model), + adapter_data={ + **plan["blocks"][0], + "declared_shared_state": shared, + "runtime_capacity_tokens": RUNTIME_CAPACITY_TOKENS, + }, + ) + + runtime._capture_sink = RowPruningSink(role="single", gpu_uuid=gpu_uuid) + runtime._capture_after_layers = frozenset({15, 31}) + # wrappers bind runtime._emit dynamically (which reads _capture_sink), + # so they are installed EXACTLY ONCE; per-case isolation comes from + # swapping the sink below. + arm_full_capture(runtime, runtime._capture_sink) + + results = [] + out_root = Path(args.out_dir) + out_root.mkdir(parents=True, exist_ok=True) + + for case in cases: + case_dir = out_root / case["case_id"] + case_dir.mkdir(parents=True, exist_ok=True) + runtime._capture_sink = RowPruningSink(role="single", gpu_uuid=gpu_uuid) + sink = runtime._capture_sink + + prompt = list(case["token_ids"]) + generated: list[int] = [] + margins: list[dict] = [] + nan_inf_total = 0 + t0 = time.perf_counter() + for step in range(GENERATED_TOKENS): + runtime.reset_session_state() + replay = prompt + generated + capture_now = step in CAPTURE_POSITIONS + runtime._capture_step = step if capture_now else None + token, logits = runtime.prefill(replay, None, 0) + row = logits # final-row FP32 consumer logits [vocab] + nan_inf_total += int( + torch.isnan(row).sum().item() + torch.isinf(row).sum().item()) + top2 = torch.topk(row, 2) + margins.append({ + "step": step, + "top1_index": int(top2.indices[0].item()), + "top1_value_hex": float(top2.values[0].item()).hex(), + "top2_index": int(top2.indices[1].item()), + "top2_value_hex": float(top2.values[1].item()).hex(), + "margin_hex": float( + top2.values[0].item() - top2.values[1].item()).hex(), + }) + generated.append(int(token)) + del logits, row + + manifest = sink.save(str(case_dir), args.tag) + summary = { + "schema": "inferswarm.issue76.single-case-run/1", + "attempt_id": args.attempt_id, + "case_id": case["case_id"], + "case_sha256": case["case_sha256"], + "prompt_sha256": case["prompt_sha256"], + "token_ids_sha256": case["token_ids_sha256"], + "generated_token_ids": generated, + "step_margins": margins, + "min_top1_margin_hex": min( + float.fromhex(m["margin_hex"]) for m in margins).hex(), + "nan_inf_count": nan_inf_total, + "producer": producer, + "gpu_uuid": gpu_uuid, + "role": "single", + "capture_positions": list(CAPTURE_POSITIONS), + "capture_manifest": manifest, + "wall_seconds": time.perf_counter() - t0, + } + path = case_dir / f"case-{args.tag}.json" + if path.exists(): + raise SystemExit(f"refusing to overwrite {path}") + path.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n") + results.append(summary) + print(json.dumps({ + "case_id": case["case_id"], "status": "CASE_COMMITTED", + "tokens": generated, + "nan_inf": nan_inf_total, + "records": manifest["record_count"]}), flush=True) + # sink replaced per case at loop head + + index = { + "schema": "inferswarm.issue76.single-run-index/1", + "attempt_id": args.attempt_id, + "producer": producer, + "gpu_uuid": gpu_uuid, + "tag": args.tag, + "case_count": len(results), + "cases": [ + {"case_id": r["case_id"], "case_sha256": r["case_sha256"], + "generated_token_ids": r["generated_token_ids"], + "nan_inf_count": r["nan_inf_count"], + "min_top1_margin_hex": r["min_top1_margin_hex"]} + for r in results + ], + } + index_path = out_root / f"index-{args.tag}.json" + if index_path.exists(): + raise SystemExit(f"refusing to overwrite {index_path}") + index_path.write_text(json.dumps(index, indent=2, sort_keys=True) + "\n") + print(json.dumps({"status": "RUN_COMPLETE", "cases": len(results)})) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/inferswarm_76/stage_entry.py b/benchmarks/inferswarm_76/stage_entry.py new file mode 100644 index 000000000..d121d114b --- /dev/null +++ b/benchmarks/inferswarm_76/stage_entry.py @@ -0,0 +1,190 @@ +"""#76 local stage entry: stages 1-2 of the chain with per-case capture. + +Reuses the accepted R6 stage-process protocol (benchmarks.inferswarm_r6. +stage_chain._stage_entry shape) with these evidence-side differences: + +- ``CASE_ARM {out_dir, tag, gpu_uuid, after_layers}`` swaps in a fresh + ``RowPruningSink`` for the next case (the R6 ARM_CAPTURE armed one sink + for the whole connection); +- ``SAVE_CAPTURE {suffix}`` persists the current sink (same as R6); +- the #76 capture wrappers (o_proj checkpoints) are installed once at + startup via ``arm_full_capture``; they read the CURRENT sink through + ``runtime._capture_sink`` so per-case swapping just works. +""" + +from __future__ import annotations + +import os + + +def _stage_entry(*, role, adapter_data, model_path, connection): + import traceback + + try: + import torch + + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.layers.rotary import set_rope_device + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + set_rope_device(torch.device("cuda:0")) + + from benchmarks.inferswarm_76.capture import ( + RowPruningSink, + arm_full_capture, + ) + from benchmarks.inferswarm_r6.stage_runtime import GemmaDenseStage + + runtime = GemmaDenseStage( + role=role, model_path=model_path, adapter_data=dict(adapter_data) + ) + sink: RowPruningSink | None = None + capture_out_dir = None + capture_tag = None + + def _arm(message) -> None: + nonlocal sink, capture_out_dir, capture_tag + capture_out_dir = message["out_dir"] + capture_tag = message["tag"] + sink = RowPruningSink( + role=role, gpu_uuid=message.get("gpu_uuid") + ) + runtime._capture_sink = sink + runtime._capture_after_layers = frozenset( + int(x) for x in message.get("after_layers", []) + ) + + # wrappers read runtime._capture_sink dynamically; arming with a + # placeholder sink here only installs the bound methods once. + runtime._capture_sink = RowPruningSink(role=role) + arm_full_capture(runtime, runtime._capture_sink) + runtime._capture_sink = None + + connection.send({ + "op": "READY", + "role": role, + "runtime_report": runtime.report("P4_ready_for_resident_execution"), + }) + while True: + message = connection.recv() + op = message["op"] + if op == "CASE_ARM": + _arm(message) + connection.send({"op": "ACK"}) + elif op == "SAVE_CAPTURE": + if sink is None or capture_out_dir is None: + raise RuntimeError("SAVE_CAPTURE without CASE_ARM") + manifest = sink.save( + capture_out_dir, f"{capture_tag}-{message['suffix']}" + ) + connection.send({"op": "ACK", "manifest": manifest}) + elif op == "PREFILL": + if role != "first" and message.get("hidden") is not None: + message["hidden"] = message["hidden"].to( + device="cuda:0", dtype=torch.bfloat16 + ) + if message.get("capture_step") is not None: + runtime._capture_step = int(message["capture_step"]) + if role == "first": + hidden, _ = runtime.prefill( + message["token_ids"], None, message["position"] + ) + if message.get("capture_step") is not None: + runtime._capture_step = None + connection.send({"op": "BOUNDARY_PAYLOAD", + "hidden": hidden.cpu()}) + else: + out = runtime.prefill( + None, message["hidden"], message["position"] + ) + if message.get("capture_step") is not None: + runtime._capture_step = None + connection.send({"op": "BOUNDARY_PAYLOAD", + "hidden": out[0].cpu()}) + elif op == "DECODE": + if role != "first" and message.get("hidden") is not None: + message["hidden"] = message["hidden"].to( + device="cuda:0", dtype=torch.bfloat16 + ) + if role == "first": + hidden, _ = runtime.decode( + message["token_id"], message["position"] + ) + connection.send({"op": "BOUNDARY_PAYLOAD", + "hidden": hidden.cpu()}) + else: + out = runtime.decode(message["hidden"], message["position"]) + connection.send({"op": "BOUNDARY_PAYLOAD", + "hidden": out[0].cpu()}) + elif op == "REPORT": + connection.send({"op": "REPORT", "report": runtime.report()}) + elif op == "RESET": + runtime.reset_session_state() + connection.send({"op": "ACK"}) + elif op == "SHUTDOWN": + connection.send({"op": "ACK", "report": runtime.report()}) + return + else: + raise RuntimeError(f"unknown stage op {op!r}") + except BaseException as exc: + try: + connection.send({ + "op": "ERROR", + "type": type(exc).__name__, + "message": str(exc), + "traceback": traceback.format_exc(), + }) + except (BrokenPipeError, EOFError, OSError): + pass + + +class I76StageClient: + """Control pipe to one spawn-isolated #76 stage process.""" + + def __init__(self, context, *, role, adapter_data, model_path, gpu_index: int): + parent, child = context.Pipe() + env = {**os.environ, "CUDA_VISIBLE_DEVICES": str(gpu_index)} + self.parent = parent + self.role = role + self.process = context.Process( + target=_stage_entry_env, + args=(env,), + kwargs={ + "role": role, + "adapter_data": adapter_data, + "model_path": model_path, + "connection": child, + }, + ) + self.process.start() + + def recv(self): + return self.parent.recv() + + def send(self, message): + self.parent.send(message) + + def request(self, message): + self.send(message) + response = self.recv() + if isinstance(response, dict) and response.get("op") == "ERROR": + raise RuntimeError(f"stage {self.role} error: {response}") + return response + + def shutdown(self): + try: + self.send({"op": "SHUTDOWN"}) + self.parent.recv() + except (BrokenPipeError, EOFError, OSError): + pass + self.process.join(timeout=30) + if self.process.is_alive(): + self.process.terminate() + + +def _stage_entry_env(env: dict, *, role, adapter_data, model_path, connection): + os.environ.clear() + os.environ.update(env) + _stage_entry(role=role, adapter_data=adapter_data, model_path=model_path, + connection=connection) diff --git a/tests/research/test_inferswarm_76_harness.py b/tests/research/test_inferswarm_76_harness.py new file mode 100644 index 000000000..f87d8b965 --- /dev/null +++ b/tests/research/test_inferswarm_76_harness.py @@ -0,0 +1,196 @@ +"""CPU-only unit/source-contract tests for the #76 execution harness. + +These run on any host (no torch import at collection): pure-reducer math, +case identity verification, envelope construction, and source-contract +checks on the frozen checkpoint map. Torch-dependent pieces are exercised +by node-side qualification. +""" + +from __future__ import annotations + +import json +import math +import sys +import unittest +from pathlib import Path + +REPO = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO / "benchmarks")) + +from inferswarm_76 import ( # noqa: E402 + CAPTURE_POSITIONS, + CHECKPOINT_FAMILY_MAP, + ENVELOPES, + ENVELOPE_CHECKPOINT_IDS, + canonical_json_bytes, + conservative_case_family, + envelopes_from_case_metrics, + load_corpus, + nearest_rank_higher, + sha256_bytes, + tensor_metrics, + verify_case_identity, +) + + +class FrozenContractTests(unittest.TestCase): + def test_fifteen_envelopes(self): + self.assertEqual(len(ENVELOPES), 15) + self.assertEqual(len(ENVELOPE_CHECKPOINT_IDS), 10) + + def test_checkpoint_families_match_frozen_map(self): + # Mirrors inferswarm@f394dc9 manifests/checkpoint-family-map.json + self.assertEqual( + CHECKPOINT_FAMILY_MAP["embedding-output"], + ("local-bf16-backend-operation-output", "bfloat16"), + ) + self.assertEqual( + CHECKPOINT_FAMILY_MAP["full-final-row-fp32-consumer-logits"], + ("fp32-consumer-logits", "float32"), + ) + families = {f for f, _ in CHECKPOINT_FAMILY_MAP.values()} + self.assertEqual(len(families), 5) + + def test_capture_positions_frozen(self): + self.assertEqual(CAPTURE_POSITIONS, (0, 1, 3, 7)) + + +class ReducerMathTests(unittest.TestCase): + def test_metrics_known_values(self): + m = tensor_metrics([1.0, 2.0, 3.0, 4.0], [1.5, 2.0, 2.5, 4.0]) + self.assertEqual(m["max-absolute-difference"], 0.5) + expect_rms = math.sqrt(math.fsum([0.25, 0.25]) / 4) + self.assertEqual(m["rms-difference"], expect_rms) + # errors [0.5, 0, 0.5, 0] sorted [0, 0, 0.5, 0.5]; ceil(.99*4)=4 -> 0.5 + self.assertEqual(m["p99-absolute-error"], 0.5) + + def test_p99_nearest_rank_higher(self): + # N=100: ceil(0.99*100)=99 -> 99th smallest (one-based) = second- + # largest element of [1.0 x99, 2.0] is 1.0; the max alone is rank 100. + self.assertEqual(nearest_rank_higher([1.0] * 99 + [2.0]), 1.0) + self.assertEqual(nearest_rank_higher([2.0] + [1.0] * 99), 1.0) + # N=100 with two distinct high values: rank 99 hits the 2nd-largest + self.assertEqual(nearest_rank_higher([0.0] * 98 + [1.0, 2.0]), 1.0) + self.assertEqual(nearest_rank_higher([0.0] * 100), 0.0) + with self.assertRaises(ValueError): + nearest_rank_higher([]) + with self.assertRaises(ValueError): + nearest_rank_higher([1.0, -0.1]) + + def test_case_family_max_per_metric(self): + out = conservative_case_family([ + {"max-absolute-difference": 1.0, "rms-difference": 3.0, + "p99-absolute-error": 2.0}, + {"max-absolute-difference": 4.0, "rms-difference": 1.0, + "p99-absolute-error": 0.5}, + ]) + self.assertEqual(out["max-absolute-difference"], 4.0) + self.assertEqual(out["rms-difference"], 3.0) + self.assertEqual(out["p99-absolute-error"], 2.0) + + def test_envelopes_hex_serialization(self): + rows = [ + {"max-absolute-difference": v, "rms-difference": v / 2, + "p99-absolute-error": v / 4} + for v in (0.5, 1.25, 0.125) + ] + out = envelopes_from_case_metrics({ + "embedding-output": rows[0], + "layer-0-o-proj-input": rows[1], + "layer-0-o-proj-output": rows[2], + "global-layer-15-attention-o-proj-output": rows[0], + "post-global-layer-15-residual": rows[1], + "post-global-layer-31-residual": rows[2], + "post-global-layer-47-residual": rows[0], + "final-normalized-hidden-state": rows[1], + "full-final-row-bf16-logits": rows[2], + "full-final-row-fp32-consumer-logits": rows[0], + }) + self.assertEqual(len(out), 15) + for value in out.values(): + float.fromhex(value) # exact round trip + self.assertEqual( + float.fromhex(out["local-bf16-backend-operation-output:" + "max-absolute-difference"]), + 1.25, + ) + + def test_missing_family_fails_closed(self): + with self.assertRaises(ValueError): + envelopes_from_case_metrics({ + "embedding-output": { + "max-absolute-difference": 1.0, + "rms-difference": 1.0, + "p99-absolute-error": 1.0, + } + }) + + def test_domain_size_mismatch_rejected(self): + with self.assertRaises(ValueError): + tensor_metrics([1.0, 2.0], [1.0]) + + def test_nonfinite_rejected(self): + with self.assertRaises(ValueError): + tensor_metrics([1.0], [float("nan")]) + + +class CaseIdentityTests(unittest.TestCase): + CASE = { + "case_id": "c74-01-01-01", + "case_sha256": "286556cf76004ad524933ead5d1768aae6d24aae6199db11500bf8315d8995fe", + "content_class": "ordinary-prose", + "length_regime": [4, 8], + "prompt_sha256": "00985f1fa382d3485e136c3ff120fe1b8f2a169a56b97423d596569e0b96f1bb", + "prompt_text": "wind bird stone wind", + "token_count": 4, + "token_ids": [15879, 8001, 10810, 6573], + "token_ids_sha256": "7255f8af076acf1952ac8fa6125abe5370fe0529c3eea9f28e51749c9fba5cc6", + } + + def test_valid_case_verifies(self): + out = verify_case_identity(dict(self.CASE)) + self.assertEqual(out["case_id"], "c74-01-01-01") + + def test_mutated_prompt_fails(self): + bad = dict(self.CASE) + bad["prompt_text"] = "wind bird stone windx" + with self.assertRaises(ValueError): + verify_case_identity(bad) + + def test_mutated_tokens_fail(self): + bad = dict(self.CASE) + bad["token_ids"] = [15879, 8001, 10810, 6574] + with self.assertRaises(ValueError): + verify_case_identity(bad) + + def test_mutated_identity_fails(self): + bad = dict(self.CASE) + bad["content_class"] = "multilingual-text" + with self.assertRaises(ValueError): + verify_case_identity(bad) + + +class SourceContractTests(unittest.TestCase): + """Structural contracts over the harness source (torch-free hosts).""" + + def test_margin_definition_documented_and_frozen(self): + source = (REPO / "benchmarks/inferswarm_76/__init__.py").read_text() + self.assertIn("matched-reference-top1-margin", source) + + def test_no_threshold_constants_in_harness(self): + # The harness must not bake any numeric acceptance threshold; limits + # arrive only from the derived frozen threshold manifest. + for name in ("reference_runner.py", "chain_runner.py", "reducer.py"): + source = (REPO / "benchmarks/inferswarm_76" / name).read_text() + for token in ("0.25", "FROZEN_THRESHOLD"): + self.assertNotIn( + token, source, + f"{name} must not contain threshold literal {token}") + + def test_row_pruning_preserves_final_row_domain(self): + source = (REPO / "benchmarks/inferswarm_76/capture.py").read_text() + self.assertIn("final_row_from_bf16_record", source) + + +if __name__ == "__main__": + unittest.main() From f07189f4dd4fd8748a9b06cdbd6acc951998e9a0 Mon Sep 17 00:00:00 2001 From: Eric Sutphen Date: Thu, 3 Sep 2026 21:02:26 -0400 Subject: [PATCH 2/6] issue #76: resolve identity-only subset manifests against the full corpus sentinel-subset.json carries identity hashes without prompt text; runners now cross-check subset rows against the frozen calibration corpus before execution. --- benchmarks/inferswarm_76/chain_runner.py | 10 ++--- benchmarks/inferswarm_76/reference_runner.py | 45 +++++++++++++++++--- 2 files changed, 43 insertions(+), 12 deletions(-) diff --git a/benchmarks/inferswarm_76/chain_runner.py b/benchmarks/inferswarm_76/chain_runner.py index 39aa8b626..e6fef2661 100644 --- a/benchmarks/inferswarm_76/chain_runner.py +++ b/benchmarks/inferswarm_76/chain_runner.py @@ -30,7 +30,6 @@ from benchmarks.inferswarm_76 import ( CAPTURE_POSITIONS, GENERATED_TOKENS, - load_corpus, verify_case_identity, ) @@ -50,6 +49,8 @@ def main(argv=None) -> int: parser.add_argument("--model", required=True) parser.add_argument("--plan", required=True) parser.add_argument("--corpus", required=True) + parser.add_argument("--resolve-corpus", default=None, + help="full calibration corpus for identity-only subsets") parser.add_argument("--case-ids", default=None) parser.add_argument("--last-stage-host", default="10.0.0.219") parser.add_argument("--last-stage-port", type=int, default=18485) @@ -66,11 +67,8 @@ def main(argv=None) -> int: print(json.dumps({"status": "BLOCKED_DIRTY_SOURCE"})) return 2 - raw = json.loads(Path(args.corpus).read_text()) - if isinstance(raw, list): - cases = [verify_case_identity(row) for row in raw] - else: - cases = load_corpus(args.corpus) + from benchmarks.inferswarm_76.reference_runner import resolve_cases + cases = resolve_cases(args.corpus, args.resolve_corpus) if args.case_ids: wanted = set(args.case_ids.split(",")) cases = [c for c in cases if c["case_id"] in wanted] diff --git a/benchmarks/inferswarm_76/reference_runner.py b/benchmarks/inferswarm_76/reference_runner.py index 6bfba3053..f999be094 100644 --- a/benchmarks/inferswarm_76/reference_runner.py +++ b/benchmarks/inferswarm_76/reference_runner.py @@ -30,11 +30,46 @@ CAPTURE_POSITIONS, GENERATED_TOKENS, RUNTIME_CAPACITY_TOKENS, - load_corpus, verify_case_identity, ) +def verify_subset_row(row: dict) -> dict: + """Identity check for a subset-manifest row (no text fields).""" + for field in ("case_id", "case_sha256", "prompt_sha256", + "token_ids_sha256"): + if field not in row: + raise SystemExit(f"subset row missing {field}") + return row + + +def resolve_cases(corpus_path: str, resolve_corpus: str | None) -> list[dict]: + """Load case rows from a corpus manifest, raw list, or identity-only + subset (resolved against the full corpus via hash cross-check).""" + raw = json.loads(Path(corpus_path).read_text()) + if isinstance(raw, list): + return [verify_case_identity(row) for row in raw] + corpus_cases = raw.get("cases") + if corpus_cases and "prompt_text" in corpus_cases[0]: + return [verify_case_identity(row) for row in corpus_cases] + # identity-only subset manifest (e.g. sentinel-subset.json): rows carry + # case_id/case_sha256/prompt_sha256/token_ids_sha256 but no text; + # resolve them against the full calibration corpus. + if resolve_corpus is None: + raise SystemExit("identity-only corpus requires --resolve-corpus") + full = json.loads(Path(resolve_corpus).read_text()) + by_id = {row["case_id"]: verify_case_identity(row) for row in full["cases"]} + cases = [] + for row in corpus_cases: + wanted = verify_subset_row(row) + resolved = by_id[row["case_id"]] + for field in ("case_sha256", "prompt_sha256", "token_ids_sha256"): + if wanted[field] != resolved[field]: + raise SystemExit(f"{row['case_id']}: subset {field} mismatch") + cases.append(resolved) + return cases + + def _producer(repo: Path) -> dict: sha = subprocess.check_output( ["git", "-c", f"safe.directory={repo}", "-C", str(repo), @@ -50,6 +85,8 @@ def main(argv=None) -> int: parser.add_argument("--model", required=True) parser.add_argument("--corpus", required=True, help="JSON file: corpus manifest OR a list of case rows") + parser.add_argument("--resolve-corpus", default=None, + help="full calibration corpus for identity-only subsets") parser.add_argument("--case-ids", default=None, help="comma-separated subset of case ids to run") parser.add_argument("--out-dir", required=True) @@ -68,11 +105,7 @@ def main(argv=None) -> int: print(json.dumps({"status": "BLOCKED_DIRTY_SOURCE"})) return 2 - raw = json.loads(Path(args.corpus).read_text()) - if isinstance(raw, list): - cases = [verify_case_identity(row) for row in raw] - else: - cases = load_corpus(args.corpus) + cases = resolve_cases(args.corpus, args.resolve_corpus) if args.case_ids: wanted = set(args.case_ids.split(",")) cases = [c for c in cases if c["case_id"] in wanted] From cc7b520f3187d973025b50d57a5bd05a7402a894 Mon Sep 17 00:00:00 2001 From: Eric Sutphen Date: Thu, 3 Sep 2026 21:06:48 -0400 Subject: [PATCH 3/6] issue #76: case-scoped wire client for the last-stage service --- benchmarks/inferswarm_76/chain_runner.py | 4 +- benchmarks/inferswarm_76/wire_client.py | 162 +++++++++++++++++++++++ 2 files changed, 164 insertions(+), 2 deletions(-) create mode 100644 benchmarks/inferswarm_76/wire_client.py diff --git a/benchmarks/inferswarm_76/chain_runner.py b/benchmarks/inferswarm_76/chain_runner.py index e6fef2661..552045ff7 100644 --- a/benchmarks/inferswarm_76/chain_runner.py +++ b/benchmarks/inferswarm_76/chain_runner.py @@ -77,7 +77,7 @@ def main(argv=None) -> int: raise SystemExit(f"unknown case ids: {sorted(missing)}") from benchmarks.inferswarm_76.stage_entry import I76StageClient - from benchmarks.inferswarm_r6.wire_client import RemoteLastStageClient + from benchmarks.inferswarm_76.wire_client import I76LastStageClient plan = json.loads(Path(args.plan).read_text()) shared = plan.get("declared_shared_state") @@ -100,7 +100,7 @@ def main(argv=None) -> int: ) ) stages.append( - RemoteLastStageClient( + I76LastStageClient( host=args.last_stage_host, port=args.last_stage_port, experiment_id=plan["digest"], diff --git a/benchmarks/inferswarm_76/wire_client.py b/benchmarks/inferswarm_76/wire_client.py new file mode 100644 index 000000000..b42bbe1f4 --- /dev/null +++ b/benchmarks/inferswarm_76/wire_client.py @@ -0,0 +1,162 @@ +"""#76 wire client: R6 RemoteLastStageClient + case-scoped ops.""" + +from __future__ import annotations + +from typing import Any + +from freetoken.research.r4_wire import ( + WIRE_PROTOCOL_ID, + encode_frame, + recv_frame, + send_exact, +) + + +class I76LastStageClient: + role = "last" + + def __init__( + self, + *, + host: str, + port: int, + experiment_id: str, + session_id: int = 1, + connect_timeout: float = 30.0, + ) -> None: + import socket + + self._host = host + self._port = int(port) + self._identity = { + "protocol": WIRE_PROTOCOL_ID, + "experiment_id": experiment_id, + } + self._session_id = int(session_id) + self._sock = socket.create_connection( + (host, int(port)), timeout=connect_timeout + ) + self._sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + send_exact( + self._sock, + encode_frame({ + "kind": "hello", + "protocol": WIRE_PROTOCOL_ID, + "experiment_id": experiment_id, + "session_id": self._session_id, + }), + ) + header, _payload = recv_frame(self._sock, self._identity) + if header.get("op") != "HELLO_ACK": + raise RuntimeError(f"last-stage hello rejected: {header}") + + def send(self, message): + pass # StageClient API compatibility + + def recv(self): + raise RuntimeError("I76LastStageClient is request/response only") + + def _control(self, op: str, **fields: Any) -> dict: + header = { + "kind": "request", + "protocol": WIRE_PROTOCOL_ID, + "experiment_id": self._identity["experiment_id"], + "session_id": self._session_id, + "op": op, + **fields, + } + send_exact(self._sock, encode_frame(header)) + response, _ = recv_frame(self._sock, self._identity) + return response + + def _boundary( + self, *, operation: str, position: int, hidden, capture_step: int | None = None + ) -> dict: + import torch + + from freetoken.research.r4_wire import payload_checksum + + token_count = int(hidden.shape[0]) + payload = ( + hidden.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes() + ) + header = { + "kind": "request", + "protocol": WIRE_PROTOCOL_ID, + "experiment_id": self._identity["experiment_id"], + "session_id": self._session_id, + "op": "BOUNDARY", + "operation": operation, + "position": int(position), + "token_count": token_count, + "dtype": "bfloat16", + "layout": "plane-major-contiguous", + "payload_len": len(payload), + "payload_sha256": payload_checksum(payload), + } + if capture_step is not None: + header["capture_step"] = int(capture_step) + send_exact(self._sock, encode_frame(header, payload)) + response, _ = recv_frame(self._sock, self._identity) + if response.get("op") != "TOKEN_RESULT": + raise RuntimeError(f"last-stage rejected boundary: {response}") + return response + + def request(self, message): + op = message["op"] + if op == "PREFILL": + return self._boundary( + operation="prefill", + position=int(message["position"]), + hidden=message["hidden"], + capture_step=message.get("capture_step"), + ) + if op == "DECODE": + return self._boundary( + operation="decode", + position=int(message["position"]), + hidden=message["hidden"], + ) + if op == "CASE_BEGIN": + response = self._control("CASE_BEGIN", + case_id=message["case_id"]) + if response.get("op") != "CASE_ACK": + raise RuntimeError(f"CASE_BEGIN rejected: {response}") + return response + if op == "CASE_SAVE": + response = self._control("CASE_SAVE", tag=message["tag"]) + if response.get("op") != "SAVE_ACK": + raise RuntimeError(f"CASE_SAVE rejected: {response}") + return response + if op == "RESET": + send_exact( + self._sock, + encode_frame({ + "kind": "hello", + "protocol": WIRE_PROTOCOL_ID, + "experiment_id": self._identity["experiment_id"], + "session_id": self._session_id, + }), + ) + header, _ = recv_frame(self._sock, self._identity) + if header.get("op") != "HELLO_ACK": + raise RuntimeError(f"last-stage re-hello rejected: {header}") + return {"op": "ACK"} + if op == "REPORT": + return {"op": "REPORT", "report": {"remote": True}} + if op == "SHUTDOWN": + try: + self._sock.close() + except OSError: + pass + return {"op": "ACK"} + raise RuntimeError(f"unsupported wire op {op!r}") + + def shutdown(self): + try: + self.request({"op": "SHUTDOWN"}) + except Exception: # noqa: BLE001 + try: + self._sock.close() + except OSError: + pass From 3612dc8b4362bbedb233de2efbca067f03c6e3fe Mon Sep 17 00:00:00 2001 From: Eric Sutphen Date: Thu, 3 Sep 2026 21:07:07 -0400 Subject: [PATCH 4/6] issue #76: define wire protocol id locally (r4_wire does not export it) --- benchmarks/inferswarm_76/wire_client.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/benchmarks/inferswarm_76/wire_client.py b/benchmarks/inferswarm_76/wire_client.py index b42bbe1f4..5c6930153 100644 --- a/benchmarks/inferswarm_76/wire_client.py +++ b/benchmarks/inferswarm_76/wire_client.py @@ -5,12 +5,13 @@ from typing import Any from freetoken.research.r4_wire import ( - WIRE_PROTOCOL_ID, encode_frame, recv_frame, send_exact, ) +WIRE_PROTOCOL_ID = "inferswarm.r4.boundary-wire/1" + class I76LastStageClient: role = "last" From 615a665c4d0d0f8b25c8d0119719b2f64a612152 Mon Sep 17 00:00:00 2001 From: Eric Sutphen Date: Thu, 3 Sep 2026 21:10:13 -0400 Subject: [PATCH 5/6] issue #76: per-stage capture filenames (stage 2 overwrote stage 1's bundle) --- benchmarks/inferswarm_76/chain_runner.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/benchmarks/inferswarm_76/chain_runner.py b/benchmarks/inferswarm_76/chain_runner.py index 552045ff7..d498b1c2f 100644 --- a/benchmarks/inferswarm_76/chain_runner.py +++ b/benchmarks/inferswarm_76/chain_runner.py @@ -188,8 +188,9 @@ def main(argv=None) -> int: # persist per-case captures from stages 1-2 and the remote stage manifests = {} for stage_index, stage in enumerate(stages[:-1]): - ack = stage.request({"op": "SAVE_CAPTURE", - "suffix": args.tag}) + ack = stage.request( + {"op": "SAVE_CAPTURE", + "suffix": f"{args.tag}-stage{stage_index + 1}"}) manifests[f"stage{stage_index + 1}"] = ack["manifest"] save_ack = stages[-1].request( {"op": "CASE_SAVE", "tag": args.tag}) From 9f06d81c6e9519d7ddf68b302c1c2aa14993387d Mon Sep 17 00:00:00 2001 From: Eric Sutphen Date: Thu, 3 Sep 2026 21:12:17 -0400 Subject: [PATCH 6/6] issue #76: reducer accepts checkpoint@position keys --- benchmarks/inferswarm_76/__init__.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/benchmarks/inferswarm_76/__init__.py b/benchmarks/inferswarm_76/__init__.py index 16c4e0981..577f2bc73 100644 --- a/benchmarks/inferswarm_76/__init__.py +++ b/benchmarks/inferswarm_76/__init__.py @@ -234,9 +234,13 @@ def conservative_case_family( def envelopes_from_case_metrics( per_checkpoint: dict[str, dict[str, float]], ) -> dict[str, str]: - """15 frozen envelope strings (exact hex binary64) for one case.""" + """15 frozen envelope strings (exact hex binary64) for one case. + + Keys may be bare checkpoint ids or "@". + """ by_family: dict[str, list[dict[str, float]]] = {f: [] for f in FAMILIES} - for checkpoint_id, metrics in per_checkpoint.items(): + for key, metrics in per_checkpoint.items(): + checkpoint_id = key.split("@", 1)[0] family = CHECKPOINT_FAMILY_MAP[checkpoint_id][0] by_family[family].append(metrics) envelopes: dict[str, str] = {}