| |
| """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) |
| |
| 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)) |
|
|
