code/wal_tat/fixed_lowbit.py

"""Inference-only groupwise projections used for controlled PTQ baselines."""
from __future__ import annotations

from dataclasses import dataclass
import hashlib
import json
from pathlib import Path
from typing import Mapping, Optional, Sequence

import torch
import torch.nn as nn
import torch.nn.functional as F

from .binary import (
    ProxyBinaryLinear,
    ProxyBinaryMatrix,
    unpack_binary_codes,
    weighted_binary_project,
)
from .binary_packing import read_packed_binary_matrix
from .quantization import (
    weighted_symmetric_nz4_project,
    weighted_symmetric_q4_project,
)
from .scoring import exact_diagonal_ternary_project
from .packing import read_packed_matrix, unpack_ternary_codes
from .proxy import ProxyTernaryLinear, ProxyTernaryMatrix
from .whisper import get_module, set_module


class FixedGroupwiseLinear(nn.Module):
    """Materialized reference forward for arbitrary signed groupwise codes."""

    def __init__(
        self,
        codes: torch.Tensor,
        scales: torch.Tensor,
        *,
        in_features: int,
        compute_dtype: torch.dtype,
        bias: Optional[torch.Tensor] = None,
        allowed_codes: Optional[Sequence[int]] = None,
    ):
        super().__init__()
        if codes.ndim != 3 or scales.shape != codes.shape[:2]:
            raise ValueError("codes must be [out, groups, group_size] with matching scales")
        full_in_features = codes.shape[1] * codes.shape[2]
        if not 0 < in_features <= full_in_features:
            raise ValueError("in_features is incompatible with grouped codes")
        if full_in_features - in_features >= codes.shape[2]:
            raise ValueError("padding must be smaller than one group")
        if allowed_codes is not None:
            allowed = torch.tensor(tuple(allowed_codes), device=codes.device)
            if not torch.isin(codes, allowed).all():
                raise ValueError("codes contain a value outside the declared codebook")
        self.register_buffer("codes", codes.detach().to(torch.int8).clone())
        self.register_buffer("group_scales", scales.detach().float().clone())
        if bias is None:
            self.bias = None
        else:
            self.register_buffer("bias", bias.detach().clone())
        self.in_features = int(in_features)
        self.out_features = int(codes.shape[0])
        self.group_size = int(codes.shape[2])
        self.compute_dtype = compute_dtype
        self.register_buffer(
            "_evaluation_weight",
            (self.codes.float() * self.group_scales.unsqueeze(-1))
            .reshape(self.out_features, -1)[:, : self.in_features]
            .to(compute_dtype),
            persistent=False,
        )

    def effective_weight(self) -> torch.Tensor:
        return self._evaluation_weight.to(self.compute_dtype)

    def forward(self, value: torch.Tensor) -> torch.Tensor:
        bias = None if self.bias is None else self.bias.to(value.dtype)
        return F.linear(value, self._evaluation_weight.to(value.dtype), bias)


class FixedPartialGroupwiseLinear(nn.Module):
    """Reference forward with strict low-bit groups and exact BF16 fallback."""

    def __init__(
        self,
        codes: torch.Tensor,
        scales: torch.Tensor,
        committed_mask: torch.Tensor,
        base_weight: torch.Tensor,
        *,
        in_features: int,
        compute_dtype: torch.dtype,
        bias: Optional[torch.Tensor] = None,
        allowed_codes: Optional[Sequence[int]] = None,
    ):
        super().__init__()
        if codes.ndim != 3 or scales.shape != codes.shape[:2]:
            raise ValueError(
                "codes must be [out, groups, group_size] with matching scales"
            )
        if committed_mask.shape != codes.shape[:2] or committed_mask.dtype != torch.bool:
            raise ValueError("committed_mask must match the grouped code shape")
        if not committed_mask.any():
            raise ValueError("partial low-bit matrix must commit at least one group")
        full_in_features = codes.shape[1] * codes.shape[2]
        expected_base_shape = (codes.shape[0], in_features)
        if tuple(base_weight.shape) != expected_base_shape:
            raise ValueError(
                f"base_weight shape is {tuple(base_weight.shape)}, "
                f"expected {expected_base_shape}"
            )
        if not 0 < in_features <= full_in_features:
            raise ValueError("in_features is incompatible with grouped codes")
        if full_in_features - in_features >= codes.shape[2]:
            raise ValueError("padding must be smaller than one group")
        if allowed_codes is not None:
            allowed = torch.tensor(tuple(allowed_codes), device=codes.device)
            active = codes[committed_mask]
            if not torch.isin(active, allowed).all():
                raise ValueError(
                    "committed codes contain a value outside the declared codebook"
                )
        self.register_buffer("codes", codes.detach().to(torch.int8).clone())
        self.register_buffer("group_scales", scales.detach().float().clone())
        self.register_buffer("committed_mask", committed_mask.detach().clone())
        self.register_buffer(
            "base_weight", base_weight.detach().to(compute_dtype).clone()
        )
        if bias is None:
            self.bias = None
        else:
            self.register_buffer("bias", bias.detach().clone())
        self.in_features = int(in_features)
        self.out_features = int(codes.shape[0])
        self.group_size = int(codes.shape[2])
        self.compute_dtype = compute_dtype
        padded_base = F.pad(
            self.base_weight.float(), (0, full_in_features - self.in_features)
        ).view_as(self.codes)
        low_bit = self.codes.float() * self.group_scales.unsqueeze(-1)
        mixed = torch.where(
            self.committed_mask.unsqueeze(-1), low_bit, padded_base
        )
        self.register_buffer(
            "_evaluation_weight",
            mixed.reshape(self.out_features, -1)[:, : self.in_features].to(
                compute_dtype
            ),
            persistent=False,
        )

    def effective_weight(self) -> torch.Tensor:
        return self._evaluation_weight.to(self.compute_dtype)

    def forward(self, value: torch.Tensor) -> torch.Tensor:
        bias = None if self.bias is None else self.bias.to(value.dtype)
        return F.linear(value, self._evaluation_weight.to(value.dtype), bias)


@dataclass(frozen=True)
class ProjectionStatistics:
    name: str
    precision: str
    group_size: int
    weights: int
    weighted_error: float
    code_counts: Mapping[str, int]


@torch.no_grad()
def project_linear_module(
    linear: nn.Linear,
    *,
    precision: str,
    group_size: int,
    input_second_moment: Optional[torch.Tensor] = None,
) -> tuple[FixedGroupwiseLinear, float, Mapping[str, int]]:
    """Project one dense linear into a fixed B1/T3/NZ4/Q4 reference layer."""
    moment = (
        torch.ones(linear.in_features, device=linear.weight.device)
        if input_second_moment is None
        else input_second_moment.to(linear.weight.device)
    )
    if precision == "b1":
        codes, scales, error = weighted_binary_project(
            linear.weight, moment, group_size=group_size
        )
        allowed = (-1, 1)
    elif precision == "t3":
        codes, scales, error = exact_diagonal_ternary_project(
            linear.weight, moment, group_size=group_size
        )
        allowed = (-1, 0, 1)
    elif precision == "nz4":
        codes, scales, error = weighted_symmetric_nz4_project(
            linear.weight, moment, group_size=group_size
        )
        allowed = (-3, -1, 1, 3)
    elif precision == "q4":
        codes, scales, error = weighted_symmetric_q4_project(
            linear.weight, moment, group_size=group_size
        )
        allowed = tuple(range(-8, 8))
    else:
        raise ValueError("precision must be one of b1, t3, nz4, q4")
    values, counts = torch.unique(codes.cpu(), return_counts=True)
    histogram = {str(int(value)): int(count) for value, count in zip(values, counts)}
    fixed = FixedGroupwiseLinear(
        codes,
        scales,
        in_features=linear.in_features,
        compute_dtype=linear.weight.dtype,
        bias=linear.bias,
        allowed_codes=allowed,
    )
    return fixed, float(error.sum().item()), histogram


@torch.no_grad()
def install_fixed_projection(
    model: nn.Module,
    module_names: Sequence[str],
    *,
    precision: str,
    group_size: int,
    input_moments: Optional[Mapping[str, torch.Tensor]] = None,
) -> tuple[ProjectionStatistics, ...]:
    """Replace declared linears with deterministic projected references."""
    results = []
    for name in module_names:
        linear = get_module(model, name)
        if not isinstance(linear, nn.Linear):
            raise TypeError(f"{name} is not an nn.Linear")
        fixed, error, histogram = project_linear_module(
            linear,
            precision=precision,
            group_size=group_size,
            input_second_moment=(
                None if input_moments is None else input_moments.get(name)
            ),
        )
        fixed.to(device=linear.weight.device)
        set_module(model, name, fixed)
        results.append(
            ProjectionStatistics(
                name=name,
                precision=precision,
                group_size=group_size,
                weights=linear.weight.numel(),
                weighted_error=error,
                code_counts=histogram,
            )
        )
    return tuple(results)


@torch.no_grad()
def install_fixed_checkpoint(
    model: nn.Module,
    checkpoint: Mapping,
) -> tuple[str, ...]:
    """Install a pilot/checkpoint code+scale mapping into a fresh dense model."""
    allowed_by_precision = {
        "b1": (-1, 1),
        "t3": (-1, 0, 1),
        "nz4": (-3, -1, 1, 3),
        "q4": tuple(range(-8, 8)),
    }
    checkpoint_precision = str(checkpoint["precision"])
    if checkpoint_precision != "mixed" and checkpoint_precision not in allowed_by_precision:
        raise ValueError("checkpoint precision is unsupported")
    names = []
    for name, entry in checkpoint["matrices"].items():
        precision = str(entry.get("precision", checkpoint_precision))
        if precision not in allowed_by_precision:
            raise ValueError(f"{name} checkpoint precision is unsupported")
        linear = get_module(model, name)
        if not isinstance(linear, nn.Linear):
            raise TypeError(f"{name} is not an nn.Linear in the fresh model")
        codes = torch.as_tensor(entry["codes"], device=linear.weight.device)
        scales = torch.as_tensor(entry["scales"], device=linear.weight.device)
        bias = entry.get("bias")
        if bias is None:
            bias = linear.bias
        committed_mask = entry.get("committed_mask")
        if committed_mask is None:
            fixed = FixedGroupwiseLinear(
                codes,
                scales,
                in_features=linear.in_features,
                compute_dtype=linear.weight.dtype,
                bias=bias,
                allowed_codes=allowed_by_precision[precision],
            )
        else:
            if precision != "t3":
                raise ValueError(
                    f"{name} partial checkpoint currently supports T3 only"
                )
            mask = torch.as_tensor(
                committed_mask, device=linear.weight.device, dtype=torch.bool
            )
            stored_base = entry.get("base_weight")
            base_weight = (
                linear.weight
                if stored_base is None
                else torch.as_tensor(
                    stored_base,
                    device=linear.weight.device,
                    dtype=linear.weight.dtype,
                )
            )
            fixed = FixedPartialGroupwiseLinear(
                codes,
                scales,
                mask,
                base_weight,
                in_features=linear.in_features,
                compute_dtype=linear.weight.dtype,
                bias=bias,
                allowed_codes=allowed_by_precision[precision],
            )
        fixed = fixed.to(device=linear.weight.device)
        set_module(model, name, fixed)
        names.append(name)
    return tuple(names)


@torch.no_grad()
def install_trainable_lowbit_checkpoint(
    model: nn.Module,
    checkpoint: Mapping,
    *,
    temperature: float = 0.35,
    initial_proxy_magnitude: float = 0.75,
    initial_proxy_magnitudes: Optional[Mapping[str, float]] = None,
    initial_zero_proxy_boundaries: Optional[Mapping[str, float]] = None,
    fake_fp16_scale: bool = True,
) -> Mapping[str, ProxyBinaryMatrix | ProxyTernaryMatrix]:
    """Install a B1/T3 checkpoint as exact-hard-forward trainable proxies.

    This deliberately rejects Q4 and other rescue formats.  It is the bridge
    used by global recovery: the deployed codebook is already active in the
    forward pass while gradients may move proxy codes across its boundaries.
    """
    if temperature <= 0:
        raise ValueError("temperature must be positive")
    if initial_proxy_magnitude <= 0:
        raise ValueError("initial_proxy_magnitude must be positive")
    per_matrix_magnitudes = dict(initial_proxy_magnitudes or {})
    per_matrix_zero_boundaries = dict(initial_zero_proxy_boundaries or {})
    unknown_names = set(per_matrix_magnitudes) - set(checkpoint["matrices"])
    if unknown_names:
        raise ValueError(
            f"initial proxy magnitudes name unknown matrices: {sorted(unknown_names)}"
        )
    unknown_zero_names = set(per_matrix_zero_boundaries) - set(
        checkpoint["matrices"]
    )
    if unknown_zero_names:
        raise ValueError(
            "initial zero proxy boundaries name unknown matrices: "
            f"{sorted(unknown_zero_names)}"
        )
    checkpoint_precision = str(checkpoint["precision"])
    matrices: dict[str, ProxyBinaryMatrix | ProxyTernaryMatrix] = {}
    for name, entry in checkpoint["matrices"].items():
        precision = str(entry.get("precision", checkpoint_precision))
        if precision not in {"b1", "t3"}:
            raise ValueError(
                f"{name} uses {precision}; global strict-low-bit recovery "
                "accepts only b1/t3"
            )
        linear = get_module(model, name)
        if not isinstance(linear, nn.Linear):
            raise TypeError(f"{name} is not an nn.Linear in the fresh model")
        codes = torch.as_tensor(entry["codes"], device=linear.weight.device)
        scales = torch.as_tensor(
            entry["scales"], device=linear.weight.device, dtype=torch.float32
        )
        committed_mask = entry.get("committed_mask")
        if committed_mask is not None:
            if precision != "t3":
                raise ValueError(
                    f"{name} partial trainable checkpoint currently supports T3 only"
                )
            committed_mask = torch.as_tensor(
                committed_mask, device=linear.weight.device, dtype=torch.bool
            )
        bias = entry.get("bias")
        if bias is None:
            bias = linear.bias
        elif not isinstance(bias, torch.Tensor):
            bias = torch.as_tensor(bias)
        if bias is not None:
            bias = bias.to(device=linear.weight.device, dtype=linear.weight.dtype)
        proxy_magnitude = float(
            per_matrix_magnitudes.get(name, initial_proxy_magnitude)
        )
        if proxy_magnitude <= 0:
            raise ValueError(f"{name} initial proxy magnitude must be positive")
        if precision == "t3" and proxy_magnitude < 0.5:
            raise ValueError(
                f"{name} initial proxy magnitude must be at least 0.5 for T3"
            )
        if precision == "b1":
            if name in per_matrix_zero_boundaries:
                raise ValueError(
                    f"{name} cannot use an initial zero proxy boundary in B1"
                )
            matrix = ProxyBinaryMatrix(
                codes,
                scales,
                compute_dtype=linear.weight.dtype,
                temperature=temperature,
                initial_proxy_magnitude=proxy_magnitude,
                fake_fp16_scale=fake_fp16_scale,
            )
            wrapper = ProxyBinaryLinear(matrix, bias)
        else:
            matrix = ProxyTernaryMatrix(
                codes,
                scales,
                compute_dtype=linear.weight.dtype,
                temperature=temperature,
                committed_mask=committed_mask,
                master_weight=(
                    linear.weight
                    if entry.get("base_weight") is None
                    else torch.as_tensor(
                        entry["base_weight"],
                        device=linear.weight.device,
                        dtype=linear.weight.dtype,
                    )
                ),
                initial_proxy_magnitude=proxy_magnitude,
                initial_zero_proxy_boundary=per_matrix_zero_boundaries.get(name),
                fake_fp16_scale=fake_fp16_scale,
            )
            wrapper = ProxyTernaryLinear(matrix, bias)
        wrapper.to(device=linear.weight.device)
        set_module(model, name, wrapper)
        matrices[name] = matrix
    return matrices


def _sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        while chunk := stream.read(1 << 20):
            digest.update(chunk)
    return digest.hexdigest()


def _artifact_member(directory: Path, name: str) -> Path:
    path = (directory / name).resolve()
    if path.parent != directory.resolve():
        raise ValueError("artifact member must be a direct child of the manifest directory")
    return path


@torch.no_grad()
def install_packed_t3_manifest(
    model: nn.Module,
    manifest_path: str | Path,
) -> tuple[str, ...]:
    """Load strict, fully committed T3 matrices from a WAL Q2 manifest.

    This is a correctness/reference loader. ``FixedGroupwiseLinear`` still
    materializes an evaluation tensor, so deployment speed requires a native
    packed kernel; however, every value installed here originates from the
    serialized two-bit codes and FP16 scales rather than the training checkpoint.
    """
    manifest_path = Path(manifest_path).expanduser().resolve()
    manifest = json.loads(manifest_path.read_text())
    if manifest.get("schema_version") != 1 or manifest.get("precision") != "t3":
        raise ValueError("manifest is not a supported strict T3 artifact")
    directory = manifest_path.parent
    names: list[str] = []
    for entry in manifest.get("entries", ()):
        name = str(entry["name"])
        linear = get_module(model, name)
        if not isinstance(linear, nn.Linear):
            raise TypeError(f"{name} is not an nn.Linear in the fresh model")
        matrix_path = _artifact_member(directory, str(entry["matrix_file"]))
        if matrix_path.stat().st_size != int(entry["matrix_bytes"]):
            raise ValueError(f"{name} packed matrix size does not match the manifest")
        if _sha256_file(matrix_path) != str(entry["matrix_sha256"]):
            raise ValueError(f"{name} packed matrix SHA-256 does not match the manifest")
        packed = read_packed_matrix(matrix_path)
        if packed.committed_groups != packed.total_groups:
            raise ValueError("strict T3 manifest cannot contain BF16 fallback groups")
        if packed.shape != (linear.out_features, linear.in_features):
            raise ValueError(f"{name} packed shape does not match the fresh model")
        groups_per_row = packed.total_groups // packed.shape[0]
        codes = unpack_ternary_codes(
            packed.codes_packed,
            packed.committed_groups * packed.group_size,
        ).view(packed.shape[0], groups_per_row, packed.group_size)
        scales = packed.scales_fp16.view(packed.shape[0], groups_per_row)

        bias = None
        if entry.get("bias_file") is not None:
            bias_path = _artifact_member(directory, str(entry["bias_file"]))
            if bias_path.stat().st_size != int(entry["bias_bytes"]):
                raise ValueError(f"{name} bias size does not match the manifest")
            if _sha256_file(bias_path) != str(entry["bias_sha256"]):
                raise ValueError(f"{name} bias SHA-256 does not match the manifest")
            raw = bias_path.read_bytes()
            bias = torch.frombuffer(bytearray(raw), dtype=torch.bfloat16).clone()
            if bias.numel() != linear.out_features:
                raise ValueError(f"{name} bias shape does not match the fresh model")
        elif linear.bias is not None:
            raise ValueError(f"{name} manifest omitted a required bias")

        fixed = FixedGroupwiseLinear(
            codes,
            scales,
            in_features=linear.in_features,
            compute_dtype=linear.weight.dtype,
            bias=bias,
            allowed_codes=(-1, 0, 1),
        ).to(device=linear.weight.device)
        set_module(model, name, fixed)
        names.append(name)
    if not names:
        raise ValueError("manifest contains no matrices")
    return tuple(names)


@torch.no_grad()
def install_packed_lowbit_manifest(
    model: nn.Module,
    manifest_path: str | Path,
) -> tuple[str, ...]:
    """Install a strict schema-v2 B1/T3 overlay from serialized payloads.

    This is a correctness loader. It materializes a reference evaluation tensor;
    native speed and memory savings still require a packed B1/T3 kernel.
    """
    manifest_path = Path(manifest_path).expanduser().resolve()
    manifest = json.loads(manifest_path.read_text())
    schema_policy = (
        manifest.get("schema_version"),
        manifest.get("precision_policy"),
    )
    if schema_policy not in {
        (2, "strict_b1_t3_only"),
        (3, "strict_b1_t3_partial_bf16_fallback"),
    }:
        raise ValueError("manifest is not a supported strict mixed low-bit artifact")
    directory = manifest_path.parent
    names: list[str] = []
    for entry in manifest.get("entries", ()):
        name = str(entry["name"])
        precision = str(entry["precision"])
        if precision not in {"b1", "t3"}:
            raise ValueError(f"{name} uses unsupported packed precision {precision}")
        linear = get_module(model, name)
        if not isinstance(linear, nn.Linear):
            raise TypeError(f"{name} is not an nn.Linear in the fresh model")
        matrix_path = _artifact_member(directory, str(entry["matrix_file"]))
        if matrix_path.stat().st_size != int(entry["matrix_bytes"]):
            raise ValueError(f"{name} packed matrix size does not match the manifest")
        if _sha256_file(matrix_path) != str(entry["matrix_sha256"]):
            raise ValueError(f"{name} packed matrix SHA-256 does not match the manifest")

        if precision == "b1":
            packed_b1 = read_packed_binary_matrix(matrix_path)
            if packed_b1.shape != (linear.out_features, linear.in_features):
                raise ValueError(f"{name} packed shape does not match the fresh model")
            groups_per_row = packed_b1.total_groups // packed_b1.shape[0]
            codes = unpack_binary_codes(
                packed_b1.codes_packed,
                packed_b1.total_groups * packed_b1.group_size,
            ).view(packed_b1.shape[0], groups_per_row, packed_b1.group_size)
            scales = packed_b1.scales_fp16.view(packed_b1.shape[0], groups_per_row)
            allowed_codes = (-1, 1)
        else:
            packed_t3 = read_packed_matrix(matrix_path)
            if packed_t3.shape != (linear.out_features, linear.in_features):
                raise ValueError(f"{name} packed shape does not match the fresh model")
            groups_per_row = packed_t3.total_groups // packed_t3.shape[0]
            mask = packed_t3.committed_mask().view(
                packed_t3.shape[0], groups_per_row
            )
            active_codes = unpack_ternary_codes(
                packed_t3.codes_packed,
                packed_t3.committed_groups * packed_t3.group_size,
            ).view(packed_t3.committed_groups, packed_t3.group_size)
            codes = torch.zeros(
                (
                    packed_t3.shape[0],
                    groups_per_row,
                    packed_t3.group_size,
                ),
                dtype=torch.int8,
            )
            codes[mask] = active_codes
            scales = torch.ones(
                (packed_t3.shape[0], groups_per_row), dtype=torch.float16
            )
            scales[mask] = packed_t3.scales_fp16
            allowed_codes = (-1, 0, 1)

        bias = None
        if entry.get("bias_file") is not None:
            bias_path = _artifact_member(directory, str(entry["bias_file"]))
            if bias_path.stat().st_size != int(entry["bias_bytes"]):
                raise ValueError(f"{name} bias size does not match the manifest")
            if _sha256_file(bias_path) != str(entry["bias_sha256"]):
                raise ValueError(f"{name} bias SHA-256 does not match the manifest")
            bias = torch.frombuffer(
                bytearray(bias_path.read_bytes()), dtype=torch.bfloat16
            ).clone()
            if bias.numel() != linear.out_features:
                raise ValueError(f"{name} bias shape does not match the fresh model")
        elif linear.bias is not None:
            raise ValueError(f"{name} manifest omitted a required bias")

        if precision == "t3" and packed_t3.committed_groups != packed_t3.total_groups:
            grouped_base = torch.zeros(
                (
                    packed_t3.shape[0],
                    groups_per_row,
                    packed_t3.group_size,
                ),
                dtype=torch.bfloat16,
            )
            grouped_base[~mask] = packed_t3.fallback_bf16.view(
                packed_t3.total_groups - packed_t3.committed_groups,
                packed_t3.group_size,
            )
            base_weight = grouped_base.reshape(packed_t3.shape[0], -1)[
                :, : linear.in_features
            ]
            fixed = FixedPartialGroupwiseLinear(
                codes,
                scales,
                mask,
                base_weight,
                in_features=linear.in_features,
                compute_dtype=linear.weight.dtype,
                bias=bias,
                allowed_codes=allowed_codes,
            ).to(device=linear.weight.device)
        else:
            fixed = FixedGroupwiseLinear(
                codes,
                scales,
                in_features=linear.in_features,
                compute_dtype=linear.weight.dtype,
                bias=bias,
                allowed_codes=allowed_codes,
            ).to(device=linear.weight.device)
        set_module(model, name, fixed)
        names.append(name)
    if not names:
        raise ValueError("manifest contains no matrices")
    return tuple(names)