swiftLLM

SwiftLLM

Этот проект все еще находится в стадии разработки. Некоторые функции могут быть еще не реализованы, а документация может быть неполной.

Крошечная, но мощная система вывода LLM, предназначенная для исследовательских целей.

Производительность, эквивалентная vLLM, всего с 2 тыс. строк кода (2% от vLLM).

Почему SwiftLLM

Существует множество платформ с открытым исходным кодом для обслуживания LLM, включая HuggingFace Transformers, vLLM, LightLLM, DistServe и DeepSpeed-MII. Почему SwiftLLM?

Причина в том, что эти структуры предназначены для производства , а не для исследований . Они оснащены многочисленными функциями, такими как поддержка более 100 моделей, различные аппаратные поддержки, LoRA, квантование, мультимодальность, кэширование префиксов, поиск луча и так далее. Несмотря на то, что это универсальное решение для производства, их кодовая база слишком велика и сложна для понимания и изменения (например, vLLM имеет более 100 тысяч строк кода), что затрудняет их использование в исследовательских целях. Кроме того, их историческое бремя также является проблемой.

SwiftLLM представляет собой крошечную, но мощную систему вывода LLM, предназначенную для исследовательских целей . «Маленький» означает, что он сохраняет только те функции, которые необходимы для исследований, «мощный» означает, что он не имеет компромиссов в производительности, и, наконец, «быстрый» означает, что его легко понять и изменить. Несмотря на поддержку базовых функций (см. список ниже) и возможность достижения производительности, эквивалентной vLLM, кодовая база SwiftLLM составляет менее 2 тысяч строк кода (~ 2% от vLLM), написанного на Python и OpenAI Triton (DSL для написания Ядра CUDA), что упрощает чтение, изменение, отладку, тестирование, расширение и позволяет легко интегрировать ваши новые и блестящие исследовательские идеи.

Список функций

В настоящее время SwiftLLM поддерживает следующие функции:

• Итерационное планирование и выборочная пакетная обработка (предложено в Orca)
• PagedAttenton (предлагается в vLLM, статья)
• Модели LLaMA/LLaMA2/LLaMA3 (ссылка) и их варианты
• Совмещенное предварительное заполнение и декодирование (предложено в SARATHI)
• Вспышка внимания (предлагается в FlashAttention и FlashAttention-2)
• Постраничное внимание v2 (также называемое Flash-декодированием, предлагается здесь)

В будущем мы планируем добавить поддержку следующих функций:

• Тензорный параллелизм и конвейерный параллелизм

Чтобы сохранить небольшую кодовую базу, мы не будем поддерживать следующие функции. Если вы хотите использовать их в своем исследовательском проекте, возможно, вам придется реализовать их самостоятельно:

• Квантование
• ЛоРА
• Мультимодальный
• Модели, не соответствующие архитектуре LLaMA.
• Методы выборки, кроме жадной выборки
• Аппаратное обеспечение поддерживает графические процессоры, отличные от NVIDIA (но его можно легко перейти на другое оборудование, если OpenAI Triton поддерживает их)

Помните, что SwiftLLM НЕ является универсальным решением для производства. Рекомендуется рассматривать его как «фундамент» для вашего исследовательского проекта, и вам, возможно, придется реализовать некоторые функции самостоятельно. Мы призываем вас, мой дорогой исследователь, прочитать код, понять его, изменить и расширить в соответствии с вашими исследовательскими потребностями.

Архитектура

Архитектуру SwiftLLM можно разделить на две основные части: плоскость управления и плоскость данных .

Короче говоря, плоскость управления решает, «что вычислять» или «как планировать», тогда как плоскость данных решает, «как вычислять» или «как реализовать», и выполняет конкретные вычисления. Они работают по принципу «главный-работник»: плоскость управления действует как мастер, который выполняет планирование и координацию высокого уровня и отправляет задания в плоскость данных, которая действует как рабочий, выполняющий вычисления низкого уровня.

Код плоскости управления находится в каталоге swiftllm/server и включает такие компоненты, как Engine , Scheduler , сервер API и TokenizationEngine . Код плоскости данных находится в каталоге swiftllm/worker , включая описания графа вычислений (в swiftllm/worker/model.py ), реализацию слоев в модели (в swiftllm/layers ) и ядра OpenAI Triton ( вы можете представить «ядра» как функции, выполняемые на графическом процессоре) (в swiftllm/kernels ).

В качестве примера возьмем игрушечный API-сервер (расположенный в swiftllm/server/api_server.py ):

• При запуске он использует EngineConfig для создания Engine .
• После этого движок инициализируется через Engine.initialize , где он создает Scheduler , TokenizationEngine и набор (на данный момент только один, поскольку тензорный параллелизм не поддерживается) воркеров. Затем он дает команду рабочему выполнить profile_num_blocks для расчета количества блоков графического процессора, после чего движок дает команду всем рабочим процессам выделить свой кэш KV и подкачку KV.
• Наконец, цикл событий активируется через Engine.start_all_event_loops . На каждом этапе цикла механизм запрашивает у планировщика следующий пакет запросов для вычисления, дает команду исполнителю выполнить замену, а затем отправляет пакет исполнителю для вычисления.
• Сервер API слушает запросы пользователей и взаимодействует с движком для их выполнения.

В настоящее время плоскость управления ( Engine ) и плоскость данных ( LlamaModel ) находятся на одном узле. После реализации тензорного параллелизма/конвейерного параллелизма плоскость данных может быть распределена по нескольким узлам.

Как использовать

Мы предлагаем два способа использования SwiftLLM: использовать одновременно плоскость управления и плоскость данных или использовать только плоскость данных.

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

Стройте и запускайте

Сначала давайте настроим среду:

• Вы можете начать с чистой среды conda с Python >= 3.9 или использовать существующую. Не забудьте активировать среду conda, если вы ее используете.
• Установите ПиТорч. Обратите внимание, чтобы выбрать правильную версию в зависимости от вашего оборудования.
• Установите packaging через pip install packaging

И далее идет установка:

• Клонируйте этот репозиторий через git clone https://github.com/interestingLSY/swiftLLM.git
• cd в репозиторий ( cd swiftLLM ) и установите другие зависимости с помощью pip install -r requirements.txt .
• PyTorch может установить для вас стабильную версию OpenAI Triton. Если вы хотите использовать ночную версию для достижения максимальной производительности, но с потенциальными проблемами, вы можете удалить ее и установить ночную версию.
• Запустите pip install -e . для установки SwiftLLM в вашу среду.
• Установите некоторые привязки C с помощью pip install -e csrc

Вот несколько примеров:

• В настоящее время SwiftLLM не поддерживает автоматическую загрузку весов из HuggingFace. Возможно, вам придется сначала клонировать или загрузить веса моделей из HuggingFace. Поддерживаются форматы .bin и .safetensors . Предположим, что вес вашей модели хранится в /data/to/weight/ .
• В качестве примера автономного обслуживания вы можете попробовать python3 examples/offline.py --model-path /data/to/weight . В этом примере используется только плоскость данных. Если вы планируете использовать SwiftLLM без плоскости управления, это хорошая отправная точка.
• Для примера онлайн-обслуживания, использующего Engine , вы можете попробовать python3 examples/online.py --model-path /data/to/weight . Это отличный пример, если вы планируете использовать как плоскость управления, так и плоскость данных.
• Более сложный пример можно найти в swiftllm/server/api_server.py . Он запускает сервер API и предоставляет интерфейс, подобный vLLM, для онлайн-обслуживания.

Производительность

Несмотря на свои крошечные размеры (крошечные тоже могут быть очаровательными!), SwiftLLM не идет на компромисс с производительностью. Мы оценили SwiftLLM в нескольких сценариях и продемонстрировали, что SwiftLLM может достичь эквивалентной или даже лучшей производительности по сравнению с vLLM.

Единая прямая операция

Первый сценарий — это «одна прямая операция», когда мы снабжаем модель пакетом входных данных и позволяем ей генерировать один выходной токен (эквивалент одной «прямой» операции). Это основная операция вывода LLM (как онлайн, так и офлайн), поэтому ее производительность имеет решающее значение.

Здесь мы используем модель LLaMA-3 7B с графическим процессором NVIDIA A100 80G PCIE/RTX 4090 с точностью FP16. Результаты показаны ниже (чем ниже, тем лучше):

Видно, что SwiftLLM может достичь производительности, эквивалентной (или даже превосходящей) производительности vLLM при тех же настройках.

Онлайн-обслуживание

Второй сценарий — «онлайн-обслуживание», когда мы запускаем сервер API, берем образцы запросов из реального набора данных и позволяем модели генерировать завершения. Это сценарий, в котором LLM используется в реальных приложениях, таких как чат-боты или автодополнение кода.

Здесь мы используем набор данных ShareGPT для выборки подсказок и используем процесс Пуассона с разными лямбда-выражениями для имитации различных скоростей поступления запросов. Результаты показаны ниже (чем ниже, тем лучше):

Видно, что на A100 80G PCIE SwiftLLM может достичь производительности, эквивалентной vLLM, а на RTX 4090 SwiftLLM значительно превосходит vLLM (в основном из-за того, что наша плоскость управления имеет меньшие накладные расходы).