""" PC端模拟器: 验证 bitflow_moe.py 的推理逻辑与 PyTorch 模型一致 模拟 micro:bit 上的 int8 解量化 + GLA + MoE + 采样 全流程 """ import struct, math, random, sys # ========== 加载 int8 权重 ========== with open("bitflow_moe_weights.bin", "rb") as f: nt = struct.unpack('= 128: v -= 256 r[j] = v * s return r SMW = [] for e in range(NE): b0 = 2 + e * 12 SMW.append(dqS(b0)) # norm1.w SMW.append(dqS(b0 + 1)) # norm1.b SMW.append(dqS(b0 + 2)) # log_decay SMW.append(dqS(b0 + 8)) # norm2.w SMW.append(dqS(b0 + 9)) # norm2.b rf1b = dqS(RB + 1) rf2b = dqS(RB + 3) nfw = dqS(NB) nfb = dqS(NB + 1) # ========== 融合解量化+矩阵乘 ========== def dqmv(i, ni, no, x, out): s = SC[i]; o = OFF[i] for r in range(no): ss = 0.0; b = r * ni for c in range(ni): v = W[o + b + c] if v >= 128: v -= 256 ss += v * s * x[c] out[r] = ss def mvW(w, ni, no, x, out): for r in range(no): ss = 0.0; b = r * ni for c in range(ni): ss += w[b + c] * x[c] out[r] = ss def ln(w, b, x, o): m = sum(x) / D v = sum((x[i] - m) ** 2 for i in range(D)) / D iv = 1.0 / math.sqrt(v + 1e-5) for i in range(D): o[i] = w[i] * (x[i] - m) * iv + b[i] def e1(x): return x + 1.0 if x >= 0 else math.exp(x) def sg(x): if x >= 0: e = math.exp(-x); return 1.0 / (1.0 + e) e = math.exp(x); return e / (1.0 + e) def ge(x): c = 0.7978845608 return x * sg(2.0 * c * (x + 0.044715 * x * x * x)) def sm(l, t=0.3): m = -1e9 for i in range(VS): v = l[i] - bcnt[i] * 2.0 if v > m: m = v s = 0.0 for i in range(VS): v = l[i] - bcnt[i] * 2.0 e = math.exp((v - m) / t) bpr[i] = e; s += e r = random.random() * s; c = 0.0 for i in range(VS): c += bpr[i] if c >= r: return i return VS - 1 # ========== 状态 ========== S = [[0.0] * (D * D) for _ in range(NL)] z = [[0.0] * D for _ in range(NL)] bx = [0.0] * D; bh = [0.0] * D bq = [0.0] * D; bk = [0.0] * D bv = [0.0] * D; bgd = [0.0] * D bnum = [0.0] * D; by = [0.0] * D bout = [0.0] * D; bff = [0.0] * FD bpr = [0.0] * VS; bcnt = [0] * VS brh = [0.0] * RD; brl = [0.0] * NE blg = [0.0] * VS bsum = [0.0] * D; brx = [0.0] * D; bct = [0] # ========== GLA ========== def gs(e, layer, h): b0 = 2 + e * 12; dec = SMW[e * 5 + 2] dqmv(b0 + 3, D, D, h, bq) dqmv(b0 + 4, D, D, h, bk) dqmv(b0 + 5, D, D, h, bv) dqmv(b0 + 6, D, D, h, bgd) for i in range(D): bq[i] = e1(bq[i]); bk[i] = e1(bk[i]) for i in range(D): bgd[i] = sg(bgd[i]) * dec[i] sl = S[layer]; zl = z[layer] for i in range(D): si = i * D; gd = bgd[i]; ki = bk[i] for j in range(D): sl[si + j] = gd * sl[si + j] + ki * bv[j] for i in range(D): zl[i] = bgd[i] * zl[i] + bk[i] for j in range(D): bnum[j] = 0.0 for i in range(D): si = i * D; qi = bq[i] for j in range(D): bnum[j] += qi * sl[si + j] den = 1e-6 for i in range(D): den += bq[i] * zl[i] iv = 1.0 / den for j in range(D): by[j] = bnum[j] * iv dqmv(b0 + 7, D, D, by, bout) def bs(e, x, layer): b0 = 2 + e * 12 ln(SMW[e * 5], SMW[e * 5 + 1], x, bh) gs(e, layer, bh) for i in range(D): x[i] += bout[i] ln(SMW[e * 5 + 3], SMW[e * 5 + 4], x, bh) dqmv(b0 + 10, D, FD, bh, bff) for i in range(FD): bff[i] = ge(bff[i]) dqmv(b0 + 11, FD, D, bff, bout) for i in range(D): x[i] += bout[i] def route(x): dqmv(RB, D, RD, x, brh) for i in range(RD): brh[i] = max(0.0, brh[i] + rf1b[i]) dqmv(RB + 2, RD, NE, brh, brl) for i in range(NE): brl[i] += rf2b[i] bi = 0; bv2 = brl[0] for i in range(1, NE): if brl[i] > bv2: bv2 = brl[i]; bi = i return bi, brl[:] def fw(ctx, gen_flag=True): t = min(len(ctx) - 1, CT - 1); ti = ctx[-1] s0 = SC[0]; o0 = OFF[0]; s1 = SC[1]; o1 = OFF[1] for i in range(D): v0 = W[o0 + ti * D + i] if v0 >= 128: v0 -= 256 v1 = W[o1 + t * D + i] if v1 >= 128: v1 -= 256 bx[i] = v0 * s0 + v1 * s1 # 累加运行均值 (匹配训练时router的x.mean(dim=1)) for i in range(D): bsum[i] += bx[i] bct[0] += 1 for i in range(D): brx[i] = bsum[i] / bct[0] ei, gate = route(brx) for layer in range(NL): bs(ei, bx, layer) ln(nfw, nfb, bx, bh) if gen_flag: dqmv(HI, D, VS, bh, blg) return blg, ei, gate return None, ei, gate def rs(): for l in range(NL): for i in range(D * D): S[l][i] = 0.0 for i in range(D): z[l][i] = 0.0 for i in range(D): bsum[i] = 0.0 bct[0] = 0 def gen(ids, n): NL_ch = VS - 1 ids = list(ids) + [NL_ch] sl = len(ids); c = list(ids); rs() for i in range(VS): bcnt[i] = 0 # 种子处理 for i in range(len(c) - 1): fw(c[:i + 1], False) # 生成 expert_names = ["事实", "推理", "问答"] for _ in range(n): lg, ei, gate = fw(c) nx = sm(lg) bcnt[nx] += 1 c.append(nx) gen_ids = c[sl:] gen = "".join(I2CH.get(i, '?') for i in gen_ids) answer = gen.split("\n")[0] if "\n" in gen else gen return answer, ei # ========== 测试 ========== seeds = [ "你好", "你好吗", "谢谢", "再见", "你叫什么名字", "你喜欢什么", "你喜欢猫吗", "你好棒", "你是谁", "你会什么", "你开心吗", "早安", "今天好吗", "你在做什么", "猫好吗", "你喜欢什么颜色", "猫是什么", "猫在做什么", "猫吃什么", "天是什么颜色", "什么会飞", "什么会跑", "什么会游", "什么吃草", ] print("=" * 60) print("PC模拟器验证 (模拟micro:bit int8推理)") print("=" * 60) correct = 0; total = 0 for seed in seeds: ctx = [CH2I[c] for c in seed if c in CH2I] if not ctx: continue random.seed(42) # 固定随机种子便于复现 answer, ei = gen(ctx, 20) expert_names = ["事实", "推理", "问答"] print(f"Q: {seed:12s} [E{ei}={expert_names[ei]}]") print(f"A: {answer}") print() total += 1 # 简单验证: 答案非空且包含相关字符 if answer and len(answer) > 0: correct += 1 print(f"\n结果: {correct}/{total} 有效输出") # 内存估算 mem_w = len(W) mem_S = NL * D * D * 4 mem_z = NL * D * 4 mem_buf = (9 * D + FD + VS + VS + RD + NE + VS + D + D) * 4 # +bsum, brx mem_SMW = NE * 5 * D * 4 mem_small = (RD + NE + D + D) * 4 # rf1b, rf2b, nfw, nfb mem_lists = 3 * nt * 8 # LEN, SC, OFF as Python lists total_mem = mem_w + mem_S + mem_z + mem_buf + mem_SMW + mem_small + mem_lists print(f"\n内存估算 (micro:bit):") print(f" W (int8权重): {mem_w:>6d} B = {mem_w/1024:.1f}KB") print(f" S (GLA状态): {mem_S:>6d} B = {mem_S/1024:.1f}KB") print(f" z (归一化): {mem_z:>6d} B = {mem_z/1024:.1f}KB") print(f" 缓冲区: {mem_buf:>6d} B = {mem_buf/1024:.1f}KB") print(f" SMW (小权重): {mem_SMW:>6d} B = {mem_SMW/1024:.1f}KB") print(f" 路由/norm_f: {mem_small:>6d} B = {mem_small/1024:.1f}KB") print(f" LEN/SC/OFF: {mem_lists:>6d} B = {mem_lists/1024:.1f}KB") print(f" 总计: {total_mem:>6d} B = {total_mem/1024:.1f}KB") print(f" micro:bit可用: ~108KB (128KB-20KB解释器)") print(f" 余量: {108*1024 - total_mem:>6d} B = {(108*1024 - total_mem)/1024:.1f}KB")