307 lines
8.7 KiB
Python
307 lines
8.7 KiB
Python
"""
|
|
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('<H', f.read(2))[0]
|
|
LEN = list(struct.unpack('<%dH' % nt, f.read(2 * nt)))
|
|
SC = list(struct.unpack('<%df' % nt, f.read(4 * nt)))
|
|
W = f.read()
|
|
|
|
OFF = []
|
|
o = 0
|
|
for l in LEN:
|
|
OFF.append(o)
|
|
o += l
|
|
|
|
# ========== 参数 ==========
|
|
VS = 149; D = 40; CT = 48; FD = 80; NE = 3; NL = 2; RD = 32
|
|
RB = 2 + NE * 12 # 38
|
|
NB = RB + 4 # 42
|
|
HI = NB + 2 # 44
|
|
|
|
# ========== 词表 ==========
|
|
VOCAB_CHARS = (
|
|
"猫狗鸟鱼猪牛马鸡虫羊兔鸭鹅"
|
|
"坐跑飞游看吃睡走听来叫爬跳唱画读做"
|
|
"在上里到"
|
|
"的是好了"
|
|
"吗什么怎哪"
|
|
"你我他她们它"
|
|
"蓝绿红黄白"
|
|
"天地水火花草树木日月星云春夏秋冬"
|
|
"大小多高远快"
|
|
"有没会不"
|
|
"人头脚手眼腿"
|
|
"垫肉书子"
|
|
"一三只两"
|
|
"热冷落开"
|
|
"自然动植物质色样颜得都"
|
|
"你好吗谢谢再见早晚安静"
|
|
"很棒爱喜欢乐伤心气"
|
|
"名字叫谁和问答应"
|
|
"聊天话今天说"
|
|
"客用"
|
|
"bit"
|
|
"\n"
|
|
)
|
|
CH2I = {c: i for i, c in enumerate(VOCAB_CHARS)}
|
|
I2CH = {i: c for i, c in enumerate(VOCAB_CHARS)}
|
|
|
|
# ========== 预解量化 ==========
|
|
def dqS(i):
|
|
s = SC[i]; o = OFF[i]; n = LEN[i]; r = [0.0] * n
|
|
for j in range(n):
|
|
v = W[o + j]
|
|
if v >= 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")
|