在售价 8 美元的微控制器上运行一个 2890 万参数的大语言模型

开放求职 · 𝕏 slvDev · LinkedIn

这是一个拥有 2890 万个参数的语言模型，运行在 ESP32-S3 上——这是一块售价约 8 美元的微控制器。它在芯片本身上运行，无需向服务器发送任何数据，并以大约每秒 9 个 token 的速度将每个单词写入连接到芯片的小屏幕。此前人们在这类芯片上运行的语言模型只有 26 万个参数，因此这个模型的参数量大约是前者的 100 倍。它之所以能装得下，是因为模型的大部分数据存储在闪存而非 RAM 中，这利用了 Google Gemma 模型中的一项名为“逐层嵌入”的技术。

数据概览

参数量	存储 2890 万个参数（其中 2500 万个存储在闪存查找表中）
芯片	ESP32-S3，约 8 美元，配备 512KB SRAM、8MB PSRAM 和 16MB 闪存
速度	端到端约 9.5 tok/s（纯计算速度 9.7 tok/s）
连接性	无，一切均在设备上运行
模型大小	4 位量化下为 14.9MB

为什么这很难，以及它如何做到

微控制器的快速内存非常有限。ESP32-S3 只提供 512KB 的 SRAM。通常情况下，整个模型必须能够从该内存中访问，这导致你只能使用极小的模型，这也是为什么此前这类芯片上的模型只有 26 万个参数。

解决方法是根本不再将模型放入快速内存。语言模型的大部分参数位于嵌入表中，模型从该表中读取数据而非进行计算。因此，你可以将那个包含 2500 万行的表留在慢速闪存中，每个 token 只需从中提取所需的几行（约 450 字节），而执行实际计算的小部分则留在快速内存中。这样一来，运行这个大模型的成本几乎为零，因为你从未加载其大部分内容。它只是待在闪存里，每次被少量采样。

这个想法来自 Google 的逐层嵌入技术，源自 Gemma 3n 和 Gemma 4。在这里，它运行在微控制器的内存布局上，而非手机或 GPU。据我所知，此前还没有人尝试过在如此小的芯片上应用这项技术。

  SRAM  (fast, tiny)   the "thinking" core, used on every token
  PSRAM (medium)       the output head and working memory
  FLASH (huge, slow)   the 25M-param table, about 6 rows read per token (~450 B)

它能做什么，以及不能做什么

该模型基于 TinyStories 数据集训练，因此它能够撰写简短、简单的故事，并且大多能保持连贯性。它不会回答问题、遵循指令、编写代码或掌握事实知识。这一限制源于模型中负责推理的部分规模较小，而记忆技巧并不能改变这一点。这里真正有趣的是其架构——将一个大模型塞进一块微小的芯片，而非一个拥有 2890 万参数的模型能说出什么。

自行运行

固件、接线以及烧录步骤均位于 `firmware/esp32_llm/README.md` 文件中。训练、消融实验和量化代码在 `src/` 和 `experiments/` 目录下。完整方法、消融实验以及片上测量结果已写入 `RESULTS.md` 文件。

致谢

TinyStories 是该模型训练所用的数据集：它包含简短的合成故事，其内容足够简单，以至于小模型也能学会连贯地写作（Ronen Eldan 和 Yuanzhi Li，微软研究院，arXiv:2305.07759）。另一半功劳归于逐层嵌入（Per-Layer Embeddings），这是 Google 在 Gemma 模型中的设计，正是它让一个大模型得以适配到一块小芯片上。

Andrej Karpathy 的 llama2.c 项目让许多人——包括我在内——相信，你完全可以训练一个小型语言模型，并用纯 C 语言来运行它。本项目正是由此发展而来。

实际过程

我特意将这段混乱的历史留在了仓库中。其中包括我在自己参数统计中发现的一个错误，该错误曾导致早期数据被夸大，以及在我修正该错误后得出的更正结果。提交历史记录和 `RESULTS.md` 文件展示了数据的变化过程及其原因。