| |
| """用 NumPy 重算一整層分頁 MoE，並與 sdq-chat 的輸出比對。 |
| |
| `sdq-chat` 在 SDQ_MOE_DUMP=<prefix> SDQ_MOE_DUMP_LAYER=<il> 時，會從自訂 MoE op |
| 內寫出： |
| int32 meta[6] = { n_embd, n_tok, n_expert_used, n_expert, 1, 2 } |
| float cur [n_tok][n_embd] (FFN 輸入) |
| int32 ids [n_tok][n_expert_used] (被路由到的 expert) |
| float wts [n_tok][n_expert_used] (路由權重) |
| float out [n_tok][n_embd] (SDQ MoE 輸出) |
| |
| 本腳本從 GGUF 讀出同一批 expert 權重，以 fp32 重算 SiLU-gated expert 前向並比對， |
| 一次檢查 pager + 自訂 kernel + 權重縮放整條路徑。 |
| """ |
| import importlib.util |
| import os |
| import struct |
| import sys |
|
|
| import numpy as np |
|
|
| _HERE = os.path.dirname(os.path.abspath(__file__)) |
| spec = importlib.util.spec_from_file_location('vl', os.path.join(_HERE, 'verify_epq.py')) |
| vl = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(vl) |
|
|
| _ARCH = 'openpangu-v2' |
|
|
|
|
| def dequant_mat(raw, n_rows, ne0, tt): |
| nrow = vl.row_bytes(ne0, tt) |
| return np.stack([vl.dequant(raw[r * nrow:(r + 1) * nrow], ne0, tt) for r in range(n_rows)]) |
|
|
|
|
| def main(): |
| dump = sys.argv[1] |
| il = int(sys.argv[2]) |
| kv, ts = vl.parse_headers() |
| n_embd = int(kv[f'{_ARCH}.embedding_length']) |
| n_ff = int(kv[f'{_ARCH}.expert_feed_forward_length']) |
| n_expert = int(kv[f'{_ARCH}.expert_count']) |
|
|
| raw = open(dump, 'rb').read() |
| meta = struct.unpack('<6i', raw[:24]) |
| n_tok, n_used = meta[1], meta[2] |
| off = 24 |
| cur = np.frombuffer(raw, dtype=np.float32, count=n_tok * n_embd, offset=off).reshape(n_tok, n_embd) |
| off += 4 * n_tok * n_embd |
| ids = np.frombuffer(raw, dtype=np.int32, count=n_tok * n_used, offset=off).reshape(n_tok, n_used) |
| off += 4 * n_tok * n_used |
| wts = np.frombuffer(raw, dtype=np.float32, count=n_tok * n_used, offset=off).reshape(n_tok, n_used) |
| off += 4 * n_tok * n_used |
| got = np.frombuffer(raw, dtype=np.float32, count=n_tok * n_embd, offset=off).reshape(n_tok, n_embd) |
| print(f'layer {il}: n_tok={n_tok} n_expert_used={n_used} n_expert={n_expert} ' |
| f'n_embd={n_embd} n_ff={n_ff}') |
|
|
| def seg(nm, ie): |
| si, off_t, dims, tt = ts[f'blk.{il}.{nm}'] |
| nrow = vl.row_bytes(dims[0], tt) |
| nb = nrow * dims[1] |
| with open(vl.SHARDS[si], 'rb') as f: |
| f.seek(off_t + ie * nb) |
| return f.read(nb), nrow, tt, dims[0] |
|
|
| ref = np.zeros((n_tok, n_embd), dtype=np.float32) |
| experts = sorted(set(int(e) for e in ids.reshape(-1))) |
| print(f'distinct experts: {len(experts)}') |
| for ie in experts: |
| g_raw, g_nrow, g_t, g_ne0 = seg('ffn_gate_exps.weight', ie) |
| u_raw, _, u_t, u_ne0 = seg('ffn_up_exps.weight', ie) |
| d_raw, _, d_t, d_ne0 = seg('ffn_down_exps.weight', ie) |
| Wg = dequant_mat(g_raw, n_ff, g_ne0, g_t) |
| Wu = dequant_mat(u_raw, n_ff, u_ne0, u_t) |
| Wd = dequant_mat(d_raw, n_embd, d_ne0, d_t) |
| rows = [t for t in range(n_tok) if ie in ids[t]] |
| x = cur[rows].T |
| w = np.array([wts[t][list(ids[t]).index(ie)] for t in rows], dtype=np.float32) |
| g = Wg @ x |
| h = (g / (1.0 + np.exp(-g))) * (Wu @ x) |
| y = Wd @ h |
| ref[rows] += (y * w).T |
| del Wg, Wu, Wd |
|
|
| denom = max(float(np.abs(ref).max()), 1e-9) |
| diff = np.abs(got - ref) |
| corr = float(np.corrcoef(got.reshape(-1), ref.reshape(-1))[0, 1]) |
| print(f'max|got-ref| = {diff.max():.6g} max|ref| = {denom:.6g} ' |
| f'max_rel = {diff.max() / denom:.6g} corr = {corr:.8f}') |
| ok = corr > 0.99 and diff.max() / denom < 0.05 |
| print('PASS' if ok else 'FAIL') |
| return 0 if ok else 1 |
|
|
|
|
| if __name__ == '__main__': |
| sys.exit(main()) |
|
|
