File size: 8,109 Bytes
3fd1a35 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | #!/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()
|