ESP32‑S3 集群以 1.58 位 BitNet 量化運行 0.5B 語言模型

簡要重點

一個七節點的 ESP32‑S3 集群可運行具備 0.5 B 參數的語言模型,採用 1.58 位(BitNet)三值量化,將模型分割至多個微控制器並透過高速 SPI 串接鏈路進行通訊。這顯示即使是一般物聯網晶片也能參與非平凡語言模型的分散式推論。


專案功能說明

ESP32‑S3 LLM 集群 實作了分散式推論流程,將 0.5 B 參數模型切分至七個 ESP32‑S3 電腦板。其中一塊板作為 主控節點,負責詞元化、嵌入查找、最終 RMS‑norm、語言模型頭以及貪婪採樣;其餘六塊板則作為 運算節點,依序執行四個 Transformer 模塊(注意力與 MLP)。主控節點與運算節點之間透過雙向 SPI 串接鏈路通訊,使隱藏狀態向量(FP32)能沿著流程前進,並回傳至主控節點以產生最終輸出。

架構亮點

  • 主控節點:執行 BPE 詞元化器、INT4 嵌入表(約 14 MB 存於快閃記憶體),以及最終後處理。
  • 運算節點:每塊節點主機四個 Transformer 層,包含:
    • RMSNorm(FP16 → FP32 標準化)
    • 1.58 位三值注意力(Q/K/V/O 投影)搭配 RoPE
    • KV 快取儲存在外部 PSRAM
    • 1.58 位三值 MLP(gate、up、down 投影)
  • 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. 此流程針對每個後續詞元重複執行。

倉儲中的圖示可視化此流程,顯示主控節點至運算節點的資料流,以及模型權重在快閃記憶體與 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)
    • 二進位檔打包以對齊快閃記憶體(pack_model_bin.py)
  • workflow.md:端到端指南,涵蓋硬體接線、燒錄與模型準備。

Hacker News 社群反應

此專案引發多樣評論,展現出熱情與質疑並存的態度:

ladyanita22:「想像一個由微型微控制器構成的巨量平行系統——Rust 可能讓平行化更安全。」(推測此概念可擴展至 RISC‑V 集群。)

tdhz77:「不久的將來,每個燈泡都將運行 Kubernetes 的 AI。」(對 AI 普及化的戲謔誇張。)

cameron_b:「壓縮讓它變成一個花哨的語言模型噪音產生器,但依然迷人。」(對實際品質表示懷疑。)

librasteve:「像 https://bil-lang.org 這樣的專案致力於平行流水線處理;我們首先需要一個 TinyGo 後端。」(指出相關語言層級的發展。)

NDlurker:「這能處理基本文字處理器的文法檢查,或為文字冒險遊戲生成世界嗎?」(探討實際應用場景。)

matthewfcarlson:「我正在開發一個類似的專案,使用 150 M 參數模型——這太棒了。」(顯示平行開發。)

sneak:「有哪些低成本晶片能運行 LLM?Mac Mini 是最便宜的嗎?還是有專用 AI 晶片可用於 Pi HAT?」(尋求硬體替代方案。)

整體而言,社群肯定其技術創新,同時質疑其可擴展性、效能與現實應用價值。

為何重要

  • 成本效益:ESP32‑S3 晶片每顆僅數美元;七節點集群總成本低於 50 美元,遠低於單一 GPU 加速板。
  • 邊緣 AI 民主化:證明複雜語言模型推論可推至超低功耗、電池供電的裝置。
  • 量化技術創新:在真實硬體上驗證 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 集群證明,僅需少量低成本微控制器,即可透過激進的三值量化協同運行半十億參數的語言模型。儘管輸出品質可能無法與全精度模型相比,但此工作展現了一條可行的超低成本、分散式邊緣 AI 路徑。

Sources