В этом сообщении блога рассматриваются методы повышения скорости вывода моделей серии Llama 2 с помощью встроенных улучшений PyTorch, включая прямые высокоскоростные ядра, возможности преобразования Torch-компиляции и тензорное распараллеливание для распределенных вычислений. Мы достигли задержки 29 миллисекунд на токен для отдельных запросов на модели 70B LLaMa , протестированной на восьми графических процессорах A100. Мы стремимся поделиться этими знаниями с нашими коллегами и предоставили наш код для использования сообществом.
ВВЕДЕНИЕ
В настоящее время в области генеративного искусственного интеллекта наблюдается всплеск доступности гигантских языковых моделей с числом параметров, исчисляемым десятками миллиардов. Хотя распространение таких моделей неоспоримо, сохраняется серьезная проблема с их экономически эффективным внедрением. Сообщество ИИ экспериментировало с многочисленными стратегиями, добиваясь разного успеха и компромиссов. Аппаратно-ориентированные оптимизации, такие как Faster Transformer от NVIDIA, ограничены конкретным оборудованием, тогда как более общие методы, такие как ONNX, обеспечивают компромисс между эффективностью и гибкостью.
Запуск компиляции PyTorch в прошлом году ознаменовал начало сотрудничества IBM и PyTorch по совершенствованию компиляции моделей для повышения производительности вывода с целью сократить задержку для каждого токена, генерируемого этими большими моделями.
ВЫБОР МОДЕЛИ
Наши тесты сосредоточены на широко используемом семействе моделей Llama 2 . Конкретные модели, представляющие интерес для этого обсуждения, а также их соответствующие гиперпараметры подробно описаны в следующей таблице:
| Размер модели | Скрытое измерение | Количество голов | Количество слоев | Тип внимания |
| 7Б | 4096 | 32 | 32 | МХА |
| 13Б | 5120 | 40 | 40 | МХА |
| 70Б | 8192 | 64 | 80 | GQA |
Обсуждаемые нами модели предназначены только для декодеров, что означает, что они генерируют токены один за другим, обычно ускоряясь за счет кэширования ключ-значение (KV). Мы используем этот метод для наших тестов задержки и пропускной способности.
МЕТОДОЛОГИЯ ВЫВОДА
Целью нашего вывода является быстрое достижение оптимальных задержек, идя в ногу с быстрым развитием новых архитектур моделей в области искусственного интеллекта. Желательно использовать собственное решение PyTorch, поскольку оно обеспечивает самую широкую совместимость моделей. Мы выделяем четыре различных метода, которые ускоряют вывод:
- (а) слияние ядра с компиляцией,
- (б) использование более быстрых ядер,
- (в) тензорный параллелизм для крупномасштабных моделей и
- (г) квантование модели.
Для наших целей мы используем первые три — использование компиляции PyTorch в сочетании с ускоренными ядрами SDPA и специальной тензорной параллельной структурой для реализации задержек вывода 29 миллисекунд на токен в модели 70B LLaMa, измеренных на восьми графических процессорах NVIDIA A100, обслуживающих Один пользователь.
Полный газ с компиляцией!
PyTorch Compile повышает производительность за счет снижения нагрузки на ЦП за счет отслеживания и захвата графиков выполнения, в идеале объединяя процесс в одно выполнение графа между процессором и графическим процессором. Однако компиляция может привести к «разрывам графа» при работе с архитектурой модели или операциями, которые она не поддерживает. Например, сложные операции, такие как einops, в настоящее время обходятся без возможностей компиляции. Аналогично, тензорный параллелизм может привести к разрывам графа на каждом уровне, если структура тензорного параллельного параллелизма не использует отслеживаемые коллективы связи. Без устранения этих разрывов графа скомпилированные артефакты могут работать хуже или даже медленнее, чем обычное активное выполнение. Чтобы использовать весь потенциал компиляции, эти перерывы необходимо устранить.
Вот как мы решили эти проблемы, чтобы модель 70B LLaMa 2 полностью использовала компиляцию.
Первоначально, когда мы попытались скомпилировать стандартную модель Llama 2 с помощью torch.compile, она не удалась из-за неподдерживаемых сложных операций. Включив TORCH_COMPILE_DEBUG = 1, мы обнаружили, что позиционные кодировки RoPE используют комплексные числовые функции, что приводит к разрывам графика и значительному замедлению. Чтобы смягчить это, мы реструктурировали функцию RoPE, чтобы избежать torch.einsum (который изначально использовал torch.polar, также несовместимый с компиляцией), выбрав вместо этого функции torch.cos и torch.sin.
self.cached_freqs[dev_idx][alpha] = torch.stack(
[
torch.cos(freqs),
-torch.sin(freqs),
torch.sin(freqs),
torch.cos(freqs),
],
dim=2,
).view(*freqs.shape, 2, 2)Наша реализация вычисления частот
t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
t = t / self.scaling_factor
freqs = torch.einsum("i,j->ij", t, self.inv_freq)
# Different from paper, but it uses a different permutation in order to obtain the same calculation
emb = torch.cat((freqs, freqs), dim=-1)
self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)
self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)Hugging Face реализация вычисления частот.
Решив проблему RoPE, мы успешно скомпилировали модели 7B и 13B, не обнаружив никаких разрывов графика на одном графическом процессоре A100.
Для эффективного расчета механизмов внимания мы использовали SDPA — собственную реализацию PyTorch, поддерживающую трассировку (необходимую для компиляции). Чтобы обойти потенциальные разрывы графа, которые возникают при принудительном выборе одного алгоритма через контекст Python (стандартная рекомендация), мы прибегли к использованию функций torch.backends.cuda.enable_*_sdp.
attn = torch.nn.functional.scaled_dot_product_attention(
queries,
keys_e,
values_e,
attn_mask=attn_mask,
dropout_p=self.p_dropout if self.training else 0.0,
is_causal=is_causal_mask,
)Вычисление внимания с использованием SDPA
Для вычисления внимания SDPA эффективно использовался на моделях меньшего размера, но когда мы масштабировались до модели 70B, мы столкнулись с ограничениями, связанными с ограничениями памяти графического процессора. Даже при использовании чисел с плавающей запятой половинной точности модель 70B была слишком велика для одного графического процессора, что требовало использования тензорного параллелизма для вывода. При попытке скомпилировать модель 70B с помощью torch.compile мы столкнулись со 162 разрывами графа из-за двух операций полного сокращения на слой и одной операции полного сбора для прямого и обратного встраивания. Это не привело к заметному уменьшению задержки вывода.
На момент написания статьи распределенный тензорный параллелизм PyTorch не интегрировался с torch.compile, что побудило нас разработать собственный код тензорного параллелизма с нуля. Этот новый код опирался исключительно на отслеживаемые коллективные операции, чтобы быть совместимым с torch.compile. В этой версии компилятор PyTorch больше не вносил разрывы графа, и мы наблюдали существенное улучшение скорости вывода. Говоря конкретнее, мы достигли задержки вывода 29 миллисекунд на токен с помощью модели 70B Llama с использованием восьми графических процессоров A100, что означает улучшение в 2,4 раза по сравнению с базовой производительностью неоптимизированного вывода.
Аспекты обслуживания
Последний аспект, который следует учитывать, заключается в том, что простой компиляции модели с помощью torch.compile недостаточно для ее использования в реальной производственной среде. Для достижения упомянутого уровня производительности наряду с высокой пропускной способностью необходимо реализовать динамическую пакетную обработку и поддержку вложенных тензоров. Кроме того, фаза разминки имеет решающее значение для предварительной компиляции модели для последовательностей различной длины, которые разбиты на сегменты. Эти усилия продолжаются, чтобы гарантировать, что такие уровни производительности осуществимы в производственном контексте.
ЭКСПЕРИМЕНТЫ И ИЗМЕРЕНИЯ
В наших экспериментальных установках и измерениях мы использовали узлы, оснащенные 8 графическими процессорами NVIDIA A100 по 80 ГБ каждый, выполняя тесты в двух отдельных облачных средах (IBM Cloud и AWS, оба используют OpenShift). Мы оценили несколько различных методов: режим ожидания, ядро SDPA Flash, компиляцию и комбинацию компиляции и SDPA.
Для модели 70B мы использовали тензорный параллельный режим в сочетании с компиляцией и SDPA. Входные данные для этого эксперимента были зафиксированы на уровне 512 токенов, при этом было сгенерировано еще 50 токенов. Модели 7B и 13B тестировались с использованием одного графического процессора A100 для измерения задержек, а модель 70B тестировалась с использованием 8 графических процессоров A100. Специально для модели 70B мы также использовали опцию уменьшения накладных расходов, доступную в компиляции PyTorch, которая включает CudaGraphs, чтобы минимизировать накладные расходы ЦП на запуск ядра графического процессора.
Интересно, что использование CudaGraphs для моделей 7B и 13B не дало никаких преимуществ, поэтому эти результаты не включены в наш отчет. Результаты, показанные на рисунке 1, показывают, что комбинация компиляции и SDPA обеспечивает удивительно низкие задержки: модель 70B Llama 2 записывает задержку 29 мс на токен.
После дальнейшего исследования мы оцениваем влияние различной длины последовательности на производительность, в частности, путем увеличения длины последовательности с 1024 до 4096 токенов. Результаты указывают на сублинейный рост средней задержки на токен. Этот вывод важен, поскольку предполагает, что модель эффективно масштабируется с увеличением размера контекста.
Таким образом, когда модель обрабатывает большие документы, увеличение времени ответа не пропорционально увеличению размера документа, гарантируя, что время ответа останется в разумных пределах даже при росте объема обрабатываемых данных.
Заключительная часть анализа
В заключительной части нашего анализа мы сосредоточимся на взаимосвязи между размером пакетов и задержкой ответа. Данные показывают, что по мере увеличения размеров пакетов результирующие задержки также растут, но сублинейным образом. Это означает, что задержка растет не так быстро, как размер пакета, что указывает на то, что система эффективно справляется с возросшей нагрузкой.
Однако отмечается, что при тестировании модели 13B с размером пакета 8 возникает ошибка нехватки памяти (OOM), что указывает на то, что достигнут предел того, что может обработать один графический процессор для этого размера модели. Напротив, для модели 70B, которая работает на 8 графических процессорах с тензорным параллелизмом, мы не сталкиваемся с проблемами OOM. Это подчеркивает преимущество распределенных вычислений и тензорного параллелизма в управлении более крупными моделями и размерами пакетов без ограничений памяти.
ПОСЛЕДНИЕ МЫСЛИ
В заключительных замечаниях была установлена эффективность пути компиляции PyTorch в достижении удивительно низких задержек вывода для модели 70B. Предстоящие шаги включают в себя включение динамической пакетной обработки и вложенных тензоров для дальнейшего использования вышеупомянутых оптимизаций.
Благодарность выражается Эдварду Янгу, Элиасу Эллисону, Дриссу Гессусу, Уиллу Фэну, Уиллу Констеблу, Хорасу Хе, Лессу Райту и Эндрю Гу из команды PyTorch. Их тщательные проверки кода и вклад сыграли важную роль в достижении заявленных низких задержек с использованием собственной методологии PyTorch. Дополнительные похвалы адресованы расширенной команде PyTorch, которая посвятила себя постоянному совершенствованию PyTorch. Особая благодарность выражается команде SDPA за их работу по отслеживанию и включению компиляции на быстрых ядрах, а также команде компилятора за руководство в навигации и решении проблем, включая выявление ошибок в графиках CUDA NVIDIA и сообщение о них команде драйверов.
Исследование IBM.
