#!/usr/bin/env python3 """Evaluate released INT4 codes/scales without a second weight quantization. Compute is the same CUDA BF16 teacher-forcing reference as previous diagnostics, not native compressed INT4 execution and not an RK3588 speed measurement. """ import argparse import hashlib import importlib.metadata import importlib.util import json from contextlib import ExitStack from pathlib import Path import time import torch from safetensors import safe_open from transformers import AutoTokenizer from quant_loss_reference import load_model, evaluate REVISION = "65a6d1d71e01f73ba01e572992bbd69ea92c865f" def unpack4(packed, shape): assert packed.dtype == torch.int32 and len(shape) == 2 shifts = torch.arange(0, 32, 4, dtype=torch.int32, device=packed.device) return (((packed.unsqueeze(-1) >> shifts) & 15).reshape(shape[0], -1) [:, :shape[1]] - 8).to(torch.int8) def verify_unpack(): # Load the dependency's standalone helper without importing optional runtimes. dist = importlib.metadata.distribution("compressed-tensors") path = dist.locate_file("compressed_tensors/compressors/pack_quantized/helpers.py") spec = importlib.util.spec_from_file_location("ct_pack_helper", path) helper = importlib.util.module_from_spec(spec) spec.loader.exec_module(helper) codes = torch.arange(-8, 8, dtype=torch.int8).repeat(8).reshape(4, 32) packed = helper.pack_to_int32(codes, 4) assert torch.equal(unpack4(packed, codes.shape), codes) assert torch.equal(helper.unpack_from_int32(packed, 4, codes.shape), codes) return {"package_version": dist.version, "helper_sha256": hashlib.sha256(path.read_bytes()).hexdigest()}, helper def replace_weights(model, official, helper): index = json.loads((official / "model.safetensors.index.json").read_text())["weight_map"] cfg = json.loads((official / "config.json").read_text())["quantization_config"] weights = cfg["config_groups"]["group_0"]["weights"] assert cfg["format"] == "pack-quantized" assert weights["group_size"] == 32 and weights["num_bits"] == 4 and weights["symmetric"] assert not any(k.endswith(("weight_zero_point", "weight_g_idx")) for k in index) state = model.state_dict(keep_vars=True) projected = {k.removesuffix("_packed") if k.endswith("weight_packed") else k for k in index if not k.endswith(("weight_scale", "weight_shape"))} if projected != set(state): raise RuntimeError(f"tensor coverage mismatch: missing={set(state)-projected}, extra={projected-set(state)}") stats = {"revision": REVISION, "quantized_matrices": 0, "quantized_elements": 0, "raw_tensors": 0, "raw_elements": 0, "raw_differences": [], "raw_compute_casts": [], "original_weight_squared_sum": 0., "weight_error_squared_sum": 0., "bf16_reconstruction_rounding_squared_sum": 0., "scale_fp16_not_exact": 0, "compute": "original BF16 implementation; released codes * BF16 group32 scales, cast to BF16; no A8"} with ExitStack() as stack, torch.no_grad(): handles = {s: stack.enter_context(safe_open(official / s, framework="pt", device="cpu")) for s in sorted(set(index.values()))} def get(name): return handles[index[name]].get_tensor(name) for i, (name, param) in enumerate(state.items()): if name in index: value = get(name).to(param.device) if value.dtype != param.dtype: # Previous BF16 reference casts non-parameter router-bias buffers. # Match that reference arithmetic rather than silently changing it. if not name.endswith(".mlp.gate.expert_bias"): raise RuntimeError(f"raw dtype mismatch {name}: {value.dtype} vs {param.dtype}") stats["raw_compute_casts"].append({"name": name, "source": str(value.dtype), "compute": str(param.dtype)}) value = value.to(param.dtype) if not torch.equal(value, param): stats["raw_differences"].append(name) param.copy_(value) stats["raw_tensors"] += 1 stats["raw_elements"] += param.numel() else: assert ".mlp.experts." in name, name shape = tuple(get(name + "_shape").tolist()) assert shape == tuple(param.shape) packed = get(name + "_packed") if stats["quantized_matrices"] == 0: assert torch.equal(unpack4(packed, shape), helper.unpack_from_int32(packed, 4, shape)) codes = unpack4(packed.to(param.device), shape) scale = get(name + "_scale").to(param.device) assert scale.shape == (shape[0], shape[1] // 32) assert torch.isfinite(scale).all() and (scale >= 0).all() stats["scale_fp16_not_exact"] += int((scale.half().bfloat16() != scale).sum()) exact = (codes.reshape(shape[0], -1, 32).float() * scale.float().unsqueeze(-1)).reshape(shape) original = param.float() stats["original_weight_squared_sum"] += original.double().square().sum().item() stats["weight_error_squared_sum"] += (original-exact).double().square().sum().item() stats["bf16_reconstruction_rounding_squared_sum"] += (exact-exact.bfloat16().float()).double().square().sum().item() param.copy_(exact) stats["quantized_matrices"] += 1 stats["quantized_elements"] += param.numel() if (i+1) % 1000 == 0: print(json.dumps({"loaded_tensors": i+1, "total": len(state)}), flush=True) stats["expert_relative_weight_rmse"] = (stats["weight_error_squared_sum"] / stats["original_weight_squared_sum"]) ** .5 return stats def main(): p = argparse.ArgumentParser(description=__doc__) p.add_argument("--source", type=Path, required=True) p.add_argument("--official", type=Path, required=True) p.add_argument("--suite", type=Path, required=True) p.add_argument("--output-dir", type=Path, required=True) args = p.parse_args() args.output_dir.mkdir(parents=True, exist_ok=True) torch.set_num_threads(4) torch.backends.cuda.matmul.allow_tf32 = False checks, helper = verify_unpack() suite = json.loads(args.suite.read_text()) tok = AutoTokenizer.from_pretrained(args.official, trust_remote_code=True, local_files_only=True) for case in suite["cases"]: assert tok.encode(case["prompt_text"], add_special_tokens=False) == case["prompt_ids"] assert tok.encode(case["target_text"], add_special_tokens=False) == case["teacher_ids"] source_config = json.loads((args.source / "config.json").read_text()) official_config = json.loads((args.official / "config.json").read_text()) differences = {k: [source_config.get(k), official_config.get(k)] for k in set(source_config) | set(official_config) if source_config.get(k) != official_config.get(k)} allowed = {"quantization_config", "num_nextn_predict_layers", "model_type", "auto_map", "torch_dtype", "transformers_version", "_name_or_path"} if set(differences) - allowed: raise RuntimeError(f"unreviewed config differences: {differences}") print(json.dumps({"config_differences": differences, "unpack_check": checks}), flush=True) start = time.monotonic() model = load_model(args.source) stats = replace_weights(model, args.official, helper) stats.update({"unpack_check": checks, "config_differences": differences, "load_seconds": time.monotonic()-start, "suite_sha256": hashlib.sha256(args.suite.read_bytes()).hexdigest()}) (args.output_dir / "weight-audit.json").write_text(json.dumps(stats, indent=2) + "\n") evaluate(model, suite["cases"], args.output_dir, "official_int4_bf16") print("PASS: official INT4 weight reconstruction and paired likelihood evaluation", flush=True) if __name__ == "__main__": main()