PagedWeight: метод экономии памяти GPU для MoE-моделей

Исследователи из Университета Иллинойса представили PagedWeight — метод, который позволяет освобождать память GPU во время работы больших MoE-моделей. В тестах на Qwen, Mixtral и Gemma экономия памяти достигла 72%, а скорость генерации выросла до 1,94 раза.

Исследователи из Университета Иллинойса представили PagedWeight — метод освобождения памяти GPU во время работы больших MoE-моделей. Такие модели состоят из множества специализированных блоков («экспертов»), но для каждого следующего слова активируют лишь несколько из них, при этом веса всех экспертов обычно всё равно хранятся на GPU. PagedWeight делит веса на небольшие блоки, точность которых можно менять независимо. Когда памяти не хватает, система выбирает блоки, наименее чувствительные к сжатию, и переносит в обычную оперативную память часть битов, нужных для их более точной записи. На GPU остаётся сжатая версия, с которой модель продолжает работать. Когда место освобождается, недостающие биты можно загрузить обратно. Система также учитывает, как часто используется каждый эксперт, и старается осторожнее сжимать самые востребованные. Метод проверили на трёх MoE-моделях: Qwen, Mixtral и Gemma. В тестах на длинных текстах Qwen занял 9,86 ГБ вместо 35,25 ГБ и сохранил тот же средний результат — экономия составила 72%. В отдельном тесте система генерировала до 1,94 раза больше токенов в секунду, чем обычный 16-битный режим. Публичный код авторы пока не выпустили, они поделились только подходом и первыми экспериментами.

Источники 1

PagedWeight: метод экономии памяти GPU для MoE-моделей Нейроканал

Чем длиннее разговор с нейронкой, тем больше памяти видеокарты она тратит. Модели нужно хранить служебные данные обо всех предыдущих словах, и постепенно для них начинает не хватать места.

Как вы поняли, далее речь пойдёт про свежую статью с arxiv, давненько не проверял что там интересного.

Так вот, исследователи из Университета Иллинойса представили PagedWeight — способ освобождать память прямо во время работы больших MoE-моделей. Такие модели состоят из множества специализированных блоков, или «экспертов», но для каждого следующего слова включают только несколько из них. При этом веса всех экспертов обычно всё равно хранятся на GPU.

PagedWeight делит веса на небольшие блоки, точность которых можно менять независимо. Когда памяти не хватает, система выбирает блоки, наименее чувствительные к сжатию, и переносит в обычную оперативную память часть битов, нужных для их более точной записи. На GPU остаётся сжатая версия, с которой модель продолжает работать. Когда место освобождается, недостающие биты можно загрузить обратно.

Система также учитывает, как часто используется каждый эксперт, и старается осторожнее сжимать самые востребованные. Отсюда и название PagedWeight: блоками весов управляют примерно как отдельными страницами памяти.

Метод проверили на трёх MoE-моделях: Qwen, Mixtral и Gemma. В тестах на длинных текстах Qwen занял 9,86 ГБ вместо 35,25 ГБ и сохранил тот же средний результат — экономия составила 72%. В отдельном тесте система генерировала до 1,94 раза больше токенов в секунду, чем обычный 16-битный режим.

Публичный код авторы не выпустили, пока только поделились подходов и первыми экспериментами.

💻 Если пока непонятно, о чём здесь вообще речь и откуда берутся такие результаты, начать разбираться можно на курсе «Основы глубокого обучения нейросетей» в Практикуме PRO. В программе — устройство и обучение нейросетей, основные архитектуры, трансформеры и LLM.

@neuro_channel