File size: 4,512 Bytes
12eb3a5 | 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 | #!/usr/bin/env python3
"""用 epq-chat 的 SDQ_MOE_DUMP 檔，驗證「整層」MoE 的計算（真實路由輸入）。
這是最接近模型實際執行的驗證：不是隨機 activation，而是模型自己算出來的
cur / selected_experts / weights，然後檢查我們的輸出是否等於 NumPy 重算的結果。
"""
import argparse, json, struct, sys
from pathlib import Path
import numpy as np
sys.path.insert(0, str(Path(__file__).parent))
from moe_ref import quant_q8_K, read_expert
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", required=True)
ap.add_argument("--index", required=True)
ap.add_argument("--dump", required=True)
ap.add_argument("--tol", type=float, default=3e-3)
a = ap.parse_args()
b = open(a.dump, "rb").read()
n_embd, n_tok, n_used, n_expert, _, _ = struct.unpack_from("<6i", b, 0)
off = 24
cur = np.frombuffer(b, dtype=np.float32, count=n_embd * n_tok, offset=off).reshape(n_tok, n_embd)
off += 4 * n_embd * n_tok
ids = np.frombuffer(b, dtype=np.int32, count=n_used * n_tok, offset=off).reshape(n_tok, n_used)
off += 4 * n_used * n_tok
wts = np.frombuffer(b, dtype=np.float32, count=n_used * n_tok, offset=off).reshape(n_tok, n_used)
off += 4 * n_used * n_tok
got = np.frombuffer(b, dtype=np.float32, count=n_embd * n_tok, offset=off).reshape(n_tok, n_embd)
idx = json.loads(Path(a.index).read_text())
il = int(Path(a.dump).name.rsplit(".", 1)[1])
layer = idx["per_layer"][il]["experts"]
n_ff = idx["n_ff_exp"]
out = np.zeros((n_tok, n_embd), dtype=np.float32)
# Worst-case effect on one output element of a single Q8_K LSB flip in h.
# ggml accumulates the gate/up dots in float32 while this reference is exact,
# so the two can disagree by ~1e-6 on h. When that crosses a Q8_K rounding
# boundary one code moves by 1 and the layer output jumps by far more than
# float noise. See docs/CORRECTNESS.md: a 1-LSB flip alone produced 7.9e-4.
lsb = np.zeros(n_ff, dtype=np.float64)
with open(a.model, "rb") as fh:
cache = {}
for t in range(n_tok):
for k in range(n_used):
e = int(ids[t, k])
if e not in cache:
g, u, d = layer["ffn_gate_exps"], layer["ffn_up_exps"], layer["ffn_down_exps"]
Wg = read_expert(fh, g["file_offset"] + e * g["bytes_per_expert"], n_ff * n_embd, g["type"]).reshape(n_ff, n_embd)
Wu = read_expert(fh, u["file_offset"] + e * u["bytes_per_expert"], n_ff * n_embd, u["type"]).reshape(n_ff, n_embd)
Wd = read_expert(fh, d["file_offset"] + e * d["bytes_per_expert"], n_embd * n_ff, d["type"]).reshape(n_embd, n_ff)
cache[e] = (Wg, Wu, Wd)
Wg, Wu, Wd = cache[e]
xq = quant_q8_K(cur[t])
# float64 for the gate/up dots: this is the reference, so it should
# be the more accurate of the two.
go = Wg.astype(np.float64) @ xq.astype(np.float64)
uo = Wu.astype(np.float64) @ xq.astype(np.float64)
h = ((go / (1.0 + np.exp(-go))) * uo).astype(np.float32)
if t == 0 and k == 0:
ab = np.abs(h)
for bi in range(h.size // 256):
blk = ab[bi * 256:(bi + 1) * 256]
lsb[bi * 256:(bi + 1) * 256] = blk.max() / 127.0 if blk.max() > 0 else 0.0
out[t] += (Wd @ quant_q8_K(h)) * wts[t, k]
diff = np.abs(out - got)
scale = np.maximum(np.abs(out), 1e-2)
rel = (diff / scale).max()
# bound[r] = sum_j |Wd[r,j]| * d_h(block(j)) = worst case for one flipped code
bound = np.zeros((n_embd, n_ff), dtype=np.float64)
for t in range(n_tok):
for k in range(n_used):
Wd = cache[int(ids[t, k])][2]
bound += np.abs(Wd.astype(np.float64)) * lsb[None, :]
print(f"layer {il} n_tok={n_tok} n_used={n_used} ids[0]={ids[0].tolist()}")
print(f" NumPy out[0,:5]={out[0,:5]}")
print(f" SDQ out[0,:5]={got[0,:5]}")
print(f" max abs diff={diff.max():.6g} max rel={rel:.6g} 1-LSB-flip bound={bound.max():.6g}")
ok = rel <= a.tol
verdict = "PASS" if ok else ("PASS (quantization-boundary)" if diff.max() <= bound.max() else "FAIL")
print(" 結果：", verdict + (" ✅" if ok else (" ✅" if verdict != "FAIL" else " ❌")))
sys.exit(0 if ok or verdict != "FAIL" else 1)
main()
|