Я обучила собственную ИИ-модель на 30 миллионов параметров и запустила её на ESP32
Как обучить свою трансформер-модель и запустить её на крошечном ESP32: токенизация, квантизация до 4 бит и ограничения памяти.
Автор создал собственную языковую модель и запустил ее на микроконтроллере ESP32. В статье описаны архитектура модели, процесс обучения, квантование весов и результаты работы на двух платах.
Архитектура модели
Автор выбрал архитектуру декодера трансформера (decoder-only transformer). Такой же тип архитектуры используется в современных GPT-моделях. Модель предсказывает следующий токен на основе предыдущих токенов.
Модель содержит 12 слоев трансформера и словарь около 6 000 токенов. Автор ориентировался на ограниченный объем памяти ESP32, поэтому модель намеренно выбрана небольшой.
Цель проекта — чат-бот, который понимает диалоги, язык и структуру предложений.
Данные и токенизация
Для обучения использовались общедоступные наборы данных.
Текст разбивался с помощью инструмента SentencePiece. Каждый токен получает числовой идентификатор. Модель работает с числами, а не со словами. Одно слово может состоять из нескольких токенов, например слово dependable.
Данные делились на батчи. Модель обучалась на каждом батче последовательно, что позволяло ей исправлять ошибки и постепенно улучшаться.
Процесс обучения
Обучающий цикл реализован на Python:
- Модель получает текст и предсказывает следующий токен.
- Функция кросс-энтропии (cross entropy) оценивает точность предсказания.
- Обратное распространение ошибки (back propagation) определяет, как изменить веса модели.
- Градиентное отсечение (gradient clipping) ограничивает величину изменений весов.
Модель начала обучение со случайных весов. На этом этапе она выдавала случайные фрагменты текста. К 100 млн токенов качество улучшилось, но ошибки сохранялись. К 300 млн токенов модель заметно улучшилась.
Валидационная функция потерь (validation loss) снизилась по всем трем потокам данных. На диалоговых данных значение упало с 8,81 до 2. Это означает, что модель лучше предсказывает следующий токен.
Квантование весов
ESP32 располагает 32 МБ флеш-памяти и 16 МБ PS RAM. Модель имеет 3,5 млн параметров. При хранении всех параметров в 32-битном формате требуется около 122 МБ памяти, что примерно в четыре раза больше объема флеш-памяти.
Для уменьшения объема применено квантование. Большинство весов переведено в 4-битный формат. Часть весов сохранена в 16-битном формате, чтобы не потерять всю точность. Квантование снижает точность модели, но автор считает баланс между размером и качеством удовлетворительным.
Работа модели на ESP32
- Веса хранятся во флеш-памяти.
- PS RAM используется для временных данных.
- Процессор выполняет вычисления, предсказание и чтение весов.
Процесс повторяется для каждого токена до формирования полного ответа.
Автор получил ответы, сгенерированные моделью на ESP32. Модель не справляется со сложными вопросами, но дает осмысленные ответы.
Запуск на Xiao ESP32S3
Плата Xiao ESP32S3 имеет 8 МБ флеш-памяти и 8 МБ PS RAM. Даже при квантовании модель не помещается в память платы полностью.
Веса загружаются с microSD-карты, которая расположена на обратной стороне круглого дисплея. Чтение с microSD медленнее, чем из RAM. Поэтому использовались два буфера: пока процессор обрабатывает данные из одного буфера, следующая часть весов загружается во второй.
ESP32S3 имеет два ядра. Нагрузка распределялась между ними. Однако для генерации каждого токена процессор должен пройти все 12 слоев трансформера, что занимает много времени.
| Плата | Объем флеш-памяти | Объем PS RAM | Скорость генерации |
|---|---|---|---|
| ESP32 (dev kit) | 32 МБ | 16 МБ | 1,4 токена в секунду |
| Xiao ESP32S3 | 8 МБ | 8 МБ | около 0,74 токена в секунду |
Альтернативные способы
Автор отмечает, что для работы на ESP32 есть более простые способы.
Облачный API (на примере Google Gemini). Модель запускается в Google AI Studio. Пользователь получает API-ключ и добавляет его в код ESP32. Плата отправляет запросы на сервер Google и выводит ответ на экран. Модель не работает на самой плате.
Локальный запуск через llama. Программа позволяет запускать модели ИИ на компьютере. Модель загружается и работает в терминале. API-ключ добавляется в код ESP32 так же, как при использовании Gemini. Отличие состоит в том, что модель работает на локальном компьютере, а не на сервере Google.
Выбор способа зависит от целей проекта. Обучение модели с нуля и запуск на ESP32 требуют значительно больше времени и ресурсов.
Автор планирует опубликовать инструкции и дополнительные материалы на GitHub.
Введение
Автор представляет цель видео: обучить собственную ИИ-модель и запустить её на ESP32.
- Автор: Норина
- Цель: обучить собственную ИИ-модель и запустить её на ESP32
- Автор отмечает, что задача была сложной и потребовала решения многих проблем
Архитектура модели
Выбрана архитектура декодерного трансформера, которая используется в современных GPT. Модель предсказывает следующий токен на основе предыдущих.
- Архитектура: decoder-only transformer, как в современных GPT
- Принцип работы: предсказание следующего токена на основе предыдущих
- Модель имеет 12 слоёв трансформера
- Словарь около 6000 токенов
- Модель создаётся для слабого устройства с очень ограниченной памятью
Данные и токенизация
Цель — создать чат-бота, который понимает диалоги, язык и структуру предложений. Данные взяты из публично доступных наборов. Текст токенизируется с помощью SentencePiece, данные делятся на пакеты (батчи).
- Цель: создать чат-бота, понимающего диалоги и структуру предложений
- Использованы публично доступные наборы данных
- Инструмент токенизации: SentencePiece
- Токены получают числовые значения; модель работает с числами, а не словами
- Слова вроде dependable могут состоять из нескольких токенов
- Данные разделены на батчи для постепенного обучения
Код обучения
Модель предсказывает следующий токен, кросс-энтропия оценивает точность, обратное распространение ошибки корректирует веса. Используется gradient clipping для контроля изменений.
- Модель получает текст и предсказывает следующий элемент
- Кросс-энтропия (cross entropy) оценивает правильность предсказаний
- Обратное распространение ошибки (back propagation) корректирует веса модели
- Gradient clipping используется для контроля величины изменений весов
Процесс обучения
Модель начинает со случайных весов и выдаёт случайные фрагменты. К 100 миллионам токенов качество начинает улучшаться, к 300 миллионам — заметно улучшается.
- Начальное состояние: случайные веса, случайные фрагменты текста
- Около 100 миллионов токенов: начало улучшений, но ошибки остаются
- К 300 миллионам токенов: значительное улучшение
- Валидационная потеря снизилась на всех трёх потоках данных
- На разговорных данных потеря снизилась с 8.81 до 2
Запуск на ESP32
Плата ESP32 имеет 32 МБ флеш-памяти и 16 МБ PS RAM. Модель с 3.5 миллионами параметров при хранении в 32-битном формате занимает около 122 МБ, что примерно в четыре раза больше флеш-памяти. Поэтому применяется квантизация до 4 бит.
- Плата: 32 МБ флеш-памяти и 16 МБ PS RAM
- Модель: 3.5 миллиона параметров
- При 32-битном хранении нужно около 122 МБ памяти
- Это примерно в четыре раза больше флеш-памяти платы
- Квантизация: хранение весов в 4-битном формате с потерей точности
- Часть весов оставлена в 16 бит для сохранения точности
- Веса хранятся во флеш-памяти, PS RAM используется для временных данных
- Процессор генерирует токены по одному до завершения ответа
Результаты на ESP32
Модель работает на ESP32 локально, отвечает на простые вопросы, со сложными справляется хуже. Ответы признаны осмысленными.
- Модель запущена локально на ESP32
- Сложные вопросы модель обрабатывает плохо
- Ответы оцениваются как осмысленные
Запуск на Xiao ESP32S3
Автор пробует запустить модель на меньшей плате Xiao ESP32S3 с 8 МБ флеш-памяти и 8 МБ PS RAM. Веса читаются с microSD-карты, что медленнее чтения из RAM. Используются два буфера и два ядра процессора.
- Плата Xiao ESP32S3: 8 МБ флеш-памяти и 8 МБ PS RAM
- Даже с квантизацией модель не помещается целиком в плату
- Веса читаются с microSD-карты, расположенной на обратной стороне круглого дисплея
- Чтение с microSD медленнее, чем из RAM
- Используются два буфера: пока один обрабатывается, другой загружает веса
- Работа распределяется между двумя ядрами ESP32S3
- Скорость: 1.4 токена в секунду на dev kit, около 0.74 токена в секунду на Xiao
Альтернативные способы
Автор упоминает более простые способы: использование API Google Gemini через Google AI Studio и локальный запуск через Llama. Инструкции будут опубликованы на GitHub.
- Вариант 1: Google Gemini через Google AI Studio, получение API-ключа и добавление в код ESP32
- ESP32 отправляет запрос на сервер Google и получает ответ, который выводится на экран
- Вариант 2: Llama для локального запуска модели на компьютере через терминал
- Для Llama также можно получить API и подключить его к коду ESP32
- Инструкции будут опубликованы на GitHub
Заключение
Обучение модели с нуля и запуск на ESP32 — сложная задача. Выбор способа зависит от целей. Все материалы автор обещает выложить на GitHub.
- Более простые способы существуют, чем обучение модели с нуля
- Выбор подхода зависит от целей
- Все материалы будут опубликованы на GitHub
