scripts/measure_glm53_packed_student_kld.py

#!/usr/bin/env python3
"""Measure exact tokenwise KL(BF16 teacher || packed student) from sealed captures."""

from __future__ import annotations

import argparse
import io
import json
from pathlib import Path
import time

import numpy as np

from quant_pipeline.core.artifacts import atomic_write, canonical_json, prepare_empty_destination, sha256_bytes, sha256_file, write_json
from quant_pipeline.evaluation.glm53_logits import load_capture_receipt, summarize


FINAL_WINDOW_IDS = tuple(f"final-{index:04d}" for index in range(25))
FINAL_PREDICTION_POSITIONS = 25 * 2047


def _record_map(receipt):
    return {row["window_id"]: row for row in receipt["logit_files"]}


def _load_slice(path: Path, start: int, stop: int):
    from safetensors import safe_open

    with safe_open(path, framework="pt", device="cpu") as handle:
        return handle.get_slice("logits")[start:stop].clone()


def _gpu_token_kld(teacher, student, device):
    import torch

    if teacher.shape != student.shape or teacher.ndim != 2:
        raise ValueError("teacher/student logit geometry mismatch")
    teacher64 = teacher.to(device=device, dtype=torch.float64, non_blocking=True)
    student64 = student.to(device=device, dtype=torch.float64, non_blocking=True)
    if not torch.isfinite(teacher64).all() or not torch.isfinite(student64).all():
        raise ValueError("teacher/student logits must be finite")
    teacher_logp = torch.log_softmax(teacher64, dim=-1)
    student_logp = torch.log_softmax(student64, dim=-1)
    values = torch.sum(torch.exp(teacher_logp) * (teacher_logp - student_logp), dim=-1)
    matches = int(torch.count_nonzero(torch.argmax(teacher64, dim=-1) == torch.argmax(student64, dim=-1)))
    return values.cpu().numpy(), matches


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--teacher-receipt", type=Path, required=True)
    parser.add_argument("--student-receipt", type=Path, required=True)
    parser.add_argument("--student-label", choices=("uniform-k4", "uniform-k6"), required=True)
    parser.add_argument("--chunk-positions", type=int, default=16)
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--execute", action="store_true")
    args = parser.parse_args()
    if args.chunk_positions < 1:
        raise ValueError("chunk positions must be positive")
    teacher = load_capture_receipt(args.teacher_receipt.resolve(), expected_role="bf16_teacher")
    student = load_capture_receipt(args.student_receipt.resolve(), expected_role="packed_student")
    if teacher["token_panel_receipt_sha256"] != student["token_panel_receipt_sha256"]:
        raise ValueError("teacher and student captures use different sealed token panels")
    if teacher["vocab_size"] != student["vocab_size"]:
        raise ValueError("teacher and student vocabularies differ")
    if not student.get("runtime_reader_sha256"):
        raise ValueError("packed-student capture is not bound to an exact runtime reader")
    teacher_rows = _record_map(teacher)
    student_rows = _record_map(student)
    if set(teacher_rows) != set(student_rows):
        raise ValueError("teacher and student window sets differ")
    if (
        tuple(sorted(teacher_rows)) != FINAL_WINDOW_IDS
        or any(row.get("role") != "final" for row in teacher_rows.values())
        or any(row.get("role") != "final" for row in student_rows.values())
        or sum(int(row.get("prediction_positions", -1)) for row in teacher_rows.values())
        != FINAL_PREDICTION_POSITIONS
    ):
        raise ValueError("KLD qualification requires the sealed 25-window final panel only")
    for window_id, left in teacher_rows.items():
        right = student_rows[window_id]
        fields = ("document_id", "domain", "role", "token_ids_sha256", "attention_mask_sha256", "prediction_positions")
        if any(left[field] != right[field] for field in fields):
            raise ValueError(f"student capture relabels sealed window {window_id}")
    plan = {
        "schema": "quant-pipeline.glm53-packed-student-kld-plan.v1",
        "teacher_receipt_sha256": teacher["receipt_sha256"],
        "student_receipt_sha256": student["receipt_sha256"],
        "student_label": args.student_label,
        "kld_direction": "teacher_to_student",
        "chunk_positions": args.chunk_positions,
        "compute_device": args.device,
        "output": str(args.output.resolve()),
        "dry_run": not args.execute,
    }
    print(json.dumps(plan, sort_keys=True), flush=True)
    if not args.execute:
        return 0

    started = time.monotonic()
    import torch

    device = torch.device(args.device)
    if device.type != "cuda" or not torch.cuda.is_available():
        raise RuntimeError("packed-student KLD is pinned to an available CUDA device")
    torch.cuda.set_device(device)
    output = prepare_empty_destination(args.output.resolve())
    token_values = []
    top1_matches = 0
    per_window = []
    for window_id in sorted(teacher_rows):
        teacher_row = teacher_rows[window_id]
        student_row = student_rows[window_id]
        count = int(teacher_row["prediction_positions"])
        values = np.empty(count, dtype=np.float64)
        for start in range(0, count, args.chunk_positions):
            stop = min(start + args.chunk_positions, count)
            teacher_logits = _load_slice(Path(teacher_row["path"]), start, stop)
            student_logits = _load_slice(Path(student_row["path"]), start, stop)
            if teacher_logits.shape != (stop - start, int(teacher["vocab_size"])) or student_logits.shape != teacher_logits.shape:
                raise ValueError(f"logit geometry mismatch in {window_id}")
            chunk_values, chunk_matches = _gpu_token_kld(teacher_logits, student_logits, device)
            values[start:stop] = chunk_values
            top1_matches += chunk_matches
        token_values.append(values)
        per_window.append(
            {
                "window_id": window_id,
                "document_id": teacher_row["document_id"],
                "domain": teacher_row["domain"],
                "role": teacher_row["role"],
                "summary": summarize(values),
            }
        )
    all_values = np.concatenate(token_values)
    buffer = io.BytesIO()
    np.save(buffer, all_values, allow_pickle=False)
    token_path = output / "tokenwise-kld.npy"
    atomic_write(token_path, buffer.getvalue())
    domains = {}
    for domain in sorted({row["domain"] for row in per_window}):
        indices = [index for index, row in enumerate(per_window) if row["domain"] == domain]
        domains[domain] = summarize(np.concatenate([token_values[index] for index in indices]))
    overall = summarize(all_values)
    report = {
        "schema": "quant-pipeline.glm53-packed-student-kld.v1",
        "teacher_receipt_sha256": teacher["receipt_sha256"],
        "student_receipt_sha256": student["receipt_sha256"],
        "student_label": args.student_label,
        "student_checkpoint_identity_sha256": student["checkpoint_identity_sha256"],
        "runtime_reader_sha256": student["runtime_reader_sha256"],
        "token_panel_receipt_sha256": teacher["token_panel_receipt_sha256"],
        "teacher_backend_identity_sha256": teacher["backend_identity_sha256"],
        "student_backend_identity_sha256": student["backend_identity_sha256"],
        "qualification_panel_final_only": True,
        "qualification_window_count": len(per_window),
        "kld_direction": "teacher_to_student",
        "metric": "tokenwise KL over exact jointly-valid causal prediction positions",
        "compute_device": str(device),
        "compute_dtype": "float64",
        "summary": overall,
        "per_domain": domains,
        "per_window": per_window,
        "top1_agreement": top1_matches / int(all_values.size),
        "mean_kld_lt_0_06": bool(overall["mean"] < 0.06),
        "tokenwise_kld_path": str(token_path.resolve()),
        "tokenwise_kld_bytes": token_path.stat().st_size,
        "tokenwise_kld_sha256": sha256_file(token_path),
        "elapsed_seconds": time.monotonic() - started,
    }
    report["report_sha256"] = sha256_bytes(canonical_json(report))
    write_json(output / "kld-report.json", report)
    print(json.dumps({"ok": True, **report}, sort_keys=True), flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())