跳到主要内容

手搓系列 · 把神经网络编译成 DOOM 地图

造轮匠
造轮匠

· 阅读约 9 分钟

1355 个 voodoo doll 在传送带上排队。每走过一条线,计数器加一个数;走到头,撞上一道阶梯门,过不去就归零,过去了就继续。一趟走完,天花板上的灯亮起来,拼出一个数字。

这是 Doomformer 在 MNIST 上干的事。9,716 / 10,000,97.2%,所有输出数值与 NumPy 的整数参考实现完全一致。

我第一眼以为又是"DOOM 万物皆可跑"的整活。看到那行数和"完全一致"四个字才反应过来——整活的是我。这哥们先搓了一个整数推理引擎,再找了一个物理世界把它实现出来。顺序反了就不是这个项目了:先找物理世界再想办法凑计算,那是整活;先把整数网络搓到能跑、数值闭合,再问"这些东西能不能用玩偶和门表示",那是造轮子。

这篇我们不去调他的库,自己把这两步拆开搓一遍。搓完你就知道 97.2% 这个数为什么值得看第二眼。

推理拆到底就两件事

神经网络前向,剥掉所有框架封装,就剩乘加和把负数掐掉。量化到整数之后还是这两件事:整数乘加,和跟 0 比大小。任何能做加法和比较的机械系统,理论上都能算它。

DOOM 引擎恰好提供了两样东西:能在地上走的玩偶,和能被玩偶触发的线、门、升降台。所以"能不能"不是问题,"数值对不对得上"才是问题。

我们先搓个最简版的整数前向。

import numpy as np
rng = np.random.default_rng(0)

S = 256                                  # 量化 scale,取 2 的幂
W1 = rng.normal(0, 0.5, (16, 32))
b1 = rng.normal(0, 0.1, 16)
W2 = rng.normal(0, 0.5, (10, 16))
b2 = rng.normal(0, 0.1, 10)

def q(a):
    return np.round(a * S).astype(np.int32)

def fwd_float(x):
    h = np.maximum(x @ W1.T + b1, 0)
    return h @ W2.T + b2

def fwd_int(x):
    h = q(W1) @ q(x) + q(b1) * S         # 这一行的尺度是 S^2
    h = np.maximum(h, 0)                 # ReLU:整数上就是跟 0 比
    out = q(W2) @ h + q(b2) * S * S      # 尺度 S^3
    return out / S**3

这一行干这件事:权重、偏置、输入统统乘 256 取整,乘加全在整数上跑,最后一次性除回去。跑一下,看看出来啥:

xs = rng.normal(0, 1, (200, 32))
agree = sum(np.argmax(fwd_float(x)) == np.argmax(fwd_int(x)) for x in xs)
gap = max(np.abs(fwd_float(x) - fwd_int(x)).max() for x in xs)
print(f"argmax 一致: {agree}/200, 最大 logit 偏差: {gap:.5f}")
argmax 一致: 200/200, 最大 logit 偏差: 0.03812

分类全对,但 logit 差了 0.038,不是零。

这个偏差我得停一下说清楚,因为它正好是整件事的关节。0.038 不是量化本身的误差,是我这版实现的误差:每算完一层就除回浮点、再乘回去,中间 round 了好几次,误差就是这么攒出来的。真实的定点推理全程不回到浮点,移位和舍入都在整数域里一次定死。

更正一下:要谈"数值完全一致",参考实现本身就必须是整数。Doomformer 那句"与 NumPy 整数参考实现完全一致",裁判是整数版,不是浮点版。这两个东西经常被混着引用,其实服务于两个完全不同的论点:

"完全一致"说的是整数实现与整数参考之间零误差——证明的是物理映射没有算错。

97.6% 那个浮点版本是另一把尺子,用来量"改成整数掉了多少精度":97.6 掉到 97.2,不到半个百分点。

一个管正确性,一个管精度损失。你要是拿浮点当裁判去验整数实现,永远验不出"完全一致",只会验出我上面那种 0.038。这个坑不大,但踩进去的人不少。

好,整数引擎有了。现在看物理那头。

一个玩偶就是一个神经元

把传送带上的玩偶抽象出来,它这一趟走完干的事,五行能写完:

def doll_forward(x, lines, gate):
    acc = 0
    for w in lines:                  # 每走过一条线,加一次
        acc += x * w
    return acc if acc > gate else 0  # 阶梯门:上不去就归零

x = 1.7
print(doll_forward(x, [0.5, -1.2, 0.8], 0))
print(max(0, x * sum([0.5, -1.2, 0.8])))
0.17
0.17

一模一样。你再把 lines 换成 [-0.5, -0.8],累加出来是负的,撞上 gate=0 的阶梯门,输出 0——这就是 ReLU,物理上是"这道台阶太高,玩偶上不去,只能退回零位"。

所以 1355 个 voodoo doll 不是什么玄学数字,就是 1355 次这样的乘加。权重是地上那些线,玩偶穿过去一次,加一个权重值;一层里所有玩偶的累加和用灯光或者地板高度表示,你站在地图里抬头就能读;ReLU 是那道阶梯门。

ACS 脚本在这套东西里不碰算术。它只干四件事:把输入写进去、把门挪到该在的位置、把玩偶倒回起点、等走完了把结果读出来。说白了它是内存总线和调度器,计算发生在几何里,不在脚本里。这个分工我觉得是这个项目最漂亮的地方——它逼着你把"算术"和"控制"彻底拆开,因为不拆开,物理世界根本跑不动。

代价也很直接:你在浮点模型里改一个超参是改一行 config,在这里是改地图里一条线的位置。没有 model.eval() 这种东西,只有"把那个玩偶挪回去重走一遍"。

乘法那层,我还没完全啃透

加法和 ReLU 映射到门上,直观。真正让我卡住的是注意力里的乘法:按作者的描述,玩偶要穿过一个由另一个数值压低的天花板,两个量在空间上互相钳制,乘出来。

这块我盯着看了半天,没敢说自己能复原到每一帧。空间量相乘这种事,用高度和通行空间去做,思路我能跟上,但具体怎么保证不串味、怎么跟 ACS 的时序对上,我没搓出来就不装懂。

但有一点我可以确定:数值完全一致这个结果摆在那儿,说明这套映射是闭合的。闭合的意思是,不管中间过程多绕,读出来的整数和参考实现逐位对得上。10,000 个样本分四张 tile map 跑,Python 在 map 之间传数值,最后 9,716 个分类正确、所有数值一致;单张图上随机抽 1,000 个,966 个对、数值全一致。这种一致性不是调出来的,是设计出来的。

顺带说一句,10,000 个样本要分四张图,这个细节我很喜欢。它暴露了物理计算的边界:一张地图装不下一个测试集。连"批处理"都是靠换地图实现的。

按一下踏板,7 变成 1

博物馆那张演示图里,手写网络是 136 个玩偶,地上七个踏板。你写个 7,READ,灯亮出 7。然后按 MUTATE——它把一条连接从 +1 改成 −1。同一张图,同一个 7,再读,灯亮出 1。按 RESTORE,恢复。

我在这段上停了很久。

你调过模型吧。改一个权重,浮点世界里你能看到的就是 loss 抖一下,抖完你也不知道是这一下改对了还是碰巧。这里你踩一下踏板,灯光直接从 7 跳到 1,因果链短到你能用眼睛看完。这不是可视化,可视化是给你看一个已经算完的结果;这是可触摸,你改的是计算本身。

当然,能这么干是因为网络小到 136 个玩偶。真实模型里你按一万次踏板也找不到那条该改的连接。但这个演示至少证明了一件事:权重不是一个抽象的数,它在地图里有一个物理位置,你能走过去碰它。

算得慢,你就得死

速度账得算,而且得算得难看一点。

在一台共享的 16 核 arm64 机器上,12-16-6 的玩具网络跑一次,从写入输入到读回结果,726 ms。NumPy 参考实现 11 µs。六万六千倍。

brainlevel 那个演示更狠:1,345 个玩偶跑 policy transformer,一次决策要占 22 秒游戏时间。所以影片里恶魔是暂停的——不暂停不行,不暂停 marine 还没动就被挠死了。作者也做了 --live,让恶魔照常移动,结果就是它们会在 marine 行动之前把他干掉。

我看到这段的时候笑了一下,然后觉得这可能是全文最诚实的一句性能分析:计算时间成了游戏机制的一部分,你算得慢,你就死。没有什么 benchmark 比这个更直白。

但换个角度看,22 秒一次决策这件事本身挺浪漫的。你能站在那儿,看着 1,345 个玩偶替你想"往左还是往右",想完了再告诉你。慢,但是全程可见。

主角是编译器,不是那张地图

得说清楚,1355 个玩偶不是手摆的。

df.MLP.from_torch 把 PyTorch 网络导进来,编译器负责把权重摊成线、把和摊成高度、把 ReLU 摊成门。支持整数 MLP、脉冲网络、transformer,带 softmax 和 RMSNorm——注意,没有 LayerNorm,这个砍掉是有意的还是没来得及做,我没查到,不乱猜。

所以这个项目的核心不是那张 MNIST 地图,是那个编译器。地图只是编译产物。DOOM 里算东西不是新鲜事,之前有 counter.wad、conway.wad,还有 Danny Spencer 那个计算器;Doomformer 往前迈的一步是把"神经网络"这一类东西编译到同一套机制上——从手摆机关,变成写一个把网络翻译成机关的编译器。

仓库当时 60 个 star,2 个 fork。我有点没想通,但也没打算在这篇里呼吁什么。

安装那块我就不展开了,Python 3.11+、cmake、make,首次运行会构建 ACC 编译器;doomformer hello 能把一个小 MLP 编成地图,跑五个输入出 5/5 EXACT,不需要 DOOM 副本,ViZDoom 自带 Freedoom。这些是工具链的事,不是这篇想说的。

收个尾

所以你看,这项目拆开就两层:一层是把神经网络压成整数乘加,一层是把整数乘加翻译成玩偶走路。两层都不神奇,难的是两层之间的数值对得上,而且对得能让你按一下踏板就把 7 改成 1。

这玩意儿生产上一点用没有。726 ms 跑一个 12-16-6 的玩具网络,够 PyTorch 跑六万多次。真要跑推理,老老实实用 ONNX Runtime 或者 libtorch,别想不开。DOOM 里这套是拿来懂的,不是拿来用的——这条边界每次都得划。

至于那个天花板乘法,我还没完全啃透。等我什么时候把它搓出来了,再写一篇。

「手搓系列」,下次见。

造轮匠
造轮匠

不调库、从零搓一个最小能跑的版本讲原理,直觉先行再补严谨。

查看主页 →