Языковую модель на 28,9 млн параметров запустили на микроконтроллере ESP32-S3

Разработчик Слава Сербов запустил языковую модель на 28,9 млн параметров на плате с микроконтроллером ESP32-S3. Нейросеть работает полностью локально, а сгенерированный текст выводится на небольшой дисплей.

Для эксперимента разработчик использовал ESP32-S3 N16R8 с 16 Мбайт флеш?памяти и всего 512 Кбайт встроенной SRAM. Кроме того, чип оснащён двумя 32-битными ядрами Xtensa LX7 с частотой до 240 МГц.

Уместить модель на устройстве с ограниченной памятью удалось благодаря четырёхбитной квантизации и распределению весов между разными типами памяти. Финальный файл занимает всего 14,9 Мбайт. При этом из 28,9 млн параметров примерно 559 тысяч относятся к вычислительному ядру, 3,1 млн — к выходному слою, а оставшиеся 25 млн хранятся во флеш?памяти в виде таблицы послойных эмбеддингов. 

Такой подход использовали в моделях Google Gemma. Вместо загрузки сразу всей таблицы в оперативную память система извлекает только те данные, которые необходимы для текущего токена, и добавляет их на разные слои модели. В случае с ESP-32 за один шаг из флеш?памяти считывается около шести строк таблицы. Вычислительное ядро остаётся в более быстрой памяти, а выходной слой и рабочие данные — в PSRAM.

Модель для проекта обучили на TinyStories — синтетическом наборе коротких англоязычных рассказов, которые сгенерировали с помощью GPT-3.5 и GPT-4. В обучающих текстах используется преимущественно лексика, знакомая детям 3–4 лет. Итоговая модель может генерировать похожие простые рассказы для детей, но не умеет отвечать на вопросы, писать код и оперировать фактами о мире.

После оптимизации LLM на ESP-32 генерирует около 9,5 токенов в секунду с учётом вывода текста через последовательный интерфейс. Сами вычисления выполняются со скоростью 9,73 токена в секунду. Для ускорения система заранее переводит выходной слой и активации эмбеддингов в восьмибитный формат и распределяет расчёты между двумя ядрами микроконтроллера.

Ранее разработчик Дэйв Беннетт (Dave Bennett) тоже запускал LLM на ESP32, но его способа оптимизации хватило на модель с 260 тысячами параметров.