tinymind-native-8b-remote-handoff/bundle/evaluation/int6_precision_ladder.py

"""INT4/INT6 precision ladder evidence for TinyMind sparse exports."""

from __future__ import annotations

from datetime import datetime, timezone
import json
from pathlib import Path

import torch

from model.sparse_int4 import INT4SparseLinear
from model.sparse_int6 import INT6SparseLinear


def _layer_error(layer: torch.nn.Linear, sparse: torch.nn.Module, x: torch.Tensor) -> dict:
    with torch.no_grad():
        dense = layer(x)
        approx = sparse(x)
        diff = (dense - approx).float()
    return {
        "mean_abs_error": float(diff.abs().mean().item()),
        "max_abs_error": float(diff.abs().max().item()),
        "root_mean_square_error": float(torch.sqrt((diff * diff).mean()).item()),
    }


def build_int6_precision_ladder(out_dir: str | Path, seed: int = 20260523) -> dict:
    torch.manual_seed(seed)
    layer = torch.nn.Linear(128, 32, bias=False)
    x = torch.randn(8, 128)
    int4 = INT4SparseLinear.from_dense(layer)
    int6 = INT6SparseLinear.from_dense(layer)
    int4_error = _layer_error(layer, int4, x)
    int6_error = _layer_error(layer, int6, x)
    int4_bytes = int(int4.packed_weight.numel() + int4.metadata.numel() * 4 + int4.scales.numel() * 4)
    int6_bytes = int(int6.packed_weight.numel() + int6.metadata.numel() * 4 + int6.scales.numel() * 4)
    dense_bytes = int(layer.weight.numel() * 2)
    int6_wins_drift = int6_error["mean_abs_error"] <= int4_error["mean_abs_error"]
    report = {
        "schema_version": "tinymind-int6-precision-ladder-v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "seed": seed,
        "formats": {
            "int4": {
                "format": int4.format_name,
                "alias": int4.user_alias,
                "artifact_bytes_reference": int4_bytes,
                "compression_vs_bf16_dense": dense_bytes / max(int4_bytes, 1),
                "error": int4_error,
            },
            "int6": {
                "format": int6.format_name,
                "alias": int6.user_alias,
                "artifact_bytes_reference": int6_bytes,
                "compression_vs_bf16_dense": dense_bytes / max(int6_bytes, 1),
                "error": int6_error,
            },
        },
        "decision": {
            "int6_wins_drift_over_int4": int6_wins_drift,
            "int6_smaller_than_int8_payload": True,
            "int6_smaller_than_int4": False,
            "recommended_use": "Use INT6 when quality drift matters more than minimum artifact size; keep INT4 for fastest/smallest path.",
        },
        "claim_gate": {
            "world_best_precision_claim_allowed": False,
            "requires_cuda_kernel_and_task_eval": True,
        },
    }
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    json_path = out / "int6_precision_ladder_report.json"
    md_path = out / "int6_precision_ladder_report.md"
    report["json_path"] = str(json_path)
    report["markdown_path"] = str(md_path)
    json_path.write_text(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
    md_path.write_text(_markdown(report), encoding="utf-8")
    return report


def _markdown(report: dict) -> str:
    i4 = report["formats"]["int4"]
    i6 = report["formats"]["int6"]
    return "\n".join(
        [
            "# TinyMind INT6 Precision Ladder",
            "",
            f"- INT4 MAE: {i4['error']['mean_abs_error']:.6f}",
            f"- INT6 MAE: {i6['error']['mean_abs_error']:.6f}",
            f"- INT4 reference bytes: {i4['artifact_bytes_reference']}",
            f"- INT6 reference bytes: {i6['artifact_bytes_reference']}",
            f"- INT6 wins drift over INT4: {report['decision']['int6_wins_drift_over_int4']}",
            f"- Recommended use: {report['decision']['recommended_use']}",
            "- World-best precision claim allowed: false",
            "",
        ]
    )