Ускорение инференса LLM через тензорный параллелизм (подробный обзор)

Использование Ollama или нативного HF инференса может порождать стереотипы о медленной скорости генерации ответа LLM при работе на нескольких GPU, об отсутствии какого-то заметного ускорения.

Достаточно долго мне не доводилось натыкаться на факты, когда с помощью второго GPU можно ускорить генерацию одного запроса. Множество опубликованных тестов производительности делают акцент на демонстрацию увеличения общей пропускной способности при параллельной генерации. В то время как тензорный параллелизм, реализованный в профессиональных движках типа vLLM или Tensor-RT, позволяет ускорить генерацию и одного потока за счет нескольких GPU. Об этом и многом другом поговорим в данном обзоре.

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

В исследовании было показано, LLaMA-2-7B превосходит по производительности как Mistral-7B, так и LLaMA-3-8B, в то время как Mistral-7B опережает LLaMA-3-8B при различных размерах батча и количестве GPU. Это контринтуитивно по сравнению с TRT-LLM и vLLM, где модели с GQA работают лучше, чем модели с MHA. Это показывает, что llama.cpp не в состоянии в полной мере использовать преимущества механизма Grouped-Query Attention.

Несмотря на то, что llama.cpp активно развивается и внедряются различные методы оптимизации, полноценного ускорения от тензорного параллелизма нет даже через опцию --split-mode row (-sm row), которая ускоряет по разным оценкам до 20%.

Схожая проблема наблюдается и со стандартным инференсом HF через device_map. Но об этом чуть позже.

В контексте LLM тензор — это универсальный контейнер для чисел, похожий на многомерную таблицу, в котором хранится абсолютно вся информация. Когда вы пишете запрос, например, «Сколько цветов», модель сначала превращает его в простой числовой список (1D-тензор), допустим [5, 2, 7]. Затем этот список преобразуется в более сложный, «глубокий» 3D-тензор, где каждое число «обрастает» сотнями других чисел, описывающих его семантическое значение.

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

Если тензоры — это гигантские таблицы с числами, на которых работает LLM, то тензорные ядра в GPU — это специализированные «мини-фабрики» для их быстрой обработки. Они созданы для одной-единственной задачи: мгновенно перемножать целые блоки этих таблиц, а не отдельные числа по очереди. Представьте, что вместо обычного калькулятора у вас есть кнопка «умножить две таблицы целиком», — это и есть принцип работы тензорного ядра.

В 2017 году тензорные ядра впервые были представлены NVIDIA в архитектуре Volta. Первый ускоритель Tesla V100 получил впечатляющее число - 640 тензорных ядер. Почти все последующие GPU, включая геймерский класс, стали выпускаться с этим блоком вычислений. Выпущено уже 5 поколений к 2025 году. Наличие наибольшее числа тензорных ядер присуще топовым потребительским или профессиональным моделям, что также сказывается на увеличение стоимости.

Тензорные ядра часто вносят путаницу в производительность GPU. Во многих публикациях путают производительность тензорных вычислений с CUDA-ядрами. В технической документации разделяют на Peak FP16 (non Tensor) - обычно десятки TFLOPS для CUDA-ядер, а Peak FP16 Tensor - исчисляются сотнями TFLOPS для тензорных ядер.

Разные поколения тензорных ядер могут существенно отличаться по производительности. Например, у Tesla V100 640 тензорных ядер имеют пик 125 TLFOPS, а RTX 3090 TI имеет всего 336 тензорных ядер 3-го поколения - 160 TLFOPS, а у RTX 4090 всего 512 ядер 4-го поколения - 330 FP16 TFLOPS (и 660 FP6 TFLOPS). Однако, в реальности не всякий сценарий инференса дает прирост, только лишь на тензорным ядрам.

Применение модельного параллелизма (model parallelism) подразумевает наличие более 1 GPU в сервере или хотя бы по 1 GPU на двух серверах. Разные механизмы параллелизма (или комбинация механизмов) решают следующие задачи инференса:

Конвейерный параллелизм (Pipeline Parallelism) представляет собой метод распределения LLM целыми слоями на несколько GPU. Данные обрабатываются поэтапно, переходя от одного GPU к другому, что позволяет распределить память, необходимую для хранения весов модели, между несколькими устройствами. Можно использовать любое число устройств, как четное так и нечетное 1, 2, 3 4 и т.д.

Основным ограничением данного подхода является последовательный характер вычислений, приводящий к простоям GPU («пузырькам конвейера»), особенно при обработке одиночных запросов.

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

Тензорный параллелизм (Tensor Parallelism) представляет собой метод внутрислойного параллелизма, при котором каждый слой LLM разделяется горизонтально на фрагменты и распределяется между несколькими GPU.

Все GPU работают над одним и тем же токеном, затем частичные результаты вычислений агрегируются посредством операции All-Reduce. Данный подход способствует сокращению времени обработки слоя за счет пропорционального уменьшения вычислительной нагрузки на каждый GPU и обеспечения их одновременной работы.

Однако тензорный параллелизм сопряжен со значительными коммуникационными издержками, так как требует частого обмена данными между устройствами для синхронизации результатов после обработки каждого слоя, что выдвигает высокие требования к пропускной способности межпроцессорных соединений (NVLink, NVSwitch, InfiniBand). Эффективность тензорного параллелизма напрямую зависит от скорости этих соединений: при высокоскоростных соединениях он обеспечивает существенный прирост производительности, тогда как при медленных коммуникационные задержки могут нивелировать выгоды или замедлить инференс.

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

Data Parallel (DP) - еще один метод параллизма, но на уровне данных. Каждый GPU содержит модель целиком, батчи раскидываются параллельно.

DP, в отличие от tensor/pipeline parallelism, не помогает впихнуть слишком крупную модель и почти не влияет на латентность одного запроса, зато даёт самое простое и дешёвое горизонтальное масштабирование сервиса. Если модель уже помещается и цель — throughput кучи запросов, проще запустить несколько полных копий (data-parallel) вместо TP.

В vLLM можно запустить --tensor-parallel-size 1 --num-gpus 4 --worker-use-ray - Ray поднимет четыре независимых воркера-реплики (data-parallel) и роутер раздаст им запросы

Все режимы device_map в Hugging Face Accelerate/Transformers — "auto", "balanced", "balanced_low_0" и "sequential" — выполняют только последовательное размещение слоёв на разных устройствах, чтобы модель поместилась в память. Они не реализуют ни тензорный параллелизм (деление весов слоя между GPU с All-Reduce внутри слоя), ни конвейерный параллелизм. Поэтому от device_map не стоит ждать ускорения: в каждый момент времени работает лишь одна карта, а остальные ждут, пока управление дойдёт до их слоёв.

Как запустить модель Llama 3.1 405B без квантования, которой требуется от 800 Гб видеопамяти в bf16? Это можно сделать на 16 GPU H100 или на 16 GPU A100 80GB, используя конвейерный или тензорный параллелизм vLLM, а также через совместное использование конвейерного и тензорного параллелизма.

При наличии 16 GPU на 2 узлах, вы можете использовать двойной конвейерный параллелизм (PP) и восьмикратный тензорный параллелизм (TP) для оптимизации использования GPU.

Если у вас есть быстрая сеть между серверами типа Infiniband, вы можете использовать 16-сторонний тензорный параллелизм:

Конвейерный параллелизм необходим, когда узлы не соединены через Infiniband. По сравнению с 16-кратным тензорным параллелизмом без Infiniband, сочетание 2-кратного конвейерного параллелизма с 8-кратным тензорным параллелизмом приводит к увеличению пропускной способности в 6,6 раза. С другой стороны, при наличии Infiniband производительность обеих конфигураций схожа.

Теоретически, если коммуникация между GPU достаточно быстрая, время на проход одного токена сокращается почти вдвое на 2 GPU (или примерно в N раз на N GPU). В реальности ускорение меньше идеального из-за времени на синхронизацию результатов между устройствами (all-reduce) и возможных дисбалансов нагрузки. В отличие от конвейерного, тензорный параллелизм требует всегда четного числа GPU - 2, 4, 8.

Ускорение от тензорного параллелизма возможно как в случае, когда модель не помещается на один GPU, так и в случае, когда вмещается целиком. По умолчанию у vLLM стоит tensor_parallel_size = 1 и инференс идет на 1 GPU. Если модель помещается на 1 GPU, то указав tensor_parallel_size = 2 инференс 1 потока будет идти на 2 GPU с почти линейным ускорением.

Или из питона

vLLM распилит весовые матрицы по алгоритму Megatron-style TP и создаст два воркера (в системе два python процесса будут), работающих параллельно. См. документацию.

Оценка производительности инференса LLM требует понимания трех ключевых метрик: пропускной способности (Throughput), времени до первого токена (Time-to-First-Token, TTFT) и времени на генерацию выходного токена (Time-Per-Output-Token, TPOT). Каждая метрика и связанные с ней параметры показаны на схеме.

Пропускная способность — это количество токенов, которое система может сгенерировать в единицу времени при обработке максимального числа параллельных запросов (batch - max_num_seqs в VLLM). Она рассчитывается как общее количество сгенерированных токенов, деленное на общее время вывода (инференса). Высокая пропускная способность указывает на то, что система может эффективно обрабатывать большой объем запросов, что крайне важно для приложений, работающих в реальном времени, и для одновременного обслуживания множества пользователей.

Стоит заметить, что скорость генерации 1 запроса может быть весьма ограниченной, например, 50 токенов/секунду, в то время как общая пропускная способность может измеряться сотнями или тысячами токенов/секунду.

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

TPOT, также известный как межтокенная задержка (inter-token-latency, ITL), — это среднее время, затрачиваемое на генерацию каждого последующего токена после первого. Эта метрика дает представление о скорости генерации токенов моделью во время вывода. Чем ниже TPOT, тем быстрее и плавнее происходит генерация токенов.

Отслеживая и оптимизируя пропускную способность, TTFT и TPOT, специалисты могут принимать обоснованные решения о развертывании модели, распределении ресурсов и конфигурациях системы.

Пример с Batch Size 64 на 1, 2 и 4 GPU

На рисунке а показана производительность модели LLaMA-3-8B на 1, 2 и 4 GPU, а на рисунке b представлено сравнение для модели Mixtral-8x7B при использовании различных видов параллелизма в рамках одного узла. Несмотря на то, что гибридный подход к распараллеливанию обеспечивает большую гибкость, тензорный параллелизм (TP) оказывается эффективнее благодаря более высокой утилизации устройств и меньшим накладным расходам на обмен данными. Наши наблюдения показывают, что при использовании LLaMA-3-8B на 4-х GPU A100, TP в 1,30 раза быстрее гибридного подхода (TP=2, PP=2) и в 1,94 раза быстрее, чем конвейерный параллелизм (PP).

Схемы тензорного параллелизма постепенно усложняются: сначала — 1-D (Megatron-style) с делением весов только по одной оси и одной коллективной операцией all-gather/reduce-scatter; затем 2-D, где веса и активации шардируются по двум измерениям сетки устройств, уменьшая трафик примерно в √p раз; дальше 2.5-D, добавляющая к квадратной сетке дополнительную ось реплик для лучшего баланса между памятью и коммуникациями; и, наконец, 3-D, объединяющая тензорное, пайплайновое и дата-параллельное разбиения, чтобы обслуживать сотни миллиардов параметров.

Каждая следующая ступень экономит сеть или память по сравнению с предыдущей, но требует более сложных коллективных операций и поддерживающего кода; реализованы эти подходы в библиотеках вроде Megatron-LM, Colossal-AI, DeepSpeed и JAX pjit.

Схема разделения на столбцы или ряды при 1-D.

Проще всего на тензорный делить классический decoder-only Transformer с много-головым (MHA) вниманием и одинаковыми, кратными числу GPU размерами скрытых слоёв (GPT-3/NeoX, LLaMA 1-3, Mistral 7B, Falcon 40B и т.п.).

TP делит матрицы Q K V и MLP-слоёв по колонкам/строкам или по головам, и каждая карта считает «свою» долю — коммуникация нужна только для итогового All-Reduce на выходе слоя, что масштабируется почти линейно при NVLink/NVSwitch (≈ 85-95 % эффективности на 8 GPU)

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

А если 1 GPU разделен на 2 виртуальных?

Большую VRAM одной мощной карты (A100 80 GB, H100 94 GB, RTX 6000 Ada 48 GB и т.д.) можно «нарезать» на несколько виртуальных GPU (профиль vGPU, Hyper-V GPU-PV) или, на архитектурах Ampere/Hopper, на MIG-инстансы. Такое деление не ускорит одну «тяжёлую» модель через tensor parallelism: каждая виртуальная доля получает лишь часть потоковых мультипроцессоров, тензорных ядер и памяти, а добавляется ещё синхронизация между инстансами. На практике один рабочий поток почти всегда быстрее на «цельном» GPU.

Теоретически не снижает качество инференса, но может быть фактором детерминирования из-за чего вывод может незначительно отличаться по сравнению с моделью, которая помещается на GPU или при конвейерном параллелизме.

Как и в случае со скоростью инференса, использование Ollama и конвейерного параллелизма создает стереотип, что скорость шины не важна, или минимальна и можно даже на майнинговых райзерах 1x PCIE работать. Однако при тензорном параллелизме узким местом становится шина PCIE и может привести к тому, что пользы не будет или наоборот замедлит инференс. Поэтому в HPC-решениях используют SXM/NVLINK, а Infiniband для соединения узлов.

Я замерил на своем сервере и действительно при тензорном параллелизме на два GPU PCIE 3 была полностью забита до теоретического максимума на этапе prefill, загрузка снижалась лишь на этапе decode. В данном обзоре не буду делать свою выкладку по тестированию, пожалуй, лучше в отдельную статье размещу.

Пользователь с reddit показал, что инференс Mistral Large 123B 5bpw на 4x3090 в конфигурации x8 x8 x4 x1 PCI-E 3.0 давал 15-20 т/с, а на другом сервере в режиме x16 PCI-E 4.0 - достигал уже 40 т/с. В то же время принудительное занижение до 3 версии PCIE снизило лишь на 5% скорость.

Еще один замечательный пример наконец-то показывает значимость NVLINK даже для потребительских карт. 3090 карты установлены в режиме PCIE Gen4 x8:

Важно иметь ввиду, особенно при работе на домашнем сервере с multi-GPU, чтобы не спалить проводку или не накрутить счетчики: при тензорном параллелизме все GPU в системе нагружаются одновременно. Значит пиковая потребляемая мощность будет максимальной из возможной, если позволяет, конечно, шина PCIE.

В статье тензорный параллелизм (TP) сравнивается с data parallelism (DP) и pipeline parallelism (PP). TP потребляет больше энергии, чем DP (из-за коммуникаций), но меньше, чем PP (где overhead на синхронизацию слоёв выше). Например: DP: Низкий overhead, энергия ~80–90% на вычисления.TP: ~60–70% на вычисления, 30–40% на коммуникации. PP: Может быть энергоёмким из-за bubble time (простои), до 50% потерь.

Mixed-precision (FP16/bf16) снижает энергопотребление на 20–30% в TP, так как уменьшает объём данных для коммуникаций. Без этого overhead растёт (до +15% энергии).

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

Заглядывайте на мой ТГ-канал, черкану может там чего еще полезного.