Ускорение генерации текста с помощью vLLM

Библиотека vLLM обеспечивает значительное ускорение генерации текста с помощью LLM до 20 раз на больших данных. Используя технологии квантизации, PagedAttention и асинхронные запросы, она эффективно оптимизирует процесс создания моделей языка. Этот инструмент полезен для исследований и разработки, с учётом возможного снижения качества ответов при установлении большого времени ожидания
Новости 2024 06 03

Введение

Большие языковые модели (LLM) играют ключевую роль в современных задачах обработки естественного языка, обеспечивая высокую точность и качество генерации текста. Однако процесс инференса, особенно на больших наборах данных, часто становится узким местом из-за значительных временных и вычислительных затрат. В данной статье рассматривается использование библиотеки vLLM для ускорения генерации текста в LLM вплоть до 20 раз за счет применения методов оптимизации, таких как PagedAttention и непрерывный батчинг.

Библиотека vLLM

Библиотека vLLM предоставляет инструменты для эффективного развертывания языковых моделей, поддерживающих OpenAI-совместимый API. Она включает поддержку квантизированных моделей (GPTQ, AWQ и SqueezeLLM), оптимизацию памяти с помощью PagedAttention, оптимизированные CUDA ядра, непрерывный батчинг и другие функции. vLLM совместима с большинством популярных моделей из HuggingFace.

Эндпоинты API

OpenAI-совместимый сервер в vLLM реализует следующие эндпоинты:
  • list models
  • ChatCompletion
  • Completion

Развертывание API может быть выполнено следующим образом:

 
from openai import OpenAI
client = OpenAI()

num_stories = 10
prompt = "Я пошел в магазин и купил"

for _ in range(num_stories):
    response = client.completions.create(
        model=model_name,
        prompt=prompt,
        max_tokens=100,
    )
    print(prompt + response.choices[0].text)

Для статического батчинга:

 
from openai import OpenAI
client = OpenAI()

num_stories = 10
prompts = ["Я пошел в магазин и купил"] * num_stories

response = client.completions.create(
    model=model_name,
    prompt=prompts,
    max_tokens=100,
)

stories = [""] * len(prompts)
for choice in response.choices:
    stories[choice.index] = prompts[choice.index] + choice.text

for story in stories:
    print(story)

Ускорение генерации: Инференс в LLM

Инференс в LLM происходит следующим образом: на вход подается последовательность токенов, затем итеративно генерируются следующие токены до достижения eos-токена или максимальной длины последовательности. Время обработки входного промпта примерно равно времени генерации выходной последовательности, поскольку вычисления для механизма внимания остаются постоянными на протяжении всей генерации.

Влияние батчинга

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

Статический батчинг

Метод статического батчинга предполагает постоянный размер батча до полной генерации каждой последовательности. Это повышает пропускную способность, но память GPU используется неэффективно из-за времени генерации самой длинной последовательности.

Непрерывный батчинг

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

PagedAttention

Механизм PagedAttention также способствует значительному ускорению инференса. Во время авторегрессионного декодирования для каждого входного токена генерируются тензоры K и V для слоя внимания, которые затем сохраняются в памяти GPU (KV-кэш). Память, используемая для KV-кэша, может достигать значительных объемов (до 1,7 ГБ для LLaMA-13B).

PagedAttention оптимизирует использование памяти, разбивая KV-кэш каждой последовательности на блоки, которые хранятся в несмежных пространствах памяти. Это позволяет обрабатывать большие батчи последовательностей, увеличивать загрузку GPU и существенно повышать скорость инференса.

Пример асинхронной генерации

Для использования непрерывного батчинга и ускорения генерации данных с помощью vLLM был применен асинхронный подход с использованием библиотеки asyncio:

 
import openai
import asyncio

async def generate_answers(prompt):
    completion = await openai.ChatCompletion.acreate(
        model="model_name",
        messages=[{"role": "user", "content": prompt}],
        max_tokens=1024,
        request_timeout=10000
    )
    return completion["choices"][0].get("message").get("content")

async def main(data):
    tasks = []
    for idx, prompt in enumerate(data):
        task = asyncio.create_task(generate_answers(prompt))
        tasks.append(task)

    answers = await asyncio.gather(*tasks)
    return answers

results = asyncio.run(main(data))

Результаты

Использование vLLM привело к многократному сокращению времени генерации данных. На видеокарте A100 40 ГБ время генерации 100 инструкций сократилось с 6 минут 10 секунд до 1 минуты, а для 2000 инструкций — с 2 часов 50 минут до 5-10 минут.

Заключение

Методы непрерывного батчинга и PagedAttention, реализованные в vLLM, позволяют значительно ускорить процесс инференса в больших языковых моделях. Несмотря на возможное падение качества генерации некоторых ответов и необходимость выставления большого времени ожидания, эти методы обеспечивают существенное ускорение экспериментов с LLM.

Надеемся, что статья окажется полезной в вашей работе с LLM.
Поиск