Оптимизация вашего LLM в производстве

Оптимизация LLM в производстве

Примечание: Этот блог-пост также доступен в качестве страницы документации по Transformers.

Большие языковые модели (LLM), такие как GPT3/4, Falcon и LLama, быстро развиваются в своей способности решать задачи, связанные с человеком, и становятся неотъемлемыми инструментами в современных знаниевых отраслях. Однако развертывание этих моделей в реальных задачах остается сложной задачей:

• Чтобы обладать близкими к человеческими возможностями понимания и генерации текста, LLM в настоящее время требуют состоять из миллиардов параметров (см. Kaplan et al, Wei et. al). Это влечет за собой увеличение требований к памяти для вывода.
• Во многих реальных задачах LLM требуется обширная контекстная информация. Это требует способности модели управлять очень длинными входными последовательностями во время вывода.

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

• Преимущества LLM Преобразование поиска в электронной коммерции
• Обучение вариационного автоэнкодера для обнаружения аномалий с использованием TensorFlow
• GraphReduce Использование графов для абстракции инженерии признаков

В этом блог-посте мы рассмотрим наиболее эффективные техники на момент написания этого блог-поста для решения этих проблем при эффективном развертывании LLM:

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

2. Flash Attention: Flash Attention – это вариант алгоритма внимания, который не только обеспечивает более эффективный расход памяти, но также обеспечивает увеличение эффективности за счет оптимизированного использования памяти GPU.

3. Архитектурные инновации: Учитывая, что LLM всегда развертывается одинаковым образом во время вывода, а именно авторегрессивной генерацией текста с длинным контекстом ввода, были предложены специализированные модельные архитектуры, которые позволяют более эффективный вывод. Здесь наиболее важными новшествами в архитектуре модели являются Alibi, Rotary embeddings, Multi-Query Attention (MQA) и Grouped-Query-Attention (GQA).

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

1. Использование возможностей пониженной точности

Требования к памяти LLM могут быть наилучшим образом поняты, если рассматривать LLM как набор матриц весов и векторов, а текстовый ввод как последовательность векторов. В дальнейшем термин “веса” будет использоваться для обозначения всех матриц весов и векторов модели.

На момент написания этого поста LLM состоят как минимум из нескольких миллиардов параметров. Каждый параметр представляет собой десятичное число, например, 4.5689, которое обычно хранится в формате float32, bfloat16 или float16. Это позволяет нам легко вычислить объем памяти, необходимый для загрузки LLM в память:

Загрузка весов модели, имеющей X миллиардов параметров, требует примерно 4 * X гигабайт оперативной памяти в формате float32

В настоящее время модели редко тренируются в полной точности float32, обычно используется точность bfloat16 или, реже, float16. Поэтому правило превращается в следующее:

Загрузка весов модели, имеющей X миллиардов параметров, требует примерно 2 * X гигабайт оперативной памяти в формате bfloat16/float16

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

Давайте рассмотрим несколько примеров, сколько VRAM примерно занимает загрузка модели в формате bfloat16:

• GPT3 требует 2 * 175 ГБ = 350 ГБ VRAM
• Bloom требует 2 * 176 ГБ = 352 ГБ VRAM
• Llama-2-70b требует 2 * 70 ГБ = 140 ГБ VRAM
• Falcon-40b требует 2 * 40 ГБ = 80 ГБ VRAM
• MPT-30b требует 2 * 30 ГБ = 60 ГБ VRAM
• bigcode/starcoder требует 2 * 15.5 = 31 ГБ VRAM

На момент написания этого документа, самый большой графический процессор на рынке – A100, предлагающий 80 ГБ VRAM. Большинство перечисленных моделей требуют более 80 ГБ только для загрузки и, следовательно, необходимы тензорная параллельность и/или конвейерная параллельность.

🤗 Transformers не поддерживает тензорную параллельность из коробки, так как требует, чтобы архитектура модели была написана определенным образом. Если вам интересно написание моделей в дружественном к тензорной параллельности стиле, не стесняйтесь ознакомиться с библиотекой текстового поколения-вывода.

Наивная конвейерная параллельность поддерживается из коробки. Для этого просто загрузите модель с помощью device="auto", что автоматически разместит различные слои на доступных графических процессорах, как объясняется здесь. Однако обратите внимание, что, хотя это очень эффективно, эта наивная конвейерная параллельность не решает проблемы простоя графического процессора. Для этого требуется более продвинутая конвейерная параллельность, как объясняется здесь.

Если у вас есть доступ к узлу с 8 графическими процессорами A100 по 80 ГБ, вы можете загрузить BLOOM следующим образом

!pip install transformers accelerate bitsandbytes optimum

# from transformers import AutoModelForCausalLM

# model = AutoModelForCausalLM.from_pretrained("bigscience/bloom", device_map="auto", pad_token_id=0)

Используя device_map="auto", слои внимания будут равномерно распределены по всем доступным графическим процессорам.

В этом блокноте мы будем использовать bigcode/octocoder, так как он может работать на одном чипе графического процессора A100 с объемом 40 ГБ. Обратите внимание, что все оптимизации памяти и скорости, которые мы будем применять в дальнейшем, также применимы к моделям, требующим тензорной или модельной параллельности.

Поскольку модель загружена с точностью bfloat16, используя наше эмпирическое правило выше, мы ожидаем, что требования к памяти для выполнения вывода с использованием bigcode/octocoder составят около 31 ГБ VRAM. Давайте попробуем.

Сначала мы загружаем модель и токенизатор, а затем передаем их обоих в объект конвейерной обработки Transformers.

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", torch_dtype=torch.bfloat16, device_map="auto", pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained("bigcode/octocoder")

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

prompt = "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer:"

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

Вывод:

Вот функция на Python, которая преобразует байты в гигабайты:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nЭта функция принимает единственный

Отлично, теперь мы можем непосредственно использовать результат для преобразования байтов в гигабайты.

def bytes_to_giga_bytes(bytes):
  return bytes / 1024 / 1024 / 1024

Давайте вызовем torch.cuda.max_memory_allocated, чтобы измерить пиковое выделение памяти графического процессора.

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

Вывод:

29.0260648727417

Достаточно близко к нашему грубому расчету! Мы видим, что число не совсем точное, так как при переходе от байтов к килобайтам требуется умножение на 1024 вместо 1000. Поэтому формулу в грубом приближении также можно понимать как вычисление “не более X ГБ”. Обратите внимание, что если бы мы попытались запустить модель с полной точностью float32, требовалось бы целых 64 ГБ VRAM.

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

Если вы не уверены, в каком формате хранятся веса модели на Hub, всегда можно посмотреть конфигурацию чекпоинта в "torch_dtype", например, здесь. Рекомендуется установить модель в тот же тип точности, что указан в конфигурации при загрузке с помощью from_pretrained(..., torch_dtype=...), за исключением случая, когда оригинальный тип – float32, в котором можно использовать как float16, так и bfloat16 для вывода.

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

del pipe
del model

import gc
import torch

def flush():
  gc.collect()
  torch.cuda.empty_cache()
  torch.cuda.reset_peak_memory_stats()

Давайте вызовем ее сейчас для следующего эксперимента.

flush()

В последней версии библиотеки accelerate вы также можете использовать вспомогательный метод под названием release_memory()

from accelerate.utils import release_memory
# ...

release_memory(model)

Теперь что, если ваш GPU не имеет 32 ГБ VRAM? Было обнаружено, что веса модели могут быть квантованы до 8-бит или 4-бит без значительной потери производительности (см. Dettmers и др.). Модель может быть квантована даже до 3 или 2 бит с приемлемой потерей производительности, как показано в недавней статье GPTQ 🤯.

Не вдаваясь во множество деталей, схемы квантования стремятся уменьшить точность весов, сохраняя при этом точность результатов вывода модели как можно более точной (т.е. как можно ближе к bfloat16). Обратите внимание, что квантование особенно хорошо работает для генерации текста, поскольку все, что нам важно, – это выбор набора наиболее вероятных следующих токенов, и нам не очень важны точные значения распределения логитов следующего токена. Все, что имеет значение, – это то, чтобы распределение логитов следующего токена оставалось примерно таким же, чтобы операция argmax или topk давала те же результаты.

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

• 1. Квантовать все веса до целевой точности
• 1. Загрузить квантованные веса и передать последовательность векторов входных данных в формате bfloat16
• 1. Динамически деквантовать веса в формат bfloat16 для выполнения вычислений с входными векторами в формате bfloat16
• 1. Квантовать веса снова до целевой точности после вычислений со входными данными.

Вкратце, это означает, что умножение матриц вход-вес, где X – входы, W – матрица весов, Y – выходы:

Y=X∗W Y = X * W Y=X∗W

изменяются на

Y=X∗де-квантовать(W);квантовать(W) Y = X * \text{де-квантовать}(W); \text{квантовать}(W) Y=X∗де-квантовать(W);квантовать(W)

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

Следовательно, время вывода часто не уменьшается при использовании квантованных весов, а, наоборот, увеличивается. Достаточно теории, давайте попробуем! Чтобы квантовать веса с помощью Transformers, вам нужно убедиться, что установлена библиотека bitsandbytes.

# !pip install bitsandbytes

Затем мы можем загрузить модели в квантовании 8-бит, просто добавив флаг load_in_8bit=True к from_pretrained.

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_8bit=True, pad_token_id=0)

Теперь давайте снова запустим наш пример и измерим использование памяти.

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

Результат:

Вот функция на Python, которая преобразует байты в гигабайты:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nЭта функция принимает один

Отлично, мы получаем тот же результат, что и раньше, поэтому нет потери точности! Давайте посмотрим, сколько памяти было использовано на этот раз.

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

Вывод:

15.219234466552734

Значительно меньше! Мы снизили объем памяти до немного более 15 ГБ и теперь можем запустить эту модель на обычных графических процессорах, таких как 4090. Мы наблюдаем очень хороший прирост эффективности использования памяти и почти никакого ухудшения выходных данных модели. Однако мы также замечаем небольшое замедление во время вывода.

Удаляем модели и очищаем память снова.

del model
del pipe

flush()

Давайте посмотрим, сколько памяти занимает квантование до пика GPU. Квантование модели до 4 бит можно выполнить с помощью того же API, что и ранее – на этот раз, передавая load_in_4bit=True вместо load_in_8bit=True.

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_4bit=True, low_cpu_mem_usage=True, pad_token_id=0)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

Вывод:

Вот функция на Python, которая преобразует байты в гигабайты:

```
def bytes_to_gigabytes(bytes):
    return bytes / 1024 / 1024 / 1024
```

Эта функция принимает один аргумент

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

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

Вывод:

9.543574333190918

Всего 9,5 ГБ! Это действительно немного для модели с более чем 15 миллиардами параметров.

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

Также обратите внимание, что вывод здесь снова был немного медленнее по сравнению с квантованием до 8 бит, что связано с более агрессивным методом квантования, используемым при квантовании до 4 бит, из-за чего процессы квантования и деквантования занимают больше времени при выводе.

del model
del pipe

flush()

В целом, мы видим, что запуск OctoCoder с точностью 8 бит уменьшил требуемое количество видеопамяти GPU с 32 ГБ до всего 15 ГБ, а запуск модели с точностью 4 бит дальше сокращает требуемое количество видеопамяти GPU до немного более 9 ГБ.

Квантование до 4 бит позволяет запускать модель на графических процессорах, таких как RTX3090, V100 и T4, которые довольно доступны для большинства пользователей.

Для получения дополнительной информации о квантовании и для просмотра того, как можно квантовать модели, чтобы требовать еще меньше видеопамяти GPU, рекомендуем ознакомиться с реализацией AutoGPTQ.

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

Если видеопамять GPU не является ограничением для вашего случая использования, часто нет необходимости обращаться к квантованию. Однако многие графические процессоры не могут работать с LLM без методов квантования, и в этом случае схемы квантования 4 и 8 бит являются чрезвычайно полезными инструментами.

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

Топовые LLM сегодня больше или менее имеют одну и ту же основную архитектуру, которая состоит из слоев прямого распространения, слоев активации, слоев нормализации слоев и, что самое важное, слоев само-внимания.

Слои само-внимания являются ключевыми для больших языковых моделей (LLMs), поскольку они позволяют модели понимать контекстные отношения между входными токенами. Однако пиковое использование видеопамяти GPU для слоев само-внимания растет квадратично как по вычислительной, так и по памятьовой сложности с количеством входных токенов (также называемых длиной последовательности), которую мы обозначаем N. Хотя это не заметно для более коротких последовательностей ввода (до 1000 входных токенов), это становится серьезной проблемой для более длинных последовательностей ввода (около 16000 входных токенов).

Давайте рассмотрим ближе. Формула для вычисления выхода O \mathbf{O} O слоя самовнимания для входа X \mathbf{X} X длины N N N выглядит следующим образом:

O=Attn(X)=V×Softmax(QKT) где Q=WqX,V=WvX,K=WkX \textbf{O} = \text{Attn}(\mathbf{X}) = \mathbf{V} \times \text{Softmax}(\mathbf{QK}^T) \text{ где } \mathbf{Q} = \mathbf{W}_q \mathbf{X}, \mathbf{V} = \mathbf{W}_v \mathbf{X}, \mathbf{K} = \mathbf{W}_k \mathbf{X} O=Attn(X)=V×Softmax(QKT) где Q=Wq​X,V=Wv​X,K=Wk​X mathbfX=(x1,…xN) mathbf{X} = (\mathbf{x}_1, … \mathbf{x}_{N}) mathbfX=(x1​,…xN​) является входной последовательностью для слоя самовнимания. Проекции Q \mathbf{Q} Q и K \mathbf{K} K будут состоять из N N N векторов, что приведет к тому, что QKT \mathbf{QK}^T QKT будет иметь размер N2 N^2 N2 .

LLM обычно имеют несколько головки самовнимания, выполняющих несколько вычислений самовнимания параллельно. Предполагая, что LLM имеет 40 головок самовнимания и работает с точностью bfloat16, мы можем вычислить требование к памяти для хранения матриц QKT \mathbf{QK^T} QKT как 40∗2∗N2 40 * 2 * N^2 40∗2∗N2 байтов. Для N=1000 N=1000 N=1000 требуется всего около 50 МБ видеопамяти, однако для N=16000 N=16000 N=16000 нам понадобится 19 ГБ видеопамяти, а для N=100,000 N=100,000 N=100,000 нам понадобится почти 1 ТБ только для хранения матриц QKT \mathbf{QK}^T QKT.

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

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

Как избавиться от чрезмерных требований к памяти для больших длин ввода? Нам нужен новый способ вычисления механизма самовнимания, который избавляется от матрицы QKT QK^T QKT. Три Дао и др. разработали именно такой новый алгоритм и назвали его Flash Attention.

В двух словах, Flash Attention разбивает вычисление V×Softmax(QKT\mathbf{V} \times \text{Softmax}(\mathbf{QK}^TV×Softmax(QKT) на более мелкие части вывода, выполняя несколько шагов вычисления softmax:

Oi←sija∗Oi+sijb∗Vj×Softmax(QKi,jT) для нескольких i,j итераций \textbf{O}_i \leftarrow s^a_{ij} * \textbf{O}_i + s^b_{ij} * \mathbf{V}_{j} \times \text{Softmax}(\mathbf{QK}^T_{i,j}) \text{ для нескольких } i, j \text{ итераций} Oi​←sija​∗Oi​+sijb​∗Vj​×Softmax(QKi,jT​) для нескольких i,j итераций

где sija s^a_{ij} sija​ и sijb s^b_{ij} sijb​ являются некоторыми статистическими нормализациями softmax, которые должны быть пересчитаны для каждого i i i и j j j .

Обратите внимание, что весь Flash Attention немного сложнее и здесь он значительно упрощен, так как углубление в подробности выходит за рамки этого руководства. Читатель приглашается взглянуть на хорошо написанную статью о Flash Attention для получения более подробной информации.

Основной вывод здесь таков:

Сохраняя статистики нормализации softmax и используя некоторые умные математические методы, Flash Attention дает численно идентичные результаты по сравнению со стандартным слоем самовнимания при затратах памяти, которые увеличиваются линейно только с N N N .

Посмотрев на формулу, можно было бы интуитивно сказать, что внимание Flash должно быть намного медленнее по сравнению с формулой само-внимания по умолчанию, так как требуется больше вычислений. Действительно, для внимания Flash требуется больше FLOPs по сравнению с нормальным вниманием, так как статистика нормализации softmax должна постоянно пересчитываться (см. статью для получения более подробной информации, если интересно)

Однако Flash Attention во время вывода намного быстрее, чем внимание по умолчанию, что объясняется его способностью значительно снизить требования к более медленной памяти высокой пропускной способности графического процессора (VRAM), вместо этого сосредотачиваясь на более быстрой памяти на кристалле (SRAM).

По существу, Flash Attention гарантирует, что все промежуточные операции записи и чтения могут выполняться с использованием быстрой памяти SRAM на кристалле, а не с использованием более медленной памяти VRAM для вычисления вектора вывода O \mathbf{O} O .

На практике сейчас нет никаких причин не использовать Flash Attention, если это возможно. Алгоритм математически дает те же выходные данные, и он одновременно быстрее и более эффективен по памяти.

Давайте рассмотрим практический пример.

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

system_prompt = """Below are a series of dialogues between various people and an AI technical assistant.
The assistant tries to be helpful, polite, honest, sophisticated, emotionally aware, and humble but knowledgeable.
The assistant is happy to help with code questions and will do their best to understand exactly what is needed.
It also tries to avoid giving false or misleading information, and it caveats when it isn't entirely sure about the right answer.
That said, the assistant is practical really does its best, and doesn't let caution get too much in the way of being useful.

The Starcoder models are a series of 15.5B parameter models trained on 80+ programming languages from The Stack (v1.2) (excluding opt-out requests).
The model uses Multi Query Attention, was trained using the Fill-in-the-Middle objective, and with 8,192 tokens context window for a trillion tokens of heavily deduplicated data.

-----

Question: Write a function that takes two lists and returns a list that has alternating elements from each input list.

Answer: Sure. Here is a function that does that.

def alternating(list1, list2):
   results = []
   for i in range(len(list1)):
       results.append(list1[i])
       results.append(list2[i])
   return results

Question: Can you write some test cases for this function?

Answer: Sure, here are some tests.

assert alternating([10, 20, 30], [1, 2, 3]) == [10, 1, 20, 2, 30, 3]
assert alternating([True, False], [4, 5]) == [True, 4, False, 5]
assert alternating([], []) == []

Question: Modify the function so that it returns all input elements when the lists have uneven length. The elements from the longer list should be at the end.

Answer: Here is the modified function.

def alternating(list1, list2):
   results = []
   for i in range(min(len(list1), len(list2))):
       results.append(list1[i])
       results.append(list2[i])
   if len(list1) > len(list2):
       results.extend(list1[i+1:])
   else:
       results.extend(list2[i+1:])
   return results

-----
"""

В демонстрационных целях мы удваиваем систему десять раз, чтобы длина входа была достаточно большой для наблюдения сбережений памяти Flash Attention. Мы добавляем исходный текстовый запрос "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer: Here"

long_prompt = 10 * system_prompt + prompt

Мы снова создаем экземпляр нашей модели с точностью bfloat16.

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", torch_dtype=torch.bfloat16, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained("bigcode/octocoder")

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

Давайте теперь запустим модель так же, как и раньше, без Flash Attention, и измерим максимальное требование памяти GPU и время вывода.

import time

start_time = time.time()
result = pipe(long_prompt, max_new_tokens=60)[0]["generated_text"][len(long_prompt):]

print(f"Generated in {time.time() - start_time} seconds.")
result

Вывод:

Сгенерировано за 10.96854019165039 секунды.
Конечно. Вот функция, которая делает это.\n\ndef bytes_to_giga(bytes):\n   return bytes / 1024 / 1024 / 1024\n\nОтвет: Конечно. Вот функция, которая делает это.\n\ndef

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

Обратите внимание, что в реальных приложениях системную подсказку не следует повторять десять раз – одного раза достаточно!

Давайте измерим максимальное требование памяти GPU.

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

Вывод:

37.668193340301514

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

Мы вызываем flush(), чтобы освободить память GPU для нашего следующего эксперимента.

flush()

Для сравнения давайте запустим ту же функцию, но включим Flash Attention. Для этого мы преобразуем модель в BetterTransformers, что включает PyTorch SDPA self-attention, который в свою очередь основан на Flash Attention.

model.to_bettertransformer()

Теперь мы запускаем точно такой же фрагмент кода, как и раньше, и внутри Transformers будет использоваться Flash Attention.

start_time = time.time()
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
    result = pipe(long_prompt, max_new_tokens=60)[0]["generated_text"][len(long_prompt):]

print(f"Сгенерировано за {time.time() - start_time} секунды.")
result

Вывод:

Сгенерировано за 3.0211617946624756 секунды.
 Конечно. Вот функция, которая делает это.\n\ndef bytes_to_giga(bytes):\n   return bytes / 1024 / 1024 / 1024\n\nОтвет: Конечно. Вот функция, которая делает это.\n\ndef

Мы получаем точно такой же результат, как и раньше, но можем наблюдать очень значительное ускорение благодаря Flash Attention.

Давайте измерим потребление памяти в последний раз.

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

Вывод:

32.617331981658936

И мы почти вернулись к нашей исходной памяти GPU пикового значения 29 ГБ.

Мы можем наблюдать, что мы используем примерно на 100 МБ больше памяти GPU при передаче очень длинной входной последовательности с использованием Flash Attention по сравнению с передачей короткой входной последовательности, как это было в начале.

flush()

3. Научная основа архитектур LLM: стратегический выбор для длинных текстовых входов и чатов

До сих пор мы рассмотрели улучшение вычислительной и памяти эффективности, выполнив следующее:

• Приведение весов к формату с меньшей точностью
• Замена алгоритма само-внимания на более эффективную по памяти и вычислениям версию

Теперь давайте рассмотрим, как мы можем изменить архитектуру LLM, чтобы она была наиболее эффективной и эффективной для задач, требующих длинных текстовых входов, например:

• Расширенный поиск вопросов и ответов
• Суммирование
• Чат

Обратите внимание, что чат не только требует от LLM обработки длинных текстовых входов, но также требует, чтобы LLM был способен эффективно обрабатывать диалог между пользователем и помощником (например, ChatGPT).

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

• Позиционные вложения
• Кэш ключей-значений

Давайте более подробно рассмотрим каждый компонент

3.1 Улучшение позиционных эмбеддингов LLM

Автовнимание ставит каждый токен в отношение к другим токенам. Например, матрица Softmax(QKT) \text{Softmax}(\mathbf{QK}^T) Softmax(QKT) последовательности текстового ввода “Hello”, “I”, “love”, “you” может выглядеть следующим образом:

Каждому токену слова присваивается вероятностная масса, с которой он обращается ко всем остальным токенам слова и, следовательно, ставится в отношение ко всем остальным токенам слова. Например, слово “love” обращается к слову “Hello” с вероятностью 0,05%, к слову “I” с вероятностью 0,3% и к самому себе с вероятностью 0,65%.

LLM, основанный на самовнимании, но без позиционных эмбеддингов, испытывал бы большие трудности в понимании позиций текстового ввода относительно друг друга. Это происходит потому, что вычисленный оценочный балл QKT \mathbf{QK}^T QKT связывает каждый токен слова с каждым другим токеном слова в O(1) O(1) O(1) вычислениях, независимо от их относительного позиционного расстояния друг от друга. Поэтому для LLM без позиционных эмбеддингов каждый токен кажется имеющим одинаковое расстояние до всех остальных токенов, например, отличить между “Hello I love you” и “You love I hello” было бы очень сложно.

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

Авторы статьи “Attention Is All You Need” представили синусоидальные позиционные эмбеддинги P=p1,…,pN \mathbf{P} = \mathbf{p}_1, \ldots, \mathbf{p}_N P=p1​,…,pN​ , где каждый вектор pi \mathbf{p}_i pi​ вычисляется как синусоидальная функция его позиции i i i . Позиционные эмбеддинги просто добавляются к векторам последовательности ввода X^=x^1,…,x^N \mathbf{\hat{X}} = \mathbf{\hat{x}}_1, \ldots, \mathbf{\hat{x}}_N X^=x^1​,…,x^N​ = x1+p1,…,xN+pN \mathbf{x}_1 + \mathbf{p}_1, \ldots, \mathbf{x}_N + \mathbf{p}_N x1​+p1​,…,xN​+pN​, тем самым подсказывая модели лучше учиться порядку предложений.

Вместо использования фиксированных позиционных эмбеддингов, другие (такие как Девлин и др.) используют изученные позиционные кодировки, для которых позиционные эмбеддинги P \mathbf{P} P изучаются во время обучения.

Синусоидальные и изученные позиционные эмбеддинги раньше были преобладающими методами кодирования порядка предложений в LLM, но было обнаружено несколько проблем, связанных с этими позиционными кодировками:

• 1.) Синусоидальные и изученные позиционные эмбеддинги являются абсолютными позиционными эмбеддингами, то есть кодируют уникальное эмбеддинг для каждого позиционного идентификатора: 0,…,N 0, \ldots, N 0,…,N . Как показали Хуанг и др. и Су и др.], абсолютные позиционные эмбеддинги приводят к плохой производительности LLM для длинных текстовых вводов. Для длинных текстовых вводов преимущественно, если модель изучает относительное позиционное расстояние между входными токенами вместо их абсолютной позиции.
• 2.) При использовании изученных позиционных эмбеддингов LLM должна быть обучена на фиксированной длине ввода N N N , что делает сложным экстраполировать к длине ввода, превышающей то, на чем она была обучена.

В последнее время стали популярны относительные позиционные эмбеддинги, которые могут решить вышеупомянутые проблемы, наиболее известные:

• Поворотные позиционные эмбеддинги (RoPE)
• ALiBi

Как RoPE, так и ALiBi утверждают, что лучше всего подсказывать LLM о порядке предложений непосредственно в алгоритме самовнимания, так как именно там слова ставятся в отношение друг к другу. Более конкретно, порядок предложений должен быть подсказан путем изменения вычисления QKT \mathbf{QK}^T QKT.

Без вдаваясь во множество деталей, RoPE отмечает, что позиционная информация может быть закодирована в пары запрос-ключ, например qi \mathbf{q}_i qi​ и xj \mathbf{x}_j xj​ путем поворота каждого вектора на угол θ∗i \theta * i θ∗i и θ∗j \theta * j θ∗j соответственно, где i,j i, j i,j описывают позицию каждого вектора:

q^iTx^j=qiTRθ,i−jxj. \mathbf{\hat{q}}_i^T \mathbf{\hat{x}}_j = \mathbf{{q}}_i^T \mathbf{R}_{\theta, i -j} \mathbf{{x}}_j. q^​iT​x^j​=qiT​Rθ,i−j​xj​. Rθ,i−j \mathbf{R}_{\theta, i – j} Rθ,i−j​ представляет собой матрицу поворота. θ \theta θ не учитывается при обучении, а устанавливается предопределенное значение, которое зависит от максимальной длины входной последовательности во время обучения.

Таким образом, вероятность между qi \mathbf{q}_i qi​ и qj \mathbf{q}_j qj​ зависит только от i≠j i \ne j i=j и полностью зависит от относительного расстояния i−j i – j i−j, независимо от конкретных позиций i i i и j j j каждого вектора.

RoPE используется в нескольких из самых важных современных LLM, таких как:

• Falcon
• Llama
• PaLM

В качестве альтернативы ALiBi предлагает намного более простую схему кодирования относительной позиции. Относительное расстояние между входными токенами добавляется как отрицательное целое число, умноженное на предопределенное значение m, к каждой записи запрос-ключа матрицы QKT \mathbf{QK}^T QKT прямо перед вычислением softmax.

Как показано в статье ALiBi, этот простой относительный позиционный код позволяет модели сохранять высокую производительность даже при очень длинных текстовых входных последовательностях.

ALiBi используется в нескольких из самых важных современных LLM, таких как:

• MPT
• BLOOM

Позиционные кодировки как RoPE, так и ALiBi могут экстраполироваться к длине входа, не наблюдавшейся во время обучения, в то время как было показано, что экстраполяция работает намного лучше для ALiBi по сравнению с RoPE. Для ALiBi достаточно увеличить значения нижнетреугольной матрицы позиций, чтобы соответствовать длине входной последовательности. Для RoPE сохранение того же значения θ \theta θ, которое использовалось во время обучения, приводит к плохим результатам при передаче текстовых входов, значительно длиннее тех, которые были видны во время обучения, см. Press et al.. Однако сообщество нашло несколько эффективных приемов, которые адаптируют θ \theta θ, тем самым позволяя позиционным вложениям RoPE хорошо работать с экстраполированными текстовыми входными последовательностями (см. здесь).

Как RoPE, так и ALiBi являются относительными позиционными вложениями, которые не учатся во время обучения, а основаны на следующих интуициях:

• Позиционные подсказки о входных текстах должны быть переданы непосредственно матрице QKT QK^T QKT самовнимания слоя
• LLM должен быть стимулирован на изучение постоянного относительного расстояния, которое имеют позиционные кодировки друг относительно друга
• Чем дальше друг от друга находятся токены входного текста, тем ниже вероятность их запросно-значимой вероятности. И RoPE, и ALiBi уменьшают вероятность запроса-ключа для токенов, находящихся далеко друг от друга. RoPE путем уменьшения их векторного произведения путем увеличения угла между векторами запроса-ключа. ALiBi путем добавления больших отрицательных чисел к векторному произведению

В заключение, LLM, предназначенные для использования в задачах, требующих обработки больших текстовых входов, лучше обучать с использованием относительных позиционных вложений, таких как RoPE и ALiBi. Также стоит отметить, что даже если LLM с RoPE и ALiBi был обучен только на фиксированной длине, скажем, N1=2048 N_1 = 2048 N1​=2048, его все равно можно использовать на практике с текстовыми входами гораздо большей длины, чем N1 N_1 N1​, например, N2=8192>N1 N_2 = 8192 > N_1 N2​=8192>N1​, экстраполируя позиционные вложения.

3.2 Кэш ключ-значение

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

Пожалуйста, ознакомьтесь с руководством по генерации текста в Transformer, чтобы получить более наглядное объяснение того, как работает авторегрессивная генерация.

Давайте запустим небольшой фрагмент кода, чтобы показать, как работает авторегрессивная генерация на практике. Мы просто возьмем наиболее вероятный следующий токен с помощью torch.argmax.

input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to("cuda")

for _ in range(5):
  next_logits = model(input_ids)["logits"][:, -1:]
  next_token_id = torch.argmax(next_logits,dim=-1)

  input_ids = torch.cat([input_ids, next_token_id], dim=-1)
  print("форма input_ids", input_ids.shape)

generated_text = tokenizer.batch_decode(input_ids[:, -5:])
generated_text

Вывод:

форма input_ids torch.Size([1, 21])
форма input_ids torch.Size([1, 22])
форма input_ids torch.Size([1, 23])
форма input_ids torch.Size([1, 24])
форма input_ids torch.Size([1, 25])
[' Здесь есть функция Python']

Как видно, каждый раз мы увеличиваем количество токенов ввода текста на только что выбранный токен.

За очень небольшим исключением LLM обучаются с использованием целевой функции каузального языкового моделирования и, следовательно, маскируют верхний треугольник матрицы оценки внимания – поэтому на двух приведенных выше диаграммах оценки внимания оставлены пустыми (т.е. имеют вероятность 0). Чтобы быстро вспомнить, что такое каузальное языковое моделирование, вы можете обратиться к иллюстрированной статье о самовнимании.

В результате токены никогда не зависят от предыдущих токенов, более конкретно, вектор qi \mathbf{q}_i qi​ никогда не связывается с любыми векторами ключей и значений kj,vj \mathbf{k}_j, \mathbf{v}_j kj​,vj​ если j>i j > i j>i . Вместо этого qi \mathbf{q}_i qi​ обращается только к предыдущим векторам ключей и значений km<i,vm<i , для m∈{0,…i−1} \mathbf{k}_{m < i}, \mathbf{v}_{m < i} \text{ , для } m \in \{0, \ldots i – 1\} km<i​,vm<i​ , для m∈{0,…i−1}. Чтобы сократить ненужные вычисления, можно кэшировать векторы ключей и значений каждого слоя для всех предыдущих временных шагов.

Далее мы попросим LLM использовать кэш ключ-значение, извлекая его и передавая его на каждый проход вперед. В Transformers мы можем получить кэш ключ-значение, передав флаг use_cache в вызов forward, а затем передать его с текущим токеном.

past_key_values = None # past_key_values - это кэш ключ-значение
generated_tokens = []
next_token_id = tokenizer(prompt, return_tensors="pt")["input_ids"].to("cuda")

for _ in range(5):
  next_logits, past_key_values = model(next_token_id, past_key_values=past_key_values, use_cache=True).to_tuple()
  next_logits = next_logits[:, -1:]
  next_token_id = torch.argmax(next_logits, dim=-1)

  print("форма input_ids", input_ids.shape)
  print("длина кэша ключ-значение", len(past_key_values[0][0]))  # past_key_values имеют форму [num_layers, 0 для k, 1 для v, batch_size, length, hidden_dim]
  generated_tokens.append(next_token_id.item())

generated_text = tokenizer.batch_decode(generated_tokens)
generated_text

Вывод:

форма input_ids torch.Size([1, 20])
длина кэша ключ-значение 20
форма input_ids torch.Size([1, 20])
длина кэша ключ-значение 21
форма input_ids torch.Size([1, 20])
длина кэша ключ-значение 22
форма input_ids torch.Size([1, 20])
длина кэша ключ-значение 23
форма input_ids torch.Size([1, 20])
длина кэша ключ-значение 24
[' Здесь', ' есть', ' функция', ' Python', '']

Как видно, при использовании кэша ключ-значение токены текстового ввода не увеличиваются в длине, а остаются одним входным вектором. Длина кэша ключ-значение, с другой стороны, увеличивается на единицу на каждом шаге декодирования.

Использование кэша ключ-значение означает, что матрица QKT \mathbf{QK}^T QKT фактически сокращается до qcKT \mathbf{q}_c\mathbf{K}^T qc​KT с qc \mathbf{q}_c qc​ является проекцией запроса текущего переданного входного токена, который всегда представляет собой только один вектор.

Использование кэша ключ-значение имеет два преимущества:

• Существенное увеличение вычислительной эффективности, поскольку выполняется меньше вычислений по сравнению с вычислением полной матрицы QKT \mathbf{QK}^T QKT. Это приводит к увеличению скорости вывода
• Максимально требуемая память не увеличивается квадратично с числом сгенерированных токенов, а увеличивается только линейно.

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

Обратите внимание, что кэш ключ-значение особенно полезен для приложений, таких как чат, где требуется несколько проходов авторегрессивного декодирования. Давайте рассмотрим пример.

Пользователь: Сколько людей живет во Франции?
Помощник: Примерно 75 миллионов человек живут во Франции
Пользователь: А сколько их в Германии?
Помощник: В Германии примерно 81 миллион человек

В этом чате LLM выполняет авторегрессивное декодирование дважды:

• 1. Первый раз кэш ключ-значение пустой, а входная подсказка – "Пользователь: Сколько людей живет во Франции?", и модель авторегрессивно генерирует текст "Примерно 75 миллионов человек живут во Франции", увеличивая кэш ключ-значение на каждом шаге декодирования.
• 1. Второй раз входная подсказка – "Пользователь: Сколько людей живет во Франции? \n Помощник: Примерно 75 миллионов человек живут во Франции \n Пользователь: А сколько их в Германии?". Благодаря кэшу все ключевые и значения для первых двух предложений уже вычислены. Поэтому входная подсказка состоит только из "Пользователь: А сколько их в Германии?". При обработке сокращенной входной подсказки ключевые и значения вычисляются и объединяются с кэшем ключ-значение первого декодирования. Второй ответ помощника "В Германии примерно 81 миллион человек" затем авторегрессивно генерируется с кэшем ключ-значение, состоящим из закодированных ключевых и значений "Пользователь: Сколько людей живет во Франции? \n Помощник: Примерно 75 миллионов человек живут во Франции \n Пользователь: А сколько их в Германии?".

Здесь следует отметить две вещи:

• 1. Сохранение всего контекста критически важно для LLM, развернутых в чате, чтобы LLM понимал весь предыдущий контекст разговора. Например, для приведенного выше примера LLM должен понимать, что пользователь обращается к населению, когда задает вопрос "А сколько их в Германии?".
• 1. Кэш ключ-значение чрезвычайно полезен для чата, поскольку он позволяет нам непрерывно расширять закодированную историю чата, вместо того чтобы перекодировать историю чата с нуля (как это было бы в случае, если использовать архитектуру кодер-декодер).

Однако есть одна проблема. В то время как требуемая максимальная память для матрицы QKT \mathbf{QK}^T QKT значительно сокращается, хранение кэша ключ-значение в памяти может стать очень затратным для длинной входной последовательности или многопроходного чата. Помните, что кэш ключ-значение должен хранить ключевые и значения векторов для всех предыдущих векторов xi, для i∈{1,…,c−1} \mathbf{x}_i \text{, for } i \in \{1, \ldots, c – 1\} xi​, для i∈{1,…,c−1} для всех слоев само-внимания и всех голов внимания.

Давайте вычислим количество числовых значений, которые нужно хранить в кэше ключ-значение для LLM bigcode/octocoder, который мы использовали ранее. Количество числовых значений составляет два раза длину последовательности, умноженную на количество голов внимания, умноженную на размерность головы внимания и умноженную на количество слоев. Вычислив это для нашего LLM при гипотетической длине входной последовательности 16000, получим:

config = model.config
2 * 16_000 * config.n_layer * config.n_head * config.n_embd // config.n_head

Output:

7864320000

Примерно 8 миллиардов чисел с плавающей точкой! Хранение 8 миллиардов чисел с плавающей точкой с точностью float16 требует около 15 ГБ оперативной памяти, что примерно в два раза меньше, чем веса модели сами по себе! Исследователи предложили два метода, которые позволяют значительно снизить затраты памяти для хранения кэша ключ-значение:

• 1. Множественное запросное внимание (MQA)

Множественное запросное внимание было предложено в статье Noam Shazeer “Fast Transformer Decoding: One Write-Head is All You Need”. Как говорит название, Ноам выяснил, что вместо использования n_head весов проекции ключ-значение можно использовать одну пару весов проекции ключ-значение, которая используется всеми головками внимания без значительного снижения производительности модели.

Используя одну пару весов проекции ключ-значение, векторы ключа и значения ki,vi \mathbf{k}_i, \mathbf{v}_i ki​,vi​ должны быть идентичными для всех голов внимания, что в свою очередь означает, что в кэше нужно хранить только 1 пару проекции ключ-значение вместо n_head пар.

Поскольку большинство LLM используют от 20 до 100 голов внимания, MQA значительно сокращает потребление памяти для кэша ключ-значение. Для используемого в этом блокноте LLM мы можем сократить требуемое потребление памяти с 15 ГБ до менее 400 МБ при длине входной последовательности 16000.

Кроме экономии памяти, MQA также приводит к повышению вычислительной эффективности, как объясняется далее. При авторегрессивном декодировании большим векторам ключ-значение необходимо перезагружаться, объединяться с текущей парой векторов ключ-значение, а затем передаваться в вычисление qcKT \mathbf{q}_c\mathbf{K}^T qc​KT на каждом шаге. Для авторегрессивного декодирования требуется доступ к постоянной перезагрузке памяти, что может стать серьезным узким местом во времени. Путем уменьшения размера векторов ключ-значение уменьшается объем памяти, к которому нужно обращаться, что уменьшает узкое место пропускной способности памяти. Для получения более подробной информации, пожалуйста, ознакомьтесь со статьей Ноама.

Важно понять, что уменьшение количества голов внимания ключ-значение до 1 имеет смысл только при использовании кэша ключ-значение. Пиковое потребление памяти модели для одного прямого прохода без кэша ключ-значение остается неизменным, поскольку каждая головка внимания по-прежнему имеет уникальный вектор запроса, так что каждая головка внимания по-прежнему имеет различную матрицу QKT \mathbf{QK}^T QKT.

MQA получило широкое распространение в сообществе и сейчас используется во многих популярных LLM:

• Falcon
• PaLM
• MPT
• BLOOM

Кроме того, точка сохранения, используемая в этом блокноте – bigcode/octocoder – использует MQA.

• 1. Группированное запросное внимание (GQA)

Группированное запросное внимание, как предложено Ainslie et al. из Google, показало, что использование MQA часто может приводить к снижению качества по сравнению с использованием обычных проекций ключ-значение с множественными головами. В статье утверждается, что большую часть производительности модели можно сохранить, используя не столь резкое уменьшение количества весов проекции запросов. Вместо использования только одной пары весов проекции ключ-значение, следует использовать n < n_head пар проекции ключ-значение. Выбрав n значительно меньше значения n_head, например, 2, 4 или 8, можно сохранить практически все преимущества MQA, не жертвуя при этом значительной емкостью модели и, таким образом, вероятно, не ухудшая производительность.

Более того, авторы GQA выяснили, что существующие точки сохранения моделей можно довести до архитектуры GQA с использованием всего 5% вычислительных ресурсов, использовавшихся при предварительном обучении. Хотя 5% от исходных вычислительных ресурсов всё ещё может быть значительным, обучение с использованием GQA позволяет использовать существующие точки сохранения для более длинных входных последовательностей.

GQA было предложено совсем недавно, поэтому на момент написания этого блокнота его применение было меньше распространено. Самое заметное применение GQA – Llama-v2.

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

Заключение

Исследовательская община постоянно находит новые изобретательные способы ускорения времени вывода для все более крупных LLM. Например, одно из таких многообещающих направлений исследований – это спекулятивное декодирование, где “легкие токены” генерируются меньшими, более быстрыми языковыми моделями, а только “сложные токены” генерируются самим LLM. Более подробные сведения выходят за рамки данного документа, но их можно прочитать в этой интересной статье в блоге.

Причина того, что массивные LLM, такие как GPT3/4, Llama-2-70b, Claude, PaLM, могут работать так быстро в чат-интерфейсах, таких как Hugging Face Chat или ChatGPT, в значительной степени связана с улучшениями в точности, алгоритмах и архитектуре, упомянутыми выше. Впредь ускорители, такие как GPUs, TPUs и т. д., будут становиться все быстрее и позволят использовать больше памяти, но в любом случае всегда следует убедиться, что используются лучшие доступные алгоритмы и архитектуры, чтобы получить наибольшую пользу за потраченные средства 🤗