七块 ESP32-S3 拼出跑 BitNet 模型的集群

七块 ESP32-S3 可以凑出一个跑语言模型的集群。一块做主控,负责分词和嵌入,另外六块当计算节点,每块跑四层 Transformer,加起来 24 层,节点之间用 SPI 菊花链串起来。基础模型取 Qwen2-0.5B,权重压成 1.58-bit 三值,模型和微控制器之间那道通常跨不过去的墙,被硬掰开一条缝。整个项目开源,代码、固件、Python 工具链和接线说明都在仓库里。

它能跑通,但别指望速度。每个节点推理要 1.3 秒,生成一个 token 要全链走完,加上词表裁剪和量化损失,它更适合当教学演示,看一个量化语言模型的完整流程如何在单片机上分层落地。

基础模型选 Qwen2-0.5B,24 层,隐藏维度 896,前馈 4864,用分组注意力,14 个查询头对 2 个键值头,头维度 64,KV 维度 128,上下文 512 token。分片的依据就是这 24 层,六块计算板各分四层,边界画在层与层之间,激活值沿 SPI 链往下传。

模型侧要做的三步

  1. 用 crop_tokenizer.py 把词表裁到 32000,用裁剪脚本同步裁嵌入层,再更新配置文件。
  2. 跑 1.58-bit 的量化感知微调脚本,让权重适应三值表示,再用脚本把嵌入层打包成 INT4。
  3. 用打包脚本生成固件能烧录的二进制文件。

三值权重是 BitNet 的做法,每个权重只取负一、零、正一,一个字节装四个。嵌入层和输出头保留 INT4,最后的归一化层用 FP16。裁完之后每层约 3.82 MB,一个节点放四层约 15.3 MB,KV cache 放 PSRAM,主控的嵌入约 14 MB。

硬件接线与烧录

  1. 先刷固件,主控刷 master,其余六块刷 node,烧写期间复位引脚要悬空。
  2. 刷数据,主控烧嵌入和分词器文件,六个节点按层区间各烧对应的层文件,需要手动改串口和文件名。
  3. 按菊花链接线,发送和接收各用一路 SPI,启用 DMA,一根线做链式复位,一根回传就绪信号,每块板接一个状态灯。所有板必须共地。
  4. 物理顺序要和层区间一一对应,接反了输出乱码,而且不会有任何报错。

速度和功耗的真实账

每个节点推理要 1.3 秒,生成一个 token 要全链走完。按节点数近似线性增长的规律,加一个节点就多一份延迟。这个速度只适合演示,不适合任何要交互的场景。

供电按 5V 0.3A 估,推理时约 1.53W,空闲约 1.17W。主控串口输入提示词,加 /bench 可以测速。

速度的瓶颈不在算力,在通信。激活值依次流过 SPI 链,后面的层必须等前一块算完,时序上没有并行余地。

已知问题与边界

精度是主要代价。词表裁到 32000,嵌入用 INT4,权重三值化,都会掉点。仓库文档记录了重复 token 的问题,归因于训练不充分加贪心采样。词表外的输入会输出空白。

工程上有两个静默失败的点。Python 侧的三值打包必须和 C、汇编的解包逻辑完全一致,不一致时权重会静默损坏。刷错层映射同样只表现成乱码。排查这两类问题,先从接线顺序和文件对应关系查起。

模型再大就放不下,每节点 15.3 MB 已经贴着 ESP32-S3 的 SRAM 加 PSRAM 上限。要更大模型,得换带更多内存的板子,或者减少每个节点的层数、增加节点数量。

常见问题

这和普通在 MCU 上跑模型差在哪。多数单片机推理把整个模型放进一块芯片,这里用流水线并行把层切开,靠 SPI 在芯片间传激活值,难点在通信时序和内存分配。

适合学到什么。量化流程怎么和硬件限制互相拉扯,菊花链的接线与调试,以及一个静默失败的系统要怎么定位问题。

能改造成别的模型吗。可以,但要重跑裁剪、量化感知微调和打包三步,层数和隐藏维度变化会同时影响每个节点的内存和层区间划分。

七块板要买一样的吗。主控和计算节点刷的固件不同,硬件型号一致最省事。混型号要重新核对层映射和引脚定义,出问题时不报错,只会输出乱码。