ESP32-S3 запустил TinyStories на 28.9 млн параметров при 9.88 токена/с
ESP32-S3 с 512 КБ SRAM генерирует текст локально со скоростью 9.88 токена/с, используя модель размером 14.9 МБ. В 28.9 млн параметров входят 25 млн редко выбираемых элементов Per-Layer Embeddings, поэтому основная вычислительная нагрузка остаётся небольшой. Результат определила раскладка данных между SRAM, PSRAM и flash с учётом характера доступа.
На плате ESP32-S3 N16R8 запустили языковую модель TinyStories с 28.9 млн хранимых параметров. Инференс выполняется локально на двух ядрах Xtensa LX7 с частотой до 240 МГц, сервер и сетевое подключение не требуются. Опубликованная реализация показывает 94.9 мс вычислений на токен и 9.88 токена/с при работающем последовательном выводе.
Аппаратная конфигурация включает 512 КБ встроенной SRAM, 8 МБ внешней PSRAM и 16 МБ flash. Образ квантованной модели занимает около 14.9 МБ. Параметры и временные буферы распределены по уровням памяти в зависимости от частоты и последовательности обращений: наиболее горячие данные остаются в SRAM, крупные массивы размещаются в PSRAM, редко используемая таблица хранится во flash.
Заявленные 28.9 млн параметров требуют пояснения. Около 559 тыс. приходится на плотное ядро трансформера, ещё примерно 3.1 млн содержит выходная голова. Основную часть, около 25 млн, образует таблица Per-Layer Embeddings, или PLE. На каждом шаге из неё выбирается лишь несколько строк, поэтому весь массив не участвует в вычислении очередного токена.
PLE позволяет вынести значительную долю ёмкости модели во flash и обращаться к ней по мере прохождения слоёв. По измерениям автора, для одного токена требуется около шести случайных чтений общим объёмом примерно 450 байт. Чтение 512-байтной строки из flash занимает 20.3 мкс. Синтетическая доля PLE в измеренном трафике памяти составляет около 0.7%; это число описывает трафик и не равно доле общей задержки инференса.
Остальные веса хранятся в формате int4, промежуточные представления переведены в int8. Выходная голова исполняется на двух ядрах. Измеренная пропускная способность достигает 60.7 МБ/с для последовательного чтения PSRAM и 240 МБ/с для SRAM, поэтому размещение массивов влияет на скорость сопоставимо с арифметическими оптимизациями.
Модель обучена на TinyStories и продолжает короткие синтетические рассказы. Она не рассчитана на диалог, выполнение инструкций, генерацию кода или извлечение фактов. Число 28.9 млн характеризует общий объём хранимых параметров, тогда как вычислительное ядро значительно меньше и использует разреженный доступ к PLE.
Исходный код обучения и квантования, C-runtime, прошивка ESP32-S3 и результаты аппаратного профилирования опубликованы в репозитории esp32-ai. Демонстрация показывает полезный для MCU-проектов приём: ёмкость модели можно наращивать за счёт дешёвой flash, если архитектура ограничивает число случайных чтений и оставляет интенсивно используемые данные во внутренней памяти.