scripts/build_qwen3tts_int4_stack_reference.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

import numpy as np

from build_qwen3tts_int4_attention_reference import apply_rope, softmax
from build_qwen3tts_int4_matvec_reference import load_manifest, load_tensor
from build_qwen3tts_int4_mlp_reference import find_tensor, input_vector, rms_norm, silu


def load_layer(export_dir: Path, manifest: dict[str, Any], layer: int, prefix: str) -> tuple[dict[str, Any], dict[str, np.ndarray]]:
    names = {
        "input_norm": f"{prefix}.layers.{layer}.input_layernorm.weight",
        "post_attention_norm": f"{prefix}.layers.{layer}.post_attention_layernorm.weight",
        "q_proj": f"{prefix}.layers.{layer}.self_attn.q_proj.weight",
        "k_proj": f"{prefix}.layers.{layer}.self_attn.k_proj.weight",
        "v_proj": f"{prefix}.layers.{layer}.self_attn.v_proj.weight",
        "o_proj": f"{prefix}.layers.{layer}.self_attn.o_proj.weight",
        "q_norm": f"{prefix}.layers.{layer}.self_attn.q_norm.weight",
        "k_norm": f"{prefix}.layers.{layer}.self_attn.k_norm.weight",
        "gate": f"{prefix}.layers.{layer}.mlp.gate_proj.weight",
        "up": f"{prefix}.layers.{layer}.mlp.up_proj.weight",
        "down": f"{prefix}.layers.{layer}.mlp.down_proj.weight",
    }
    tensors: dict[str, Any] = {}
    weights: dict[str, np.ndarray] = {}
    for role, name in names.items():
        component, index, tensor = find_tensor(manifest, name)
        matrix, _ = load_tensor(export_dir, tensor)
        tensors[role] = {"component": component, "tensor_index": index, "tensor_name": name, "shape": tensor["shape"]}
        weights[role] = matrix.astype(np.float32)
    return tensors, weights


def load_named_tensor(export_dir: Path, manifest: dict[str, Any], name: str) -> tuple[dict[str, Any], np.ndarray]:
    component, index, tensor = find_tensor(manifest, name)
    matrix, _ = load_tensor(export_dir, tensor)
    return {"component": component, "tensor_index": index, "tensor_name": name, "shape": tensor["shape"]}, matrix.astype(np.float32)


def run_layer(hidden: np.ndarray, weights: dict[str, np.ndarray], heads: int, kv_heads: int, head_dim: int, eps: float, rope_theta: float) -> tuple[np.ndarray, dict[str, float]]:
    seq_len = hidden.shape[0]
    normalized = np.stack([rms_norm(row, weights["input_norm"], eps) for row in hidden]).astype(np.float32)
    q = (normalized @ weights["q_proj"].T).reshape(seq_len, heads, head_dim)
    k = (normalized @ weights["k_proj"].T).reshape(seq_len, kv_heads, head_dim)
    v = (normalized @ weights["v_proj"].T).reshape(seq_len, kv_heads, head_dim)
    q = np.stack([[rms_norm(q[pos, head], weights["q_norm"], eps) for head in range(heads)] for pos in range(seq_len)]).astype(np.float32)
    k = np.stack([[rms_norm(k[pos, head], weights["k_norm"], eps) for head in range(kv_heads)] for pos in range(seq_len)]).astype(np.float32)
    q = apply_rope(q, rope_theta)
    k = apply_rope(k, rope_theta)
    group = heads // kv_heads
    output_rows = np.empty_like(hidden)
    attention_checksums = []
    attention_output_checksums = []
    mlp_output_checksums = []
    for target in range(seq_len):
        context = np.empty((heads, head_dim), dtype=np.float32)
        for head in range(heads):
            kv_head = head // group
            scores = np.array([np.dot(q[target, head], k[pos, kv_head]) / np.sqrt(np.float32(head_dim)) for pos in range(target + 1)], dtype=np.float32)
            probs = softmax(scores)
            context[head] = np.sum(probs[:, None] * v[: target + 1, kv_head], axis=0, dtype=np.float32)
            attention_checksums.append(float(np.sum(probs, dtype=np.float32)))
        attention_output = weights["o_proj"] @ context.reshape(heads * head_dim)
        residual = hidden[target] + attention_output
        mlp_input = rms_norm(residual, weights["post_attention_norm"], eps)
        gate = weights["gate"] @ mlp_input
        up = weights["up"] @ mlp_input
        mlp_output = weights["down"] @ (silu(gate) * up)
        output_rows[target] = residual + mlp_output
        attention_output_checksums.append(float(np.sum(attention_output, dtype=np.float32)))
        mlp_output_checksums.append(float(np.sum(mlp_output, dtype=np.float32)))
    stats = {
        "attention_checksum": float(np.sum(attention_output_checksums, dtype=np.float32)),
        "mlp_checksum": float(np.sum(mlp_output_checksums, dtype=np.float32)),
        "attention_probability_checksum": float(np.sum(attention_checksums, dtype=np.float32)),
        "output_checksum": float(np.sum(output_rows, dtype=np.float32)),
    }
    return output_rows, stats


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--export-dir", type=Path, default=Path("exports/qwen3tts_0p6b_sft_partial_int4_weights_20260606"))
    parser.add_argument("--config", type=Path, default=Path("job_output/qwen3tts-0p6b-official-sft-ptbr-256-partial-last2heads-20260606/output_b1_epoch1_lr2e6/checkpoint-epoch-0/config.json"))
    parser.add_argument("--layers", type=int, default=2)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--sequence-length", type=int, default=3)
    parser.add_argument("--prefix", default="talker.model")
    parser.add_argument("--final-norm-name")
    parser.add_argument("--head-prefix")
    parser.add_argument("--head-count", type=int, default=0)
    parser.add_argument("--out", type=Path, default=Path("exports/qwen3tts_0p6b_sft_partial_int4_stack2_reference_20260606.json"))
    args = parser.parse_args()
    manifest = load_manifest(args.export_dir / "manifest.json")
    config = json.loads(args.config.read_text())["talker_config"]
    hidden_size = int(config["hidden_size"])
    head_dim = int(config["head_dim"])
    heads = int(config["num_attention_heads"])
    kv_heads = int(config["num_key_value_heads"])
    eps = float(config["rms_norm_eps"])
    rope_theta = float(config["rope_theta"])
    hidden = np.stack([input_vector(hidden_size, args.seed + pos) for pos in range(args.sequence_length)]).astype(np.float32)
    layer_refs = []
    layer_stats = []
    for layer in range(args.layers):
        tensors, weights = load_layer(args.export_dir, manifest, layer, args.prefix)
        hidden, stats = run_layer(hidden, weights, heads, kv_heads, head_dim, eps, rope_theta)
        layer_refs.append({"layer": layer, "tensors": tensors})
        layer_stats.append({"layer": layer, **stats})
    target = args.sequence_length - 1
    output = hidden[target]
    data = {
        "format": "qwen3tts_int4_stack_reference_v1",
        "package": args.export_dir.name,
        "prefix": args.prefix,
        "layer_count": args.layers,
        "seed": args.seed,
        "sequence_length": args.sequence_length,
        "target_index": target,
        "eps": eps,
        "hidden_size": hidden_size,
        "heads": heads,
        "kv_heads": kv_heads,
        "head_dim": head_dim,
        "rope_theta": rope_theta,
        "layers": layer_refs,
        "layer_stats": layer_stats,
        "sequence_checksum": float(np.sum(hidden, dtype=np.float32)),
        "checksum": float(np.sum(output, dtype=np.float32)),
        "max_abs": float(np.max(np.abs(output))),
        "values": [float(value) for value in output.tolist()],
    }
    if args.final_norm_name and args.head_prefix and args.head_count > 0:
        norm_ref, norm_weight = load_named_tensor(args.export_dir, manifest, args.final_norm_name)
        normalized = rms_norm(output, norm_weight, eps)
        head_refs = []
        for head in range(args.head_count):
            head_ref, head_weight = load_named_tensor(args.export_dir, manifest, f"{args.head_prefix}.{head}.weight")
            logits = head_weight @ normalized
            head_refs.append({
                "head": head,
                "tensor": head_ref,
                "checksum": float(np.sum(logits, dtype=np.float32)),
                "max_abs": float(np.max(np.abs(logits))),
                "argmax": int(np.argmax(logits)),
                "argmax_value": float(np.max(logits)),
                "values": [float(value) for value in logits.tolist()],
            })
        data["logits"] = {
            "norm": norm_ref,
            "heads": head_refs,
            "normalized_checksum": float(np.sum(normalized, dtype=np.float32)),
            "normalized_max_abs": float(np.max(np.abs(normalized))),
        }
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n")
    print(json.dumps({"out": str(args.out), "checksum": data["checksum"], "max_abs": data["max_abs"]}, ensure_ascii=False))


if __name__ == "__main__":
    main()