КПД

Fast Inference of Mixture-of-Experts Language Models with Offloading
[Статья][Код]

Введение

Смеси экспертов (MoE), использующие лишь часть параметров на инференсе, дают значительную экономию вычислений. Однако SOTA-модели оказываются слишком тяжеловесными для многих GPU и встает вопрос эффективного оффлоадинга.

Наивная подгрузка экспертов работает довольно медленно из-за передачи большого обьема данных с ОЗУ/RAM на GPU. Возникает вопрос - можно ли оптимизировать данную процедуру?

И ответ, как можно было заранее догадаться, — да 😅.

Метод

Чтобы понять, как оптимизировать загрузку экспертов, авторы смотрят на то, как активируются эксперты. И оказывается, что в конкретном слое один и тот же эксперт может активироваться 2-4 раза подряд, потому если их оставлять в памяти, они будут доступны при прогонке следующего токена. Отсюда возникает идея поддерживать LRU кэш токенов, и сгружать экспертов на диск по необходимости.

Следующий нюанс в том, как загружать следующие блоки. Для обычных трансформеров никакой проблемы нет, но для экспертов заранее неизвестно, какие активируются. Но можно воспользоваться стратегией из De Javu, где предсказывали активные нейроны по активациям предыдущего блока, и отбирать экспертов, подавая в gating функцию на следующем блоке текущие активации (speculative loading). Это работает в силу residual структуры трансформера.

Чтобы уменьшить размер модели, используется квантизация HQQ. Не самый очевидный выбор, ибо это Data-Free квантизация, уступающая более продвинутым подходам. Но раз работает, то и ладно. QMoE, говорят, не завелось (может надо было квантовать в больше, чем 1 бит?). Attention квантуют в 4 бита, экспертов в 2 и 3 бита.

Эксперименты

Рассматриваются две конфигурации системы - условный Colab с бесплатной GPU и хороший геймерский ПК. Для GPU c 12Gb VRAM можно держать кэш из 2 экспертов, а для 16Gb - влезают 4.

Чем больше размер кэша - тем больше hit-ratio (вероятность того, что эксперт лежит в кэше) и точность удачного speculative loading.
Выбранная схема квантования более менее оптимальна, так как квантизация attention в меньшее число бит уже сильно просаживает качество, а 4 бита некритично хуже fp16. Конечные конфигурации в 4-5 раз меньше по размеру базовой модели. Некоторая просадка по качеству на бенчмарках есть, но приемлемая.

Все компоненты метода - LRU-cache, пре-подгрузка экспертов заметно ускоряют инференс. Итоговое ускорение порядка 2-2.5 раз на RTX3060/RTX3080 и до 3.5 раз на T4 по сравнению с наивной реализацией в accelerate. Для прикола замеряют еще и на A100, куда влезает квантованная модель и так (для сравнения).

Вывод

Простой и разумный подход с практически полезным результатом. По всей видимости, умный оффлоадинг - одно из самых, если не самое перспективное направление, ускорения моделях на пользовательских устройствах. Однако, букву “О” гармонично вставить в название канала не получается. Кажется, что результаты работы можно улучшить за счет более сильного алгоритма квантизации экспертов, уменьшив просадку качества.

👍9

www.tgoop.com/quant_prune_distill/177

743 viewsedited  Jan 14, 2024 at 21:03

Fast Inference of Mixture-of-Experts Language Models with Offloading
[Статья][Код]

Введение

Смеси экспертов (MoE), использующие лишь часть параметров на инференсе, дают значительную экономию вычислений. Однако SOTA-модели оказываются слишком тяжеловесными для многих GPU и встает вопрос эффективного оффлоадинга.

Наивная подгрузка экспертов работает довольно медленно из-за передачи большого обьема данных с ОЗУ/RAM на GPU. Возникает вопрос - можно ли оптимизировать данную процедуру?

И ответ, как можно было заранее догадаться, — да 😅.

Метод

Чтобы понять, как оптимизировать загрузку экспертов, авторы смотрят на то, как активируются эксперты. И оказывается, что в конкретном слое один и тот же эксперт может активироваться 2-4 раза подряд, потому если их оставлять в памяти, они будут доступны при прогонке следующего токена. Отсюда возникает идея поддерживать LRU кэш токенов, и сгружать экспертов на диск по необходимости.

Следующий нюанс в том, как загружать следующие блоки. Для обычных трансформеров никакой проблемы нет, но для экспертов заранее неизвестно, какие активируются. Но можно воспользоваться стратегией из De Javu, где предсказывали активные нейроны по активациям предыдущего блока, и отбирать экспертов, подавая в gating функцию на следующем блоке текущие активации (speculative loading). Это работает в силу residual структуры трансформера.

Чтобы уменьшить размер модели, используется квантизация HQQ. Не самый очевидный выбор, ибо это Data-Free квантизация, уступающая более продвинутым подходам. Но раз работает, то и ладно. QMoE, говорят, не завелось (может надо было квантовать в больше, чем 1 бит?). Attention квантуют в 4 бита, экспертов в 2 и 3 бита.

Эксперименты

Рассматриваются две конфигурации системы - условный Colab с бесплатной GPU и хороший геймерский ПК. Для GPU c 12Gb VRAM можно держать кэш из 2 экспертов, а для 16Gb - влезают 4.

Чем больше размер кэша - тем больше hit-ratio (вероятность того, что эксперт лежит в кэше) и точность удачного speculative loading.
Выбранная схема квантования более менее оптимальна, так как квантизация attention в меньшее число бит уже сильно просаживает качество, а 4 бита некритично хуже fp16. Конечные конфигурации в 4-5 раз меньше по размеру базовой модели. Некоторая просадка по качеству на бенчмарках есть, но приемлемая.

Все компоненты метода - LRU-cache, пре-подгрузка экспертов заметно ускоряют инференс. Итоговое ускорение порядка 2-2.5 раз на RTX3060/RTX3080 и до 3.5 раз на T4 по сравнению с наивной реализацией в accelerate. Для прикола замеряют еще и на A100, куда влезает квантованная модель и так (для сравнения).

Вывод

Простой и разумный подход с практически полезным результатом. По всей видимости, умный оффлоадинг - одно из самых, если не самое перспективное направление, ускорения моделях на пользовательских устройствах. Однако, букву “О” гармонично вставить в название канала не получается. Кажется, что результаты работы можно улучшить за счет более сильного алгоритма квантизации экспертов, уменьшив просадку качества.

BY КПД