miutti's picture
Source: kernel, desktop app, tools (snapshot of the GitHub repo)
df41178 verified
Raw History Blame Contribute Delete
3 kB
import os
"""実モデルの重み行列で、1ビット未満まで含む各コーデックを比較する。
使い方:
/usr/bin/python3 bench.py # 既定のテンソル1枚
/usr/bin/python3 bench.py --tensor blk.0.ffn_gate.weight
"""
import argparse, json, os, sys, time
ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, os.path.join(ROOT, "lib"))
import numpy as np
import gguf, wcodec as C
BLOB = (os.environ.get("MODEL_BLOB") or "/path/to/localai/ollama-models/blobs/"
"sha256-81fb60c7daa80fc1123380b98970b320ae233409f0f71a72ed7b9b0d62f40490")
def build_codecs():
return [
("基準8bit", lambda W: C.rtn(W, 8, 128)),
("基準4bit", lambda W: C.rtn(W, 4, 128)),
("基準2bit", lambda W: C.rtn(W, 2, 128)),
("1bit級", lambda W: C.ternary(W, 128)),
("1bit未満", lambda W: C.pvq(W, k=4, cb_bits=4, G=128)), # 1.0bit
("1bit未満", lambda W: C.pvq(W, k=8, cb_bits=6, G=128)), # 0.75bit
("1bit未満", lambda W: C.pvq(W, k=8, cb_bits=4, G=128)), # 0.5bit
("1bit未満", lambda W: C.pvq(W, k=8, cb_bits=4, G=256)), # 0.56→0.5bit
("1bit未満", lambda W: C.pvq(W, k=16, cb_bits=8, G=256)), # 0.5bit
("1bit未満", lambda W: C.pvq(W, k=8, cb_bits=4, G=128, random_codebook=True)),
("ハイブリッド", lambda W: C.hybrid(W, 0.01, 8, G=128)),
("ハイブリッド", lambda W: C.hybrid(W, 0.03, 8, G=128)),
("ハイブリッド", lambda W: C.hybrid(
W, 0.03, 8, G=128, sub=lambda M: C.pvq(M, k=16, cb_bits=8, G=256))),
]
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--tensor", default="v.blk.0.mlp.linear_fc1.weight")
ap.add_argument("--blob", default=BLOB)
ap.add_argument("--out", default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "results"))
a = ap.parse_args()
r = gguf.Reader(a.blob)
W, tname = r.load(a.tensor)
W = W.astype(np.float32)
print(f"対象: {a.tensor} {W.shape} 元の格納形式={tname} std={W.std():.5f}")
print(f"{'区分':<12}{'方式':<44}{'bit/重み':>9}{'重み誤差':>10}{'出力誤差':>10}{'秒':>7}")
print("-" * 94)
rows = []
for cat, fn in build_codecs():
t0 = time.time()
Wh, bpw, label = fn(W)
rw, ry = C.evaluate(W, Wh)
dt = time.time() - t0
print(f"{cat:<12}{label:<44}{bpw:>9.3f}{rw:>10.4f}{ry:>10.4f}{dt:>7.1f}")
rows.append(dict(区分=cat, 方式=label, bpw=round(bpw, 4),
重み誤差=round(float(rw), 5), 出力誤差=round(float(ry), 5)))
os.makedirs(a.out, exist_ok=True)
p = os.path.join(a.out, a.tensor.replace("/", "_") + ".json")
with open(p, "w") as f:
json.dump(dict(tensor=a.tensor, shape=list(W.shape), stored_as=tname, rows=rows),
f, ensure_ascii=False, indent=2)
print(f"\n保存: {p}")
if __name__ == "__main__":
main()