28.9MパラメータのLLMを8ドルのESP32-S3マイコンで実行する
GoogleのGemmaモデルによるPer-Layer Embeddingsにより、28.9Mパラメータの言語モデルが8ドルのESP32-S3チップ上で完全に動作し、毎秒9.5トークンのテキスト生成を実現する方法を学びます。
問題: LLMはマイコンに収まらない
マイコンは高速メモリが非常に少ない。約8ドルの人気チップESP32-S3は、わずか512KBのSRAMしか提供しない。通常、推論を実行するにはモデル全体をこの高速メモリにロードする必要があり、モデルサイズが大幅に制限される。同様のチップにおける従来の最先端技術は、わずか260,000パラメータだった。
ブレークスルー: Per-Layer Embeddings
鍵となる洞察は、GoogleのGemma 3nおよびGemma 4モデルから来ている: Per-Layer Embeddings (PLE)。言語モデルのパラメータのほとんどは埋め込みテーブル（トークンをベクトルにマッピングする大きなルックアップテーブル）に存在する。モデルはこのテーブルから読み取るが、計算は行わない。そのため、2500万行のテーブルを低速なフラッシュメモリに残し、各トークンが必要とする数行（約450バイト）だけを取り出すことができる。実際の計算を行う小さな部分は高速なSRAMに留まる。
これは、大きなモデルを実行するコストがほぼゼロであることを意味する — ほとんどの部分をロードすることはない。フラッシュに保存され、少しずつサンプリングされるだけだ。
数値
| 指標 | 値 |
|---|---|
| パラメータ | 28.9M保存（25Mはフラッシュルックアップテーブル） |
| チップ | ESP32-S3（約8ドル）512KB SRAM、8MB PSRAM、16MBフラッシュ |
| 速度 | 約9.5 tok/s エンドツーエンド（9.7 tok/s 純粋計算） |
| 接続性 | なし — すべてデバイス上で実行 |
| モデルサイズ | 14.9MB（4ビット） |
メモリ配置
このトリックは、注意深いメモリ階層によって機能する:
- SRAM（高速、小容量）: 「思考」コア — 各トークンで使用
- PSRAM（中速）: 出力ヘッドとワーキングメモリ
- FLASH（大容量、低速）: 2500万パラメータのテーブル — 各トークンあたり約6行読み取り（約450バイト）
できること（とできないこと）
モデルはTinyStories（Ronen EldanとYuanzhi Li、Microsoft Research）で学習されているため、短くシンプルなストーリーを書き、ほとんど一貫性を保つ。質問に答えたり、指示に従ったり、コードを書いたり、事実を知ったりはしない。この制限はモデルの小さな推論部分に起因する — メモリのトリックは28.9Mパラメータモデルが言える内容を変えない。
自分で実行する
完全なファームウェア、配線図、フラッシュ手順はfirmware/esp32_llm/README.mdにあります。学習、アブレーション、量子化コードはsrc/とexperiments/にあります。完全な方法、アブレーション、オンチップ測定結果はRESULTS.mdに文書化されています。
クレジット
- TinyStories: データセット — 小さなモデルでも一貫した文章を学習できるように十分シンプルな短い合成ストーリー（arXiv:2305.07759）
- Per-Layer Embeddings: GemmaモデルからのGoogleの設計。小さなチップで大きなモデルを可能にする
- llama2.c: Andrej Karpathyの作品。小さな言語モデルを学習し、プレーンなCで実行できるというアイデアに触発された
混乱した歴史
リポジトリは意図的に混乱した歴史を保存しており、初期の数値を膨らませたパラメータ計上のバグと、修正された結果を含む。コミット履歴とRESULTS.mdは、数値がどこでどのように動いたかを示している。
このプロジェクトは、適切なアーキテクチャがあれば、8ドルのマイコンでも驚くほど有能な言語モデルを実行できることを示している — 真のエッジベースのAIアプリケーションの可能性を開く。