#!/usr/bin/env python3 """Locked WikiText distillation/evaluation for QH-024 and controls.""" from __future__ import annotations import argparse import hashlib import json import math import random import sys import time from datetime import datetime, timezone from pathlib import Path import numpy as np import torch from transformers import AutoModelForMultimodalLM PROJECT = Path(__file__).resolve().parents[2] LOCK = PROJECT / "artifacts/wikitext-qh024-distillation-lock.json" TRAIN_TOKENS = PROJECT / "datasets/processed/wikitext2_qwen38/train-tokens.pt" VALIDATION_TOKENS = PROJECT / "datasets/processed/wikitext2_qwen38/validation-tokens.pt" OUTPUT_DIR = PROJECT / "artifacts/wikitext_qh024" sys.path.insert(0, str(PROJECT / "code/src")) from quantum_qwen38.hyqut_replacement import QH024Config, freeze_and_replace_qh024 def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--candidate", choices=["base", "svd512", "qh024", "cc024"], required=True) parser.add_argument("--timing-context", choices=["exclusive", "shared"], default="shared") return parser.parse_args() def sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for chunk in iter(lambda: handle.read(8 << 20), b""): digest.update(chunk) return digest.hexdigest() def seed_everything(seed: int) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) def token_block(tokens: torch.Tensor, index: int, length: int, device: torch.device) -> torch.Tensor: start = index * length return tokens[start : start + length].to(device=device, dtype=torch.long).unsqueeze(0) def distill(model, replacement, tokens: torch.Tensor, lock: dict) -> dict: protocol = lock["training_protocol"] optimizer = torch.optim.AdamW( [parameter for parameter in replacement.parameters() if parameter.requires_grad], lr=float(protocol["learning_rate"]), weight_decay=float(protocol["weight_decay"]), ) losses, gradient_norms, seconds = [], [], [] model.eval() model.config.use_cache = False order = list(int(value) for value in lock["train_indices"]) random.Random(int(lock["seed"])).shuffle(order) for index in order: input_ids = token_block(tokens, index, int(lock["sequence_length"]), model.device) replacement.set_capture_teacher_only(True) with torch.inference_mode(): model(input_ids=input_ids, use_cache=False) hidden_states, teacher_target = replacement.take_capture() replacement.set_capture_teacher_only(False) optimizer.zero_grad(set_to_none=True) started = time.perf_counter() student = replacement.student_forward(hidden_states) denominator = teacher_target.float().square().mean().clamp_min(1.0e-8) loss = (student.float() - teacher_target.float()).square().mean() / denominator loss.backward() gradient_norm = torch.nn.utils.clip_grad_norm_( [parameter for parameter in replacement.parameters() if parameter.requires_grad], float(protocol["gradient_clip_norm"]), ) optimizer.step() torch.cuda.synchronize(replacement.gamma.device) losses.append(float(loss.detach())) gradient_norms.append(float(gradient_norm)) seconds.append(time.perf_counter() - started) del input_ids, hidden_states, teacher_target, student, denominator, loss return { "steps": len(losses), "objective": protocol["objective"], "learning_rate": float(protocol["learning_rate"]), "loss_first_16_mean": float(np.mean(losses[:16])), "loss_last_16_mean": float(np.mean(losses[-16:])), "loss_min": float(np.min(losses)), "gradient_norm_mean": float(np.mean(gradient_norms)), "mean_residual_step_seconds_excluding_teacher_forward": float(np.mean(seconds)), "total_residual_step_seconds_excluding_teacher_forward": float(np.sum(seconds)), "losses": losses, "frozen_parameter_grad_tensor_count": sum( parameter.grad is not None for parameter in model.parameters() if not parameter.requires_grad ), } @torch.inference_mode() def evaluate(model, tokens: torch.Tensor, indices: list[int], length: int) -> dict: model.eval() rows, seconds = [], [] for index in indices: input_ids = token_block(tokens, index, length, model.device) started = time.perf_counter() loss = model(input_ids=input_ids, labels=input_ids, use_cache=False).loss value = float(loss) elapsed = time.perf_counter() - started rows.append({ "block_index": index, "mean_token_nll": value, "predicted_tokens": length - 1, "seconds": elapsed, }) seconds.append(elapsed) nll = float(np.mean([row["mean_token_nll"] for row in rows])) return { "metrics": { "sequence_count": len(rows), "predicted_token_count": len(rows) * (length - 1), "mean_token_nll": nll, "token_perplexity": math.exp(min(50.0, nll)), "mean_sequence_seconds": float(np.mean(seconds)), "p95_sequence_seconds": float(np.quantile(seconds, 0.95)), }, "rows": rows, } def main() -> None: args = parse_args() if not torch.cuda.is_available(): raise RuntimeError("QH-024 formal runner is GPU-only") lock = json.loads(LOCK.read_text(encoding="utf-8")) if lock["status"] != "locked_before_qh024_qh_cc_base_performance_results": raise RuntimeError("invalid QH-024 lock") if lock["test_rows_inspected_or_used_for_fitness"] != 0: raise RuntimeError("test leakage disclosure is nonzero") seed_everything(int(lock["seed"])) train_tokens = torch.load(TRAIN_TOKENS, map_location="cpu", weights_only=True) validation_tokens = torch.load(VALIDATION_TOKENS, map_location="cpu", weights_only=True) model_dir = Path((PROJECT / "records/active_model_path.txt").read_text().strip()) model = AutoModelForMultimodalLM.from_pretrained( model_dir, local_files_only=True, dtype=torch.bfloat16, device_map="balanced", low_cpu_mem_usage=True, attn_implementation="sdpa", ) base_parameters = sum(parameter.numel() for parameter in model.parameters()) replacement = None if args.candidate == "base": for parameter in model.parameters(): parameter.requires_grad_(False) injection = { "candidate": "frozen-dense-base", "base_model_parameters": base_parameters, "deployed_model_parameters": base_parameters, "net_model_parameter_reduction": 0, "trainable_parameter_count": 0, } training = {"steps": 0} else: kind = "classical" if args.candidate == "cc024" else "quantum" injection = freeze_and_replace_qh024(model, QH024Config(), core_kind=kind) replacement = model.model.language_model.layers[7].self_attn.v_proj if args.candidate == "svd512": replacement.set_residual_enabled(False) training = {"steps": 0, "role": "untrained causal zero-residual control"} else: training = distill(model, replacement, train_tokens, lock) OUTPUT_DIR.mkdir(parents=True, exist_ok=True) checkpoint_path = OUTPUT_DIR / f"wikitext-{args.candidate}-s{lock['seed']}-replacement.pt" if replacement is not None and args.candidate in {"qh024", "cc024"}: torch.save( { name: parameter.detach().cpu() for name, parameter in replacement.named_parameters() if parameter.requires_grad }, checkpoint_path, ) if replacement is not None: replacement.strip_teacher() eval_indices = [int(value) for value in lock["eval_indices"]] evaluation = evaluate(model, validation_tokens, eval_indices, int(lock["sequence_length"])) evaluation_zero = None simulator_calls_before_zero = None simulator_calls_after_zero = None if replacement is not None and args.candidate in {"qh024", "cc024"}: simulator_calls_before_zero = replacement.core.circuit_calls replacement.set_residual_enabled(False) evaluation_zero = evaluate(model, validation_tokens, eval_indices, int(lock["sequence_length"])) simulator_calls_after_zero = replacement.core.circuit_calls replacement.set_residual_enabled(True) if simulator_calls_after_zero != simulator_calls_before_zero: raise AssertionError("zero branch unexpectedly executed the residual core") run_id = f"wikitext-{args.candidate}-s{lock['seed']}-n{lock['train_count']}-e{lock['eval_count']}" payload = { "status": "ok", "run_id": run_id, "candidate": args.candidate, "model_dir": str(model_dir), "timing_context": args.timing_context, "dataset_lock": str(LOCK), "dataset_lock_sha256": sha256(LOCK), "dataset": { "name": lock["dataset"], "source_revision": lock["source_revision"], "license": lock["license"], "sequence_length": lock["sequence_length"], "train_indices": lock["train_indices"] if training["steps"] else [], "eval_indices": eval_indices, "train_tokens_sha256": sha256(TRAIN_TOKENS), "validation_tokens_sha256": sha256(VALIDATION_TOKENS), "test_rows_inspected_or_used_for_fitness": 0, }, "injection": injection, "training": training, "evaluation_active": evaluation, "evaluation_zero_residual": evaluation_zero, "simulator_calls_before_zero_evaluation": simulator_calls_before_zero, "simulator_calls_after_zero_evaluation": simulator_calls_after_zero, "replacement_audit_after_teacher_strip": replacement.audit() if replacement else None, "checkpoint": str(checkpoint_path) if checkpoint_path.exists() else None, "gpu_peak_allocated_bytes": [ torch.cuda.max_memory_allocated(index) for index in range(torch.cuda.device_count()) ], "checked_at": datetime.now(timezone.utc).isoformat(), } output = OUTPUT_DIR / f"{run_id}.json" output.write_text(json.dumps(payload, indent=2), encoding="utf-8") print(json.dumps({ "status": "ok", "run_id": run_id, "net_model_parameter_reduction": injection["net_model_parameter_reduction"], "trainable_parameters": injection["trainable_parameter_count"], "training_summary": {key: value for key, value in training.items() if key != "losses"}, "evaluation_active": evaluation["metrics"], "evaluation_zero_residual": evaluation_zero["metrics"] if evaluation_zero else None, "record": str(output), }, indent=2)) if __name__ == "__main__": main()