Запускаем 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. Буду рад критике и предложениям.