ESP32‑S3 集群使用 1.58 位 BitNet 量化运行 0.5B 语言模型

概要

一个由七个 ESP32‑S3 芯片组成的集群，通过高速 SPI 级联通信，在微控制器之间分割模型，运行一个 0.5 B 参数的语言模型，采用 1.58 位（BitNet）三值量化。这表明即使是普通的物联网芯片，也能参与非平凡大语言模型的分布式推理。

项目功能说明

ESP32‑S3 LLM 集群 实现了一个分布式推理流程，将一个 0.5 B 参数的模型分割到七个 ESP32‑S3 开发板上。其中一块板子作为 主控节点 —— 负责分词、嵌入查找、最终 RMS‑norm、语言模型头部和贪婪采样；其余六块板子作为 计算节点，每块依次执行四个 Transformer 块（注意力和 MLP）。主控与计算节点之间通过双向 SPI 级联通信，使隐藏状态向量（FP32）在流水线中正向流动，并返回主控以生成最终输出。

架构亮点

  • 主控节点：运行 BPE 分词器、INT4 嵌入表（约 14 MB 存储在 Flash 中），以及最终后处理。
  • 计算节点：每块节点托管四个 Transformer 层，包含：
    • RMSNorm（FP16 → FP32 缩放）
    • 1.58 位三值注意力（Q/K/V/O 投影）并支持 RoPE
    • KV 缓存存储在外部 PSRAM 中
    • 1.58 位三值 MLP（门控、上投影、下投影）
  • SPI 级联：每块板子配备两个 SPI 通道（CH A 用于正向发送，CH B 用于反向接收），实现低延迟、高吞吐的数据传输。
  • 量化技术：采用微软的 BitNet 1.58 位三值方案，大幅减少权重存储需求，同时保持推理质量。

工作原理——分步推理流程

  1. 输入提示 → 主控通过 BPE 进行分词并查找 INT4 嵌入。
  2. 隐藏状态（FP32） 通过 SPI CH A 发送到计算节点 1。
  3. 节点 1 处理第 0–3 层，更新隐藏状态，并转发给节点 2。
  4. 节点 2–5 重复相同流程，每块处理四个连续的 Transformer 块。
  5. 节点 6 处理最后的第 20–23 层，并通过 SPI CH B 将隐藏状态返回主控。
  6. 主控 应用最终 RMS‑norm（FP16）、共享语言模型头部（INT4），并进行贪婪采样，输出下一个词元。
  7. 该循环对每个后续词元重复执行。

仓库中的示意图展示了这一流水线，清晰呈现了主控到节点的数据流，以及模型权重在 Flash 和 PSRAM 之间的分布情况。

仓库结构

  • master_board/：主控板的 ESP‑IDF 固件，包含分词器、嵌入层、语言模型头部和双通道 SPI 驱动程序。
  • node_firmware/：计算节点的固件，包含 1.58 位三值线性层（bitlinear.cpp）、汇编优化的 MAC 内核（bitlinear_forward.S）、Qwen 注意力实现和 KV 缓存管理。
  • python_tools/：主机侧工具，用于模型准备：
    • 词表裁剪（crop_token.py）
    • 嵌入层切片（crop_model_weight.py）
    • BitNet 量化感知训练（qat_158.py）
    • 二进制打包以对齐 Flash（pack_model_bin.py）
  • workflow.md：端到端指南，涵盖硬件接线、烧录和模型准备流程。

Hacker News 社区反响

该项目引发了多种评论，反映出热情与质疑并存的复杂情绪：

ladyanita22：“想象一下，用微小的微控制器构建一个大规模并行系统——Rust 可能让并行化更安全。”（推测该概念可扩展至 RISC‑V 集群。）

tdhz77：“很快，每个灯泡里都会运行 Kubernetes 的 AI。”（对 AI 无处不在的戏谑夸张。）

cameron_b：“这种压缩让它变成一个花哨的 LLM 噪音生成器，但依然很迷人。”（对实际质量表示怀疑。）

librasteve：“像 https://bil-lang.org 这样的项目致力于并行流水线处理；我们首先需要一个 TinyGo 后端。”（指出相关语言层面的努力。）

NDlurker：“它能否处理基础文字处理器中的语法检查，或为文字游戏生成世界？”（询问实际应用场景。）

matthewfcarlson：“我正在做一个类似的项目，使用 150 M 参数模型——这太棒了。”（展示平行开发。）

sneak：“有哪些低成本芯片可以运行 LLM？Mac Mini 是最便宜的吗，还是有专为 Pi HAT 设计的 AI 芯片？”（寻求硬件替代方案。）

总体而言，社区认可其技术新颖性，但对可扩展性、性能和实际用途提出质疑。

为何重要

  • 成本效益：ESP32‑S3 芯片单价仅几美元；七节点集群总价低于 50 美元，远低于单块 GPU 加速板。
  • 边缘 AI 民主化：证明了复杂 LLM 推理可部署在超低功耗、电池供电的设备上。
  • 量化技术创新：在真实硬件上验证了 BitNet 的 1.58 位三值格式，弥合了研究量化与嵌入式部署之间的鸿沟。
  • 分布式微控制器 AI：为异构微控制器之间的流水线 AI 开辟了新设计空间，类似于早期并行计算，但面向现代深度学习负载。

快速上手

  1. 克隆仓库。
  2. 按照 workflow.md 指南，将主控和计算节点固件烧录到 ESP32‑S3 开发板上。
  3. 使用 Python 工具对兼容的 0.5 B 模型进行量化和打包（仓库提供 BitNet QAT 和 INT4 嵌入打包脚本）。
  4. 按照仓库硬件照片所示，使用 SPI 引脚将板子连接成级联结构。
  5. 为集群供电，并通过主控的串口控制台输入提示，接收生成的词元。

许可证

本项目采用 MIT 许可证发布。

总结

ESP32‑S3 LLM 集群证明，仅需一组廉价的微控制器，通过激进的三值量化，即可协同运行一个 5 亿参数的语言模型。尽管输出质量无法与全精度模型媲美，但这项工作展示了实现超低成本、分布式边缘 AI 的可行路径。

Sources