File size: 3,103 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 | #!/usr/bin/env python3
"""Export real expert matrices for a bounded group32 NPU compatibility probe."""
import argparse
import json
from pathlib import Path
import numpy as np
import torch
from safetensors import safe_open
from official_int4_reference import unpack4
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--official", type=Path, required=True)
p.add_argument("--source", type=Path, required=True)
p.add_argument("--output", type=Path, required=True)
a = p.parse_args()
a.output.mkdir(parents=True, exist_ok=True)
torch.set_num_threads(4)
indices = {str(root): json.loads((root / "model.safetensors.index.json").read_text())["weight_map"]
for root in (a.official, a.source)}
def get(root, name):
with safe_open(root / indices[str(root)][name], framework="pt") as f:
return f.get_tensor(name)
records = []
for proj in ("gate", "up", "down"):
name = f"model.layers.1.mlp.experts.0.{proj}_proj.weight"
original = get(a.source, name).float()
n, k = original.shape
q = unpack4(get(a.official, name+"_packed"), (n,k))
s = get(a.official, name+"_scale").float()
for kind in ("official_group32", "old_channel"):
if kind == "old_channel":
s = torch.maximum(original.amax(1).clamp_min(0)/7, (-original.amin(1)).clamp_min(0)/8).unsqueeze(1)
q = torch.round(original/s).clamp(-8,7).to(torch.int8)
group = 32 if kind == "official_group32" else k
case = a.output / (proj+"_"+kind)
case.mkdir(exist_ok=True)
np.array([k,n,group], dtype=np.int32).tofile(case / "shape.bin")
q.numpy().tofile(case / "codes.bin")
s.numpy().tofile(case / "scales.bin")
weight = (q.reshape(n,-1,group).float()*s[:,:,None]).reshape(n,k)
for rows in (1,128):
rng = torch.Generator().manual_seed(314159+rows)
x = torch.randn(rows,k,generator=rng)
# Explicit independent integer oracle for the added groupwise A8.
xg = x.reshape(rows,-1,group)
xs = xg.abs().amax(-1)/127
xq = torch.round(xg/xs[:,:,None]).clamp(-127,127)
oracle = torch.zeros(rows,n)
for g in range(k//group):
oracle += (xq[:,g,:] @ q[:,g*group:(g+1)*group].float().T)*xs[:,g,None]*s[:,g]
x.numpy().tofile(case / f"input-{rows}.bin")
oracle.numpy().tofile(case / f"oracle-{rows}.bin")
official_out = x @ weight.T
original_out = x @ original.T
records.append({"case":case.name,"rows":rows,"group":group,"k":k,"n":n,
"weight_only_output_relative_rmse":float(torch.linalg.vector_norm(official_out-original_out)/torch.linalg.vector_norm(original_out)),
"extra_a8_output_relative_rmse":float(torch.linalg.vector_norm(oracle-official_out)/torch.linalg.vector_norm(official_out))})
(a.output / "manifest.json").write_text(json.dumps({"description":"one actual layer1 expert, synthetic Gaussian activations; not full model or task accuracy", "cases": records},indent=2)+"\n")
print(json.dumps(records,indent=2))
|