FlexGen — движок для запуска AI-ботов, похожих на ChatGPT, на системах с одним GPU

Группа исследователей из Стендфордского университета, Калифорнийского университета в Беркли, Швейцарской высшей технической школы Цюриха, Высшей школы экономики, университета Карнеги — Меллона, а также компаний Yandex и Meta, опубликовала исходные тексты движка для выполнения крупных языковых моделей на системах с ограниченными ресурсами. Например, движок предоставляет возможность создания функциональности, напоминающей ChatGPT и Copilot, через выполнение готовой натренированной модели OPT-175B, охватывающей 175 миллиардов параметров, на обычном компьютере с игровой видеокартой NVIDIA RTX3090, оснащённой 24GB видеопамяти. Код написан на языке Python, использует фреймворк PyTorch и распространяется под лицензией Apache 2.0.

В состав входит пример скрипта для создания ботов, позволяющего загрузить одну из публично доступных языковых моделей и сразу начать общение (например, выполнив команду «python apps/chatbot.py —model facebook/opt-30b — -percent 0 100 100 0 100 0»). В качестве базовой предлагается использовать опубликованную Facebook крупную языковую модель, обученную на коллекциях BookCorpus (10 тысяч книг), CC-Stories, Pile (OpenSubtitles, Wikipedia, DM Mathematics, HackerNews и т.п.), Pushshift.io (на основе данных Reddit) и CCNewsV2 (архив новостей). Модель охватывает около 180 миллиардов токенов (800 ГБ данных). На тренировку модели было потрачено 33 дня работы кластера с 992 GPU NVIDIA A100 80GB.

При выполнении модели OPT-175B на системе с одним GPU NVIDIA T4 (16ГБ) движок FlexGen продемонстрировал производительность до 100 раз опережающую ранее предлагавшиеся решения, что делает использование крупных языковых моделей более доступными и позволяет запускать их на системах без специализированных ускорителей. При этом FlexGen может масштабироваться для распараллеливания вычислений при наличии нескольких GPU. Для сокращения размерам модели дополнительно применяется собственная схема сжатия параметров и механизм кэширования моделей.

В настоящее время FlexGen поддерживает только языковые модели OPT, но в дальнейшем разработчики также обещают добавить поддержку моделей BLOOM (176 миллиардов параметров, поддерживает 46 языков и 13 языков программирования), CodeGen (может генерировать код на 22 языках программирования) и GLM. Пример диалога с ботом на базе FlexGen и модели OPT-30B:

Human: What is the name of the tallest mountain in the world?

Assistant: Everest.

Human: I am planning a trip for our anniversary. What things can we do?

Assistant: Well, there are a number of things you can do for your anniversary. First, you can play cards. Second, you can go for a hike. Third, you can go to a museum.

Источник: opennet.ru