runtime/orbitquant_wan_a2/source_quant.py

from __future__ import annotations

import gc
import hashlib
import json
import os
import shutil
from pathlib import Path
from typing import Dict

import torch
from safetensors import safe_open
from safetensors.torch import save_file

from .nibbles import pack_uint4, unpack_uint4
from .orbitquant_math import EPS, fwht_last_dim, nearest_codes
from .rotation_bank import RotationBank

MODEL_ID = "Wan-AI/Wan2.2-Animate-2-14B-Distilled-Diffusers"
DEFAULT_SEED = 20260702

TARGET_SUFFIXES = (
    ".block.self_attn.q.weight",
    ".block.self_attn.k.weight",
    ".block.self_attn.v.weight",
    ".block.self_attn.o.weight",
    ".block.cross_attn.q.weight",
    ".block.cross_attn.k.weight",
    ".block.cross_attn.v.weight",
    ".block.cross_attn.o.weight",
    ".block.cross_attn.k_img.weight",
    ".block.cross_attn.v_img.weight",
    ".block.ffn.0.weight",
    ".block.ffn.2.weight",
)


def is_target_key(key: str) -> bool:
    return key.startswith("blocks.") and key.endswith(TARGET_SUFFIXES)


def _weight_map(root: Path) -> dict[str, str]:
    indexes = sorted(root.glob("*.safetensors.index.json"))
    if not indexes:
        raise FileNotFoundError(f"no safetensors index in {root}")
    payload = json.loads(indexes[0].read_text())
    return dict(payload["weight_map"])


def _group_by_file(wm: dict[str, str]) -> Dict[str, list[str]]:
    out: Dict[str, list[str]] = {}
    for k, f in wm.items():
        out.setdefault(f, []).append(k)
    for keys in out.values():
        keys.sort()
    return out


def _quantize_rows_to_packed(
    w: torch.Tensor,
    bank_item: dict,
    *,
    device: torch.device,
    row_chunk: int,
) -> tuple[torch.Tensor, torch.Tensor, dict]:
    """OrbitQuant offline W4 exactly in the paper's operation order.

    W' = W Pi^T
    row norm r' = ||w'||_2
    unit row = w'/r'
    nearest-centroid Lloyd-Max W4 direction

    The paper stores r' in BF16.  This runtime therefore stores BF16 row scales,
    not FP32 scales.  Codes are packed uint4 in [N,K/2] here and transposed later
    to GEMM-native [K/2,N].
    """
    if w.ndim != 2:
        raise ValueError(f"weight must be rank-2, got {tuple(w.shape)}")
    n, d = map(int, w.shape)
    h = int(bank_item["block_size"].item())
    perm = bank_item["perm"].to(device=device, dtype=torch.long)
    signs = bank_item["signs"].to(device=device, dtype=torch.float32)
    cb = bank_item["codebook"].to(device=device, dtype=torch.float32)

    packed = torch.empty((n, d // 2), dtype=torch.uint8, device="cpu")
    scales_bf16 = torch.empty((n,), dtype=torch.bfloat16, device="cpu")

    mse_sum = 0.0
    mae_sum = 0.0
    elem_count = 0
    code_hist = torch.zeros(16, dtype=torch.int64)

    for s in range(0, n, row_chunk):
        e = min(n, s + row_chunk)
        x = w[s:e].to(device=device, dtype=torch.float32)
        rot = x.index_select(-1, perm) * signs
        rot = fwht_last_dim(rot, h)
        norm = torch.linalg.vector_norm(rot, ord=2, dim=-1)
        # OrbitQuant paper stores the row-norm vector in BF16.
        norm_bf16 = norm.to(torch.bfloat16)
        # The fake-quantized weight uses the stored magnitude when reconstructing.
        norm_used = norm_bf16.float()
        unit = rot / (norm[:, None] + EPS)
        codes = nearest_codes(unit, cb)
        packed[s:e].copy_(pack_uint4(codes).cpu())
        scales_bf16[s:e].copy_(norm_bf16.cpu())

        decoded = cb[codes.long()] * norm_used[:, None]
        diff = decoded - rot
        mse_sum += float((diff * diff).sum().item())
        mae_sum += float(diff.abs().sum().item())
        elem_count += int(diff.numel())
        code_hist += torch.bincount(codes.flatten().cpu().long(), minlength=16)

        del x, rot, norm, norm_bf16, norm_used, unit, codes, decoded, diff

    return packed, scales_bf16, {
        "mse_rotated_weight": mse_sum / max(1, elem_count),
        "mae_rotated_weight": mae_sum / max(1, elem_count),
        "code_histogram": code_hist.tolist(),
        "row_scale_dtype": "bfloat16",
        "block_size": h,
    }


def _verify_one_chunk(
    w: torch.Tensor,
    packed_row: torch.Tensor,
    scales_bf16: torch.Tensor,
    bank_item: dict,
    *,
    rows: int = 2,
    device: torch.device | str = "cpu",
) -> dict:
    """Strict packed-decode audit on the quantization device.

    Quantization may run on CUDA. Recomputing nearest-centroid decisions on
    CPU is not a valid bit-exact packing test because FP32 FWHT/norm rounding
    can move a coordinate lying essentially on a Lloyd-Max decision boundary
    into the adjacent bin.

    This audit therefore recomputes the OrbitQuant codes on the same device
    that generated them, then independently unpacks the stored uint4 nibbles.

    The acceptance requirement remains exact 1.0.
    """
    rows = min(
        int(rows),
        int(w.shape[0]),
    )

    d = int(w.shape[1])

    dev = torch.device(device)

    h = int(
        bank_item["block_size"].item()
    )

    perm = bank_item["perm"].to(
        device=dev,
        dtype=torch.long,
    )

    signs = bank_item["signs"].to(
        device=dev,
        dtype=torch.float32,
    )

    cb_dev = bank_item["codebook"].to(
        device=dev,
        dtype=torch.float32,
    )

    # CPU copy is used only after code decisions have already been made.
    cb_cpu = (
        bank_item["codebook"]
        .to(
            device="cpu",
            dtype=torch.float32,
        )
        .contiguous()
    )

    src = w[:rows].to(
        device=dev,
        dtype=torch.float32,
    )

    rot = (
        src.index_select(
            -1,
            perm,
        )
        * signs
    )

    rot = fwht_last_dim(
        rot,
        h,
    )

    norm = torch.linalg.vector_norm(
        rot,
        ord=2,
        dim=-1,
    )

    unit = rot / (
        norm[:, None]
        + EPS
    )

    # Expected OrbitQuant decisions recomputed on the SAME device
    # as the original quantization.
    expected_codes = nearest_codes(
        unit,
        cb_dev,
    ).cpu()

    # Independent uint4 decode of what was actually stored.
    got_codes = unpack_uint4(
        packed_row[:rows],
        d,
    ).cpu()

    code_exact = float(
        (
            expected_codes
            == got_codes
        )
        .float()
        .mean()
        .item()
    )

    expected_scale = (
        norm
        .to(torch.bfloat16)
        .cpu()
    )

    got_scale = (
        scales_bf16[:rows]
        .cpu()
    )

    scale_exact = float(
        (
            expected_scale
            == got_scale
        )
        .float()
        .mean()
        .item()
    )

    # Reconstruct both paths from the independently decoded codes.
    expected_bf16 = (
        cb_cpu[
            expected_codes.long()
        ]
        * expected_scale.float()[:, None]
    ).to(torch.bfloat16)

    got_bf16 = (
        cb_cpu[
            got_codes.long()
        ]
        * got_scale.float()[:, None]
    ).to(torch.bfloat16)

    fake_exact = float(
        (
            expected_bf16
            == got_bf16
        )
        .float()
        .mean()
        .item()
    )

    return {
        "audit_rows": rows,
        "audit_device": str(dev),
        "code_exact_fraction": code_exact,
        "scale_exact_fraction": scale_exact,
        "fake_bf16_exact_fraction": fake_exact,
    }


def build_packed_from_official_source(
    transformer_dir: str | Path,
    output_dir: str | Path,
    *,
    seed: int = DEFAULT_SEED,
    bits: int = 4,
    device: str = "auto",
    row_chunk: int = 32,
    max_shard_gib: float = 0.75,
    overwrite: bool = False,
    model_revision: str | None = None,
) -> dict:
    src = Path(transformer_dir).resolve()
    out = Path(output_dir).resolve()
    if not src.is_dir():
        raise FileNotFoundError(src)
    wm = _weight_map(src)
    all_keys = set(wm)
    targets = sorted(k for k in all_keys if is_target_key(k))
    if len(all_keys) != 1303:
        raise RuntimeError(f"expected exactly 1303 Animate-2 transformer tensors, got {len(all_keys)}")
    if len(targets) != 480:
        raise RuntimeError(f"expected exactly 480 OrbitQuant target weights, got {len(targets)}")
    dims = sorted({int(_shape_of(src, wm, k)[1]) for k in targets})
    if dims != [5120, 13824]:
        raise RuntimeError(f"unexpected target input dimensions: {dims}")

    if out.exists() and any(out.iterdir()):
        if not overwrite:
            raise FileExistsError(out)
        shutil.rmtree(out)
    out.mkdir(parents=True, exist_ok=True)

    bank = RotationBank.build(dims, seed=seed, bits=bits)
    bank_path = out / "orbitquant_rotations.safetensors"
    bank.save(bank_path)

    if device == "auto":
        dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        dev = torch.device(device)
    shard_limit = int(float(max_shard_gib) * 2**30)
    buffer: dict[str, torch.Tensor] = {}
    buffer_bytes = 0
    shards: list[str] = []
    tensor_to_shard: dict[str, str] = {}
    target_info: dict[str, dict] = {}
    passthrough_info: dict[str, dict] = {}

    def flush():
        nonlocal buffer, buffer_bytes
        if not buffer:
            return
        name = f"orbitquant-runtime-{len(shards)+1:05d}.safetensors"
        save_file(buffer, str(out / name))
        for k in buffer:
            tensor_to_shard[k] = name
        shards.append(name)
        buffer = {}
        buffer_bytes = 0
        gc.collect()

    # Quantize target linears directly from the freshly-downloaded official weights.
    for i, key in enumerate(targets, 1):
        file = src / wm[key]
        with safe_open(str(file), framework="pt", device="cpu") as sf:
            w = sf.get_tensor(key)
        n, d = map(int, w.shape)
        packed_row, row_scale, stats = _quantize_rows_to_packed(
            w, bank.tensors[d], device=dev, row_chunk=row_chunk
        )
        audit = _verify_one_chunk(w, packed_row, row_scale, bank.tensors[d], rows=2, device=dev)
        if audit["code_exact_fraction"] != 1.0 or audit["scale_exact_fraction"] != 1.0 or audit["fake_bf16_exact_fraction"] != 1.0:
            raise RuntimeError(f"packed-source audit failed for {key}: {audit}")

        packed_runtime = packed_row.transpose(0, 1).contiguous()
        pkey = f"{key}.w4_packed_t"
        skey = f"{key}.row_scale_bf16"
        bytes_needed = packed_runtime.numel() + row_scale.numel() * row_scale.element_size()
        if buffer and buffer_bytes + bytes_needed > shard_limit:
            flush()
        buffer[pkey] = packed_runtime
        buffer[skey] = row_scale.contiguous()
        buffer_bytes += bytes_needed
        target_info[key] = {
            "official_key": key,
            "shape": [n, d],
            "input_dim": d,
            "output_dim": n,
            "packed_tensor": pkey,
            "packed_layout": "K_half_by_N",
            "scale_tensor": skey,
            "row_scale_dtype": "bfloat16",
            "mode": "direct_orbitquant_from_official_bf16_source",
            "audit": audit,
            **stats,
        }
        if i <= 5 or i % 20 == 0 or i == len(targets):
            print(
                f"[OrbitQuant fresh W4] {i:3d}/{len(targets)} {key} {n}x{d} "
                f"mse={stats['mse_rotated_weight']:.4e}"
            )
        del w, packed_row, packed_runtime, row_scale
        gc.collect()
        if dev.type == "cuda":
            torch.cuda.empty_cache()
    flush()

    target_set = set(targets)
    # Copy all non-target tensors unchanged; this makes the packed transformer self-contained.
    for fname, keys in sorted(_group_by_file(wm).items()):
        print(f"[passthrough] {fname}")
        with safe_open(str(src / fname), framework="pt", device="cpu") as sf:
            for key in keys:
                if key in target_set:
                    continue
                tensor = sf.get_tensor(key)
                bytes_needed = int(tensor.numel() * tensor.element_size())
                if buffer and buffer_bytes + bytes_needed > shard_limit:
                    flush()
                buffer[key] = tensor.contiguous()
                buffer_bytes += bytes_needed
                passthrough_info[key] = {"official_key": key, "tensor": key}
        gc.collect()
    flush()

    if len(passthrough_info) != 823:
        raise RuntimeError(f"expected 823 passthrough tensors, got {len(passthrough_info)}")

    for key, info in target_info.items():
        info["shard"] = tensor_to_shard[info["packed_tensor"]]
    for key, info in passthrough_info.items():
        info["shard"] = tensor_to_shard[info["tensor"]]

    config = src / "config.json"
    if config.is_file():
        shutil.copy2(config, out / "config.json")
    index = sorted(src.glob("*.safetensors.index.json"))[0]
    shutil.copy2(index, out / "source_transformer_index.json")

    shard_sizes = {name: (out / name).stat().st_size for name in shards}
    total_bytes = sum(shard_sizes.values())
    manifest = {
        "format": "OrbitQuant_WanAnimate2_direct_source_packed_nonuniform_W4A4_v3",
        "model": MODEL_ID,
        "source_model_revision": model_revision,
        "source_transformer_dir": str(src),
        "source_transformer_tensor_count": len(all_keys),
        "weight_bits": 4,
        "activation_bits": 4,
        "target_count": 480,
        "passthrough_count": 823,
        "full_transformer_tensor_count": 1303,
        "target_keyspace": "WanAnimate2_official_block_wrapper",
        "runtime_model_keyspace": "Wan-Video/Wan-Animate-2_official",
        "weight_storage_layout": "K_half_by_N",
        "weight_row_scale_dtype": "bfloat16",
        "activation_scale_dtype": "float32_runtime",
        "rotation_seed": int(seed),
        "rotation_bank": bank_path.name,
        "lloyd_max_density": "exact f_d(t)=Gamma(d/2)/(sqrt(pi)Gamma((d-1)/2))*(1-t^2)^((d-3)/2)",
        "lloyd_max_note": "OrbitQuant paper specifies the objective but does not publish its random seed or solver initialization/tolerance; this package uses deterministic exact-density Lloyd-Max and records its seed.",
        "packed_bytes": total_bytes,
        "shards": shards,
        "shard_sizes": shard_sizes,
        "targets": target_info,
        "passthrough_tensors": passthrough_info,
        "semantics": {
            "weight": "offline RPBH -> row L2 norm -> exact-density Lloyd-Max W4 direction; uint4 codes + BF16 row norm",
            "activation": "online RPBH -> token L2 norm -> exact-density Lloyd-Max A4; uint4 codes + FP32 runtime norm",
            "gemm": "centroids and scales dequantized to BF16 inside Triton K tile; FP32 accumulation; BF16 output",
            "non_target": "copied byte-for-byte tensor values from official BF16 transformer",
        },
    }
    manifest_path = out / "packed_manifest.json"
    manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True))
    manifest["manifest_sha256"] = hashlib.sha256(manifest_path.read_bytes()).hexdigest()
    return manifest


def _shape_of(root: Path, wm: dict[str, str], key: str) -> tuple[int, ...]:
    with safe_open(str(root / wm[key]), framework="pt", device="cpu") as sf:
        return tuple(int(x) for x in sf.get_slice(key).get_shape())