Запускаем 8B LLM в браузере: AQLM.rs
Меня долгое время интересовал запуск больших языковых моделей на пользовательских устройствах: есть что‑то в том, чтобы запустить одну из лучших языковых моделей на обычном домашнем компьютере или на мобильном телефоне, помещающемся в карман.
В этом посте я расскажу о своём пет‑проекте AQLM.rs. Я написал инференс модели Llama 3.1 8B, работающий в браузере на WebAssembly без использования GPU, с помощью алгоритма сжатия, разработанного нашей лабораторией.
Попробовать можно на сайте проекта.
Почему 8B
Запуск языковых моделей на пользовательском девайсе — далеко не новая идея. Например, такие модели, как Llama 3.2 1B и Llama 3.2 3B, изначально создавались с целью запуска на маломощных устройствах.
Модели, имеющие 8 млрд параметров, идеально подходят для демонстрации того, что может быть помещено в браузер с помощью продвинутых алгоритмов сжатия.
В разжатом виде на каждый параметр приходится по 16 бит, и, как следствие, модель 8B весит 16 ГБ. Используя стандартные 4-битные методы сжатия, такие как nf4, можно сжать её до 4 ГБ. В этом проекте я использую экстремальное сжатие в 2 бита, сжимающее тело модели в 8 раз. Для слоёв головы и эмбеддингов я всё ещё использую 4-битное и 8-битное сжатие, поэтому модель получается размером в районе 2,5 ГБ.
Плюс экстремального сжатия заключается в том, что с уменьшением размера сама модель ускоряется, поскольку скорость вычисления такого рода сильно упирается в работу с памятью. Меньше памяти — значит, больше скорость.
Отдельно хочется упомянуть, что Llama 3.1 8B, сжатая до 2 бит нашим алгоритмом, всё ещё лучше, чем Llama 3.2 3B в несжатом виде, потому что занимает в 2 раза меньше места.
Много больших матриц
Большие языковые модели состоят из матриц. Основная вычислительная нагрузка при работе модели — это обыкновенное умножение матрицы на вектор. Оптимизацией именно этой части занимаются методы сжатия. Они пытаются уменьшить размер матриц, меняя их представление на более компактное, минимизируя потери качества.
В мае 2024 года наша команда из Yandex Research вместе с коллегами из университетов IST Austria и KAUST опубликовала исследование, в котором мы описали алгоритм PV‑Tuning для улучшения методов сжатия больших языковых моделей без изменения формата сжатых весов (я писал про неё на «Хабре»).
В моём проекте в качестве базового метода, улучшаемого PV‑Tuning, я использую AQLM. В этом методе сжатым представлением является аддитивная векторная квантизация. Для 2-битной квантизации (это когда на один параметр приходится всего 2 бита, а не 16, как у изначальной модели, то есть в 8 раз меньше) каждая строчка матрицы собирается из маленьких кусков по восемь чисел. Каждый такой кусок — это сумма двух векторов из словарей по 256 элементов каждый. Получается, что для хранения значения восьми элементов матрицы надо потратить 2 раза по 8 бит на индексы. Отсюда и берётся 2 бита на параметр.
WebAssembly и Rust
Благодаря развитию WebAssembly программы для браузера стало можно писать почти на любом языке. Несколько лет назад я прошёл курс по Rust в ШАД, влюбился в язык, но никак не мог найти ему применение в своих пет‑проектах. Я не мог упустить этот шанс и реализовал весь инференс на нём.
Приятным сюрпризом было то, что на Rust написано очень много фундаментальных библиотек для инфраструктуры вокруг LLM.
Например, у Hugging Face есть формат под названием safetensors. Оказалось, что его реализация полностью написана на Rust! Каждый раз, когда кто‑нибудь в своём коде на Python использует safetensors‑модель с Hugging Face, он использует Rust.
На Rust также написан токенайзер под названием tiktoken от OpenAI, который используется в новых «ламах».
Многопоточность
Чтобы ускорить работу модели, я реализовал многопоточность с помощью веб‑воркеров. Они поддерживают двунаправленную связь между тредами через передачу сообщений. Моё решение основано на model‑parallel‑подходе: все матрицы разделяются по размерности выхода, и каждый воркер получает свой кусок каждой матрицы.
Самая сложная часть заключалась в организации взаимодействия между воркерами и основным тредом. Для этого я написал кастомный RPC‑стек для воркеров с интеропом между Rust и JavaScript.
Когда главному треду нужно умножить вектор на матрицу, он составляет запрос к каждому воркеру, сериализует его и отправляет в JavaScript‑рантайм. Затем JavaScript пересылает его воркеру, который десериализует, обрабатывает запрос и сериализует результат перед отправкой обратно. После этого JavaScript отправляет результат в главный тред, где тот десериализуется.
Этот подход позволил увеличить производительность примерно в 2 раза.
Где попробовать
Если не хочется ждать загрузки модели, посмотрите видео. А чтобы поиграться на своём компьютере, зайдите на демо.
При первом обращении дождитесь окончания загрузки модели, это может занять несколько минут. Рекомендую общаться с ней на английском языке — так она будет существенно умнее.
Исходный код проекта загружен на GitHub. Буду рад критике и предложениям.
