言語モデルをマイコンで動かそうとすると、まずメモリの壁に当たる。ESP32-S3のフラッシュは16MBで、数億のパラメーターを持つモデルはどう詰めても入らない。ならば1枚に収めるのをやめて、複数枚で分担すればいい。Low Zi Hong氏はESP32-S3を7枚つなぎ、24層のトランスフォーマーを分けて動かす仕組みを作り、量子化のスクリプトから配線図までをMITライセンスでGitHubに公開した。
構成は親機1枚と計算機6枚からなる。親機が文章をトークンに切り分け、埋め込みを引いて最初の中間状態を作る。計算機は1枚あたり4層ずつを受け持ち、計算した中間状態を次の1枚へ渡す。6枚を通ると24層すべてが終わり、親機に戻って次の単語を選ぶ。板の間はSPIのデイジーチェーンでつないだ。同種の先行プロジェクトが無線のESP-NOWを使っていたのに対し、有線にすることで遅延を減らす狙いだ。
モデルを載せる工夫は3段構えになっている。中心は1.58ビットの三値量子化で、重みを−1、0、+1の3通りだけで表す。8ビットや4ビットでは足りず、ここまで落として初めて数字が合ったと本人は書く。加えて語彙を3万2000トークンに刈り込み、埋め込みを4ビットに詰めた。もとは5億パラメーターのモデルで、語彙の削減後は約3億8600万になる。
本人が「いちばん厄介だったのはモデルの数学ではなくシステム設計だった」と振り返るのは、この分割の詰め方だ。1枚あたりの層数を少なくすれば領域が余り、多くすればプログラムの置き場所が消える。しかもフラッシュの区画表と、PC側で切り出したモデルの断片とを、バイト単位でぴったり合わせる必要がある。ずれが1つあれば、そこから先は意味のない値を読み続ける。
書いたのはESP-IDFのC++で、Arduinoの枠組みは使っていない。三値の掛け算と足し込みはアセンブリで書き、参照表も用意した。中間状態の受け渡しはDMAに任せている。消費電力は推論中で計算機6枚分が約1.53Wだった。親機はUSB Type-Cから給電する。
ただし速度は毎秒0.11トークン、1つの単語を出すのに約9秒かかる。本人はこれを「ハードウェアと数学の概念実証であって、GPTの代わりではない」と明言し、同梱のサンプルモデルが手元のパソコンの計算力不足で過学習していることも先に断っている。投稿には改善案も集まった。三値の重みを浮動小数へ変換する参照表を使えば速度がおよそ2倍になる、フラッシュの読み出しをQIOモードにすれば標準のDIOより大幅に速くなる、といった具体的な指摘だ。