一个在多个 ESP32-S3(ESP32S3)上运行的分布式流水线推理引擎执行 1.58-bit(BitNet)语言模型。该项目将切分的 0.5B 大语言模型部署在 7 个 ESP32-S3 集群中,一个作为主节点,其余为节点。主节点运行分词器与词嵌入,注意力层与多层感知机在节点上运行,主节点与节点通过高速 SPI 菊花链通信。
正文
在多个 ESP32-S3 上运行的分布式流水线推理引擎,执行 1.58 位(BitNet)语言模型。
架构
本项目在一个由 7 块 ESP32-S3 组成的集群上运行切片后的 0.5B LLM。其中一块作为主节点,其余作为计算节点。主节点负责分词器与嵌入,注意力层和 MLP 则在节点上运行。主节点与节点之间通过高速 SPI 菊花链通信。
┌─────────────────────────────────────────────────────────┐
│ 主节点 │
│ │
│ [ 提示词 ] ---> BPE 分词器 │
│ │ │
│ Token 嵌入 │
│ (INT4,Flash 中约 14MB) │
│ │ │
│ (SPI 通道 A - 发送至节点 1) │
└───────────────────────┬─────────────────────────────────┘
│ 隐藏状态向量(FP32)
▼
┌─────────────────────────────────────────────────────────┐
│ 计算节点 1 │
│ (SPI 通道 B - 接收自主节点) │
│ │
│ ► 第 0 至 3 层(4 个 Transformer 块) │
│ • RMSNorm(FP16 缩放至 FP32) │
│ • 1.58 位注意力(Q、K、V、O 投影)+ RoPE │
│ • KV 缓存(PSRAM) │
│ • 1.58 位 MLP(Gate、Up、Down 投影) │
│ │
│ (SPI 通道 A - 发送至节点 2) │
└───────────────────────┬─────────────────────────────────┘
│
……(节点 2 至 5)
│
▼
┌─────────────────────────────────────────────────────────┐
│ 计算节点 6 │
│ (SPI 通道 B - 接收自节点 5) │
│ │
│ ► 第 20 至 23 层(4 个 Transformer 块) │
│ • 相同的 1.58 位架构 │
│ │
│ (SPI 通道 A - 发回主节点) │
└───────────────────────┬─────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────┐
│ 主节点 │
│ (SPI 通道 B - 接收自节点 6) │
│ │
│ 最终 RMS Norm │
│ (FP16,位于 'fnorm' 分区的 64KB) │
│ │ │
│ LM Head(与 INT4 嵌入绑定) │
│ │ │
│ 贪婪采样 │
│ │ │
│ [ 输出 ] <--- 下一个 token ID │
└─────────────────────────────────────────────────────────┘
快速开始
请参考工作流指南开始使用本项目。
项目结构
.
├── README.md # 项目文档
├── workflow.md # 刷写、模型准备与接线的分步指南
├── .gitignore # 构建文件与二进制文件的 Git 忽略规则
│
├── docs/
│ └── images/ # 架构图与硬件照片
│
├── master_board/ # 主节点固件(ESP-IDF)
│ ├── main/
│ │ ├── main.cpp # 主节点编排、用户 I/O 与 BPE 分词器
│ │ ├── embedding.cpp # INT4 嵌入查找逻辑
│ │ ├── lm_head.cpp # LM Head 映射与贪婪采样
│ │ └── spi_bus.cpp # 主节点双通道 SPI 驱动
│ ├── partitions.csv # 自定义分区表(token、model、fnorm)
│ └── CMakeLists.txt
│
├── node_firmware/ # 计算节点固件(ESP-IDF)
│ ├── main/
│ │ ├── main.cpp # 节点工作入口与推理循环
│ │ ├── bitlinear.cpp # 1.58 位三值线性层实现
│ │ ├── bitlinear_forward.S # 1.58 位 MAC 运算的汇编优化
│ │ ├── qwen_attention.cpp # Qwen 注意力、RoPE 与 KV 缓存运行时
│ │ ├── lut_table.cpp # 用于极致优化的查找表
│ │ └── spi_bus.cpp # 菊花链 SPI DMA 收发
│ ├── partitions.csv # 节点的层分区布局
│ └── CMakeLists.txt
│
├── python_tools/ # PC 端量化与预处理套件
├── crop_token.py # 词表裁剪(缩减到 32K token)
├── crop_model_weight.py # 嵌入矩阵切片
├── qat_158.py # BitNet QAT(量化感知训练)微调
├── bit4_embedding.py # 嵌入的 INT4 权重打包器
├── pack_tokenizer_bin.py # 将分词器规则序列化为 ESP32 .bin
├── pack_model_bin.py # 将 1.58 位层分块打包以物理对齐
├── look_model_structure.py # 检查 .safetensors 的调试工具
└── flash_*.bat # 多线程快速刷写脚本
许可证
本项目基于 MIT 许可证发布,详见 LICENSE 文件。
致谢
灵感来源、相关工作与参考:
- ESP-32-s3-Story-maker-LLM —— 在 ESP32-S3 上单节点量化 LLM 部署的灵感来源。
- esp32s3-distributed-ai —— 微控制器上多节点分布式 AI 架构的灵感来源。
- BitNet —— 1.58 位三值量化的概念与架构。
评论
社区对将小型语言模型拆分到多个微控制器上运行普遍感到新奇有趣,但有人质疑过度压缩后实质只是“噪声制造器”;同时有人调侃“很快每个灯泡都跑Kubernetes运行AI”折射对技术泛在的戏谑。
认为在微控制器集群上跑LLM是迷人的极客实验
- 多名读者表示将0.5B模型分到七个ESP32或类似小芯片上令人着迷,有人长期幻想用大量简单芯片做大规模并行系统,也有人正在开发类似小项目。(认可度:较多人) > 看到0.5B模型被拆到七个ESP32上运行实在令人着迷(@nonasking_)
质疑压缩过度导致模型沦为噪声制造器
- 有读者认为当前量化压缩程度过高,让设备变成了花哨但无用的语言模型噪声源,虽承认仍有魅力但实用性低。(认可度:少数人) > 这种压缩程度让它成了个花哨的语言模型噪声制造器,不过依然迷人(@cameron_b)
探讨能实用运行LLM的低成本硬件底线
- 有读者认真询问哪些低价芯片能真正跑模型,如苹果小主机是否最低门槛,或有无核显主板、专用AI芯片扩展板;另有人好奇此类集群能否做语法检查或文本游戏内容生成。(认可度:少数人) > 说真的,能实用跑语言模型的低价芯片有哪些?苹果小主机是最低门槛了吗(@sneak)
建议用特定语言或架构优化流水线并行
- 有开发者指出这类流水线处理正适合专门并行语言,但也有人提醒后端未完善前别过度兴奋;另有设想用Rust避免并行坑、搭配Risc-V简单芯片做实验。(认可度:少数人) > 这正是某并行语言瞄准的场景,不过在后端完成前别太激动(@librasteve)
有价值的补充
- 有开发者透露自己正做类似项目,采用更少量化、仅1.5亿参数,认为效果惊人(@matthewfcarlson)
- 想用Rust配合大量Risc-V微控制器做大规模并行系统的想法被提出,虽经济性不明但具探索价值(@ladyanita22)
- 针对此类项目有专门并行语言旨在简化流水线处理,其TinyGo后端尚在构建(@librasteve)