FlexGen на практике: получится ли запустить тяжелую модель без мощной видеокарты

От таких штук нас обещает избавить FlexGen при работе с текстовыми моделями.

Ранее я рассказывал про создание виртуальной подруги на основе текстовой модели GPT-J, которую можно развернуть локально. Ожидаемый недостаток этой идеи — большие «аппетиты» текстовых моделей по отношению к вычислительным ресурсам. Собственно, мой последний эксперимент состоялся только благодаря карточке A100 в сервере.

В конце февраля группа исследователей опубликовала FlexGen — решение для запуска масштабных текстовых моделей в среде с ограниченными ресурсами. В теории оно должно помочь обойтись ресурсами персонального ПК, без добавления мощных видеокарт. Под катом расскажу о FlexGen и его применимости в существующем проекте.

О всех экспериментах пишу в личном Telegram-канале. Подписывайтесь, если хотите узнавать о них в зачатке, пока все это не разрослось в очередную большую публикацию на Хабре.

Сперва ознакомимся с описанием и целями проекта. FlexGen в первую очередь нацелен на повышение пропускной способности путем разгрузки (offloading) и сжатия при использовании пакетной обработки (batch processing). Это значит, что данный инструмент полностью раскроется в задачах, которые подразумевают обработку большого объема данных, где важно суммарное время обработки, а не время обработки каждого запроса.

Инструмент специфичный и появился недавно, поэтому имеет следующие ограничения.

• На момент написания статьи поддерживается только модели семейства OPT (продукт компании Meta, признанной экстремисткой организацией в России).
• Текстовая модель ограничивается только количеством сгенерированных токенов.

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

FlexGen можно установить двумя способами: из репозитория или через пакетный менеджер. Версия из репозитория более полная и включает в себя необязательные скрипты для выполнения тестов производительности. Я решил установить через пакетный менеджер:

pip install flexgen

Сразу после установки flexgen готов к использованию. Запрошенная модель будет скачана автоматически, если ее нет. Во время работы flexgen выгружает части моделей по пути ~/flexgen_offload_dir, и это требует быстрого взаимодействия с дисковой подсистемой. Поэтому рекомендуется выносить этот каталог на быстрые локальные диски.

Перед запуском разберем параметры flexgen. На данный момент доступны следующие:

• Имя модели. Разрешенные модели явно прописаны в коде flexgen.
• Каталог для хранения данных моделей.
• Каталог для временного хранения разгрузки (offload).
• Сжатие весов и кэша.
• Процент распределения модели по GPU и CPU.

Последний параметр состоит из шести чисел, показывающих следующее:

• сколько процентов модели на GPU,
• сколько процентов модели на CPU,
• сколько процентов кэша на GPU,
• сколько процентов кэша на CPU,
• сколько процентов активаций на GPU,
• сколько процентов активаций на CPU.

По умолчанию этот параметр равен 100 0 100 0 100 0, то есть модель полностью помещается в VRAM. Первые два числа определяют объем занимаемой видеопамяти (VRAM) и оперативной памяти (RAM) соответственно. Так, если указать 50 10 100 0 100 0, то в VRAM будет загружена только половина модели. Параметры 5 и 6 позволяют указать процент вычислений на GPU и на CPU.

Запускаем тест:

$ python3 -m flexgen.apps.completion --model facebook/opt-1.3b --compress-weight
Initialize...
Generate...
Outputs:
----------------------------------------------------------------------
0: Question: Where were the 2004 Olympics held?
Answer: Athens, Greece
Question: What is the longest river on the earth?
Answer: The Nile
Question: What is the longest river on the planet?
Answer: The Yangtze
Question: What is the biggest lake in the
----------------------------------------------------------------------
1: Extract the airport codes from this text.
Text: "I want a flight from New York to San Francisco."
Airport codes: JFK, SFO.
Text: "I want you to book a flight from Phoenix to Las Vegas."
Airport codes: PHX, LVG.
The.NET
Extract the airport codes from this text.
Text: "I want a flight from New
----------------------------------------------------------------------
Shutdown...

В качестве моделей можно выбирать любую из доступных OPT-моделей — от 125M до 175B. Обратите внимание, что самая большая модель на 175 миллиардов параметров доступна только по прямому запросу.

Данный тест не позволяет получить метрики по времени работы модели, но позволяют оценить корректность работы текстовой модели.

$ python3 -m flexgen.apps.completion --model facebook/opt-66b --percent 25 10 100 0 100 0 --compress-weight
Initialize...
Generate...
Outputs:
----------------------------------------------------------------------
0: Question: Where were the 2004 Olympics held?
Answer: Athens, Greece
Question: What is the longest river on the earth?
Answer: The Amazon River
Question: How many continents are there?
Answer: 6
Question: What is the capital of California?
Answer: Sacramento
Question
----------------------------------------------------------------------
1: Extract the airport codes from this text.
Text: "I want a flight from New York to San Francisco."
Airport codes: JFK, SFO.
Text: "I want you to book a flight from Phoenix to Las Vegas."
Airport codes: PHX, LAS.
Text: "I want to fly from Miami to Honolulu."
Airport codes: MIA, HNL.
Text:
----------------------------------------------------------------------
Shutdown...

Разница в производительности и качестве моделей видна и без точных метрик. Opt-1.3b работает секунды и выдает два неверных ответа на вопрос о самой длинной реке на планете, а opt-66 работает минуты, но дает корректные ответы.

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

Веб-интерфейс со странной генерацией от flexgen.

Пользователь oobabooga впечатлился веб-интерфейсом для txt2img-модели Stable Diffusion за авторством AUTOMATIC1111 и задался целью сделать аналогичный веб-интерфейс для текстовых моделей. Так появился text-generation-webui, который работает со множеством текстовых моделей, в том числе поддерживает FlexGen.

Text-generation-webui находится в активной разработке и, к сожалению, пока не так удобен, как его идейный вдохновитель. После нажатия на кнопку в интерфейсе вы скорее всего не увидите изменений, особенно в случае с медлительным flexgen. Это вызывает подсознательное желание сделать еще один клик.

Автоматическая установка в один клик доступна только пользователям Windows, для остальных операционных систем необходимо выполнить последовательность действий с установкой виртуального окружения Conda. Более подробно — в README.

Когда все установлено, можно скачивать интересующие модели и запускать их. Скачать модель можно встроенным скриптом:

python download-model.py facebook/opt-6.7b

Обратите внимание, что facebook/opt-6.7b — это идентификатор модели на сайте huggingface.co. Если ошибиться в написании или выбрать несуществующую модель, то будет страшная ошибка, которая предложит зарегистрироваться на сайте.

После загрузки всех интересующих моделей можно запустить веб-сервер:

python server.py --share

По умолчанию веб-интерфейс доступен по адресу 127.0.0.1:7860. Флаг share открывает доступ через сервис gradio и генерирует ссылку, доступную через интернет. В этом режиме можно поиграть со скачанными текстовыми моделями, но включить FlexGen не получится.

Для активации возможностей FlexGen необходимо прочесть соответствующую страницу документации и произвести магические манипуляции. Сперва конвертировать скачанную модель в формат NumPy:

python convert-to-flexgen.py models/facebook_opt-6.7b

Теперь можно запустить сервер с явным указанием модели и флагом flexgen. Напомню, что «виртуальную девушку я запускал на сервере GL70-1 с видеокартой NVIDIA A100:

python server.py --model opt-6.7b --share --flexgen

После выполнения этой команды появится ошибка о недоступности модели и предложение зарегистрироваться на сайте huggingface. В этот момент появляются первые сомнения о пользе этого веб-интерфейса и flexgen. Проблема заключается в особенностях обработки параметра model. При флаге flexgen параметр model принимает имя модели без префикса facebook/, что некорректно обрабатывается в других местах.

На момент написания этого текста text-generation-webui остановился на коммите 91745f63c35480ea3ce6b66e979dbd7ce364171c. Чтобы починить сервер, нужно внести изменения в файле modules/models.py. Необходимо заменить {shared.args.model_dir} на facebook/ (метод from_pretrained) в строчках 59 и 248, так как файл конфигурации и токенайзер подхватываются с сайта, а модель – локально.

Однако на этой версии генерация выглядит крайне странно: генерация переписывает сама себя. Для решения этой проблемы можно откатиться на версию двухнедельной давности, на коммит 1911504f8205bfae5203185d89d7860a0f930974, но тогда придется исправить еще и modules/text_generation.py и закомментировать появление стоп-токена на 180 строчке, ведь flexgen не поддерживает стоп-последовательности.

Кажется, на этом эксперимент можно завершать.

На текущий момент FlexGen не похож на инструмент, пригодный для практического применения. Пока хорошим решением для тех, кто хочет поработать с моделями, но не хочет тратиться на покупку видеокарты, может стать аренда облачного сервера с GPU. Вы можете взять прерываемый сервер со скидкой и платить только за потребленные ресурсы. Из плюсов — в Selectel есть специальная ОС с установленными драйверами для работы с GPU. Можно сразу использовать.

В остальном предлагаю следить за развитием проекта. Возможно, в будущем FlexGen даст возможность прикоснуться к «тяжелым» текстовым моделям локально тем, у кого нет доступа к современным видеокартам.