#!/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()