| |
| """Run one arm of the locked WikiText-train QH032 screen.""" |
|
|
| 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/qh032-train-screen-lock.json" |
| TOKENS = PROJECT / "datasets/processed/wikitext2_qwen38/train-tokens.pt" |
| OUTPUT_DIR = PROJECT / "artifacts/qh032_train_screen" |
| sys.path.insert(0, str(PROJECT / "code/src")) |
| from quantum_qwen38.qh032_pauli_observable_replacement import ( |
| ExactPauliObservableCore, QH032Config, install_qh032, unique_trainable_parameters, |
| ) |
|
|
|
|
| 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 block(tokens, index, length, device): |
| start = index * length |
| return tokens[start : start + length].to(device=device, dtype=torch.long).unsqueeze(0) |
|
|
|
|
| def set_projections(model, modules) -> None: |
| for index, module in modules.items(): |
| model.model.language_model.layers[index].self_attn.v_proj = module |
|
|
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--candidate", choices=["base", "sharedsvd1024", "qh032", "cc032"], required=True) |
| return parser.parse_args() |
|
|
|
|
| @torch.inference_mode() |
| def evaluate(model, tokens, indices, length): |
| rows, seconds = [], [] |
| for index in indices: |
| inputs = block(tokens, index, length, model.device) |
| started = time.perf_counter() |
| loss = model(input_ids=inputs, labels=inputs, use_cache=False).loss |
| rows.append({"block_index": index, "mean_token_nll": float(loss), "seconds": time.perf_counter() - started}) |
| seconds.append(rows[-1]["seconds"]) |
| nll = float(np.mean([row["mean_token_nll"] for row in rows])) |
| return { |
| "metrics": {"sequence_count": len(rows), "mean_token_nll": nll, "token_perplexity": math.exp(nll), |
| "mean_sequence_seconds": float(np.mean(seconds))}, |
| "rows": rows, |
| } |
|
|
|
|
| def train(model, replacements, originals, tokens, lock): |
| parameters = list(unique_trainable_parameters(replacements)) |
| protocol = lock["training_protocol"] |
| optimizer = torch.optim.AdamW(parameters, lr=protocol["learning_rate"], weight_decay=protocol["weight_decay"]) |
| backbone = model.model.language_model |
| config = QH032Config() |
| core = replacements[config.layer_indices[0]].core |
| calls_before = int(core.circuit_calls.item()) |
| losses, seconds, theta_grads, alpha_grads = [], [], [], [] |
| for position, index in enumerate(lock["train_indices"], start=1): |
| inputs = block(tokens, int(index), int(lock["sequence_length"]), model.device) |
| set_projections(model, originals) |
| with torch.inference_mode(): |
| teacher = backbone(input_ids=inputs, use_cache=False, return_dict=True).last_hidden_state.detach() |
| set_projections(model, replacements) |
| optimizer.zero_grad(set_to_none=True) |
| started = time.perf_counter() |
| student = backbone(input_ids=inputs, use_cache=False, return_dict=True).last_hidden_state |
| denominator = teacher.float().square().mean().clamp_min(1.0e-8) |
| loss = (student.float() - teacher.float()).square().mean() / denominator |
| loss.backward() |
| theta_grads.append(float(core.theta.grad.float().norm())) |
| alpha_grads.append(float(core.alpha.grad.float().norm())) |
| torch.nn.utils.clip_grad_norm_(parameters, protocol["gradient_clip_norm"]) |
| optimizer.step() |
| torch.cuda.synchronize() |
| losses.append(float(loss.detach())) |
| seconds.append(time.perf_counter() - started) |
| if position % 64 == 0: |
| print(json.dumps({"trained": position, "recent_loss": float(np.mean(losses[-64:])), "recent_seconds": float(np.mean(seconds[-64:]))}), flush=True) |
| calls = int(core.circuit_calls.item()) - calls_before |
| expected = len(lock["train_indices"]) * len(config.layer_indices) if isinstance(core, ExactPauliObservableCore) else 0 |
| if calls != expected: |
| raise AssertionError(f"training simulator calls {calls} != {expected}") |
| frozen_grads = sum(int(p.grad is not None) for p in model.parameters() if not p.requires_grad) |
| if frozen_grads: |
| raise AssertionError("frozen parameter received gradient") |
| edge = 32 |
| return { |
| "steps": len(losses), "loss_first_window_mean": float(np.mean(losses[:edge])), |
| "loss_last_window_mean": float(np.mean(losses[-edge:])), "loss_min": float(np.min(losses)), |
| "theta_gradient_norm_mean": float(np.mean(theta_grads)), "alpha_gradient_norm_mean": float(np.mean(alpha_grads)), |
| "mean_student_step_seconds": float(np.mean(seconds)), "simulator_calls": calls, |
| "frozen_gradient_tensor_count": frozen_grads, "final_theta_l2_norm": float(core.theta.detach().norm()), |
| "final_alpha_l2_norm": float(core.alpha.detach().norm()), |
| "final_gammas": {str(i): float(m.gamma.detach()) for i, m in replacements.items()}, |
| "losses": losses, |
| } |
|
|
|
|
| def checkpoint(replacements): |
| first = sorted(replacements)[0] |
| return { |
| "core.theta": replacements[first].core.theta.detach().cpu(), |
| "core.alpha": replacements[first].core.alpha.detach().cpu(), |
| **{f"gamma.{index}": module.gamma.detach().cpu() for index, module in replacements.items()}, |
| } |
|
|
|
|
| def main(): |
| args = parse_args() |
| if not torch.cuda.is_available(): |
| raise RuntimeError("QH032 screen is GPU-only") |
| lock = json.loads(LOCK.read_text(encoding="utf-8")) |
| if lock["status"] != "locked_before_qh032_train_only_screen_results": |
| raise RuntimeError("invalid QH032 train-screen lock") |
| random.seed(lock["seed"]); np.random.seed(lock["seed"]); torch.manual_seed(lock["seed"]); torch.cuda.manual_seed_all(lock["seed"]) |
| tokens = torch.load(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", |
| ) |
| model.eval(); model.config.use_cache = False |
| replacements = originals = None |
| if args.candidate == "base": |
| model.requires_grad_(False) |
| audit = {"candidate": "base", "net_model_parameter_reduction": 0, "trainable_parameter_count": 0} |
| training = {"steps": 0} |
| else: |
| branch = "classical" if args.candidate == "cc032" else "quantum" |
| replacements, originals, audit = install_qh032(model, QH032Config(), branch=branch) |
| if args.candidate == "sharedsvd1024": |
| for module in replacements.values(): |
| module.branch_enabled = False; module.requires_grad_(False) |
| training = {"steps": 0} |
| else: |
| training = train(model, replacements, originals, tokens, lock) |
| evaluation = evaluate(model, tokens, [int(i) for i in lock["eval_indices"]], int(lock["sequence_length"])) |
| OUTPUT_DIR.mkdir(parents=True, exist_ok=True) |
| if replacements is not None and args.candidate in {"qh032", "cc032"}: |
| torch.save(checkpoint(replacements), OUTPUT_DIR / f"{args.candidate}-screen-checkpoint.pt") |
| payload = { |
| "status": "ok", "candidate": args.candidate, "dataset_lock": str(LOCK), "dataset_lock_sha256": sha256(LOCK), |
| "dataset": {"name": lock["dataset"], "train_indices": lock["train_indices"] if training["steps"] else [], |
| "eval_indices": lock["eval_indices"], "validation_rows_inspected": 0, "test_rows_inspected": 0}, |
| "injection": audit, "training": training, "evaluation": evaluation, |
| "created_at": datetime.now(timezone.utc).isoformat(), |
| } |
| output = OUTPUT_DIR / f"screen-{args.candidate}-s{lock['seed']}-n{lock['train_count']}-e{lock['eval_count']}.json" |
| output.write_text(json.dumps(payload, indent=2), encoding="utf-8") |
| print(json.dumps({"status": "ok", "candidate": args.candidate, "nll": evaluation["metrics"]["mean_token_nll"], |
| "output": str(output), "sha256": sha256(output)}, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
|
|
