Tensor parallelism vs pipeline parallelism vs data parallelism — сравните для LLM инференса?

Краткий тезис

При инференсе больших языковых моделей (LLM) GPU|один GPU часто не вмещает модель целиком или не обеспечивает нужную пропускную способность. Три основные стратегии параллелизма — data parallelism (DP), tensor parallelism (TP) и pipeline parallelism (PP) — решают разные проблемы: DP дублирует модель для увеличения throughput, TP разрезает тензоры внутри одного слоя для уменьшения памяти на GPU, PP разрезает модель по слоям для работы с очень глубокими сетями. Для инференса LLM оптимальной считается комбинация TP внутри узла и PP между узлами, так как она балансирует коммуникационные затраты и использование памяти.

1. Зачем нужен параллелизм при инференсе LLM?

LLM (Large Language Model) — модель с сотнями миллиардов параметров (например, GPT-3 — 175B, LLaMA-2–70B). Даже один прямой проход (forward pass) требует:

• Память: веса модели (например, 70B параметров в FP16 ≈ 140 ГБ), KV-cache для длинных контекстов (ещё десятки ГБ).
• Вычисления: умножение матриц огромной размерности.

GPU|Один GPU (например, A100 80 ГБ) не может вместить модель 175B. Даже если вмещает, latency одного запроса может быть высокой, а throughput — низким. Параллелизм позволяет:

• Распределить модель по нескольким GPU (model parallelism).
• Обрабатывать несколько запросов одновременно (data parallelism).

2. Data Parallelism (DP) — реплики модели

Идея: каждый GPU хранит полную копию модели. Входящие запросы (батчи) делятся между GPU. После forward pass результаты собираются (например, через all-gather).

Для инференса:

• Плюсы: прост в реализации, не требует разрезания модели. Увеличивает throughput линейно с числом GPU (если батч достаточно велик).
• Минусы: каждый GPU держит полную модель — не решает проблему памяти. Коммуникация только для сбора результатов (небольшая). Не подходит, если модель не влезает в один GPU.

Когда использовать: когда модель уже помещается в один GPU, но нужно обрабатывать много запросов (высокий трафик). Часто комбинируется с TP/PP для больших моделей.

3. Tensor Parallelism (TP) — разрезание тензоров

Идея: один слой (например, attention или FFN) разрезается на части по измерению скрытого размера (hidden dimension). Каждый GPU вычисляет свою часть, затем результаты объединяются через all-reduce.

Пример для attention:

• Пусть hidden size = 4096, число голов = 32.
• При TP=4 каждый GPU вычисляет 8 голов (4096/4 = 1024 скрытых размеров на GPU).
• После вычисления attention heads выполняется all-reduce для объединения.

Особенности:

• Коммуникация: интенсивная (all-reduce после каждого слоя). Требует быстрой связи между GPU (NVLink, NVSwitch). Поэтому TP эффективен только внутри одного узла (8 GPU на A100).
• Память: каждый GPU хранит 1/N весов слоя, что позволяет разместить модель, не влезающую в один GPU.
• Latency: добавляет накладные расходы на коммуникацию, но для больших моделей выигрыш в памяти окупается.

Формула коммуникационных затрат (приблизительно):

Communication_time ≈ 2 * (hidden_size / TP) * (число слоёв) / bandwidth

Коэффициент 2 из-за all-reduce (reduce + broadcast).

4. Pipeline Parallelism (PP) — разрезание по слоям

Идея: модель делится на последовательные стадии (stages). Каждый GPU отвечает за несколько слоёв. Данные передаются от стадии к стадии (point-to-point коммуникация).

Пример: модель из 40 слоёв, PP=4 → GPU0: слои 1–10, GPU1: 11–20, GPU2: 21–30, GPU3: 31–40.

Для инференса:

• Плюсы: позволяет разместить очень глубокие модели. Коммуникация только между соседними стадиями — не требует NVLink, можно использовать Ethernet (между узлами).
• Минусы: возникает bubble (простой GPU) — когда одна стадия ждёт данные от предыдущей. Для инференса bubble меньше, чем для обучения, так как нет обратного прохода. Однако latency увеличивается пропорционально числу стадий (каждый запрос проходит через все стадии последовательно).
• Micro-batching: для уменьшения bubble можно разбить батч на микро-батчи и запускать их конвейером (например, GPipe, 1F1B). Но для инференса с низкой latency это редко применяется.

Когда использовать: когда модель не влезает в один узел (межузловой параллелизм). Часто комбинируется с TP.

5. Сравнение трёх стратегий

6. Комбинация TP + PP — стандарт для LLM инференса

На практике для больших моделей (70B+) используют гибридный подход:

• TP внутри узла (8 GPU) — разрезает каждый слой, уменьшая память и используя быструю NVLink.
• PP между узлами — разрезает модель по слоям, позволяя масштабироваться на десятки узлов.

Пример конфигурации для модели 175B:

• 4 узла по 8 GPU A100 (всего 32 GPU).
• TP=8 (внутри узла), PP=4 (4 узла).
• Каждый GPU хранит: (веса слоя / 8) * (число слоёв на стадию). Итоговая память — ~1/32 от полной модели.

Почему не DP DP не уменьшает память, поэтому для моделей, не влезающих в один GPU, он бесполезен. Однако DP можно добавить сверху (например, реплицировать TP+PP конфигурацию для увеличения throughput).

7. Особенности инференса vs обучения

• Нет backward pass: при инференсе не нужно хранить активации для градиентов, что снижает требования к памяти. Поэтому PP bubble меньше.
• KV-cache: при генерации токенов cache занимает много памяти. TP помогает распределить KV-cache между GPU (каждый GPU хранит свою часть).
• Sequence parallelism: разновидность TP, где разрезается не hidden dimension, а длина последовательности. Используется для очень длинных контекстов (например, в Megatron-LM).
• Latency-critical: для инференса в реальном времени (чат-боты) важна низкая latency. TP добавляет фиксированную задержку на коммуникацию, PP — последовательную задержку. Поэтому для low-latency часто выбирают TP без PP (если модель влезает в один узел).

8. Практические фреймворки

• vLLM: использует TP для распределения модели по GPU. Поддерживает PagedAttention для эффективного управления KV-cache.
• TensorRT-LLM: оптимизирует TP и PP, поддерживает in-flight batching.
• DeepSpeed Inference: автоматически выбирает TP/PP на основе профилирования.
• Megatron-LM: оригинальная реализация TP+PP для обучения, адаптирована для инференса.

Пример конфигурации для DeepSpeed Inference (модель 70B на 4 узлах по 8 GPU):

{
  "train_batch_size": 1,
  "tensor_parallel": {"enabled": true, "tp_size": 8},
  "pipeline_parallel": {"enabled": true, "pp_size": 4}
}

9. Выбор стратегии: практические рекомендации

Пет-проект для закрепления

Задача: Сравнить latency и throughput для инференса модели LLaMA-2–7B (или 13B) при использовании DP, TP, PP и TP+PP на кластере из 2 узлов по 4 GPU.

Инструменты:

• Python, PyTorch, Hugging Face Transformers.
• DeepSpeed (с поддержкой inference) или vLLM.
• NVIDIA Nsight Systems для профилирования.

Шаги:

1. Установить DeepSpeed и скачать модель (например, meta-llama/Llama-2-7b-hf).
2. Написать скрипт инференса с разными конфигами:
   • DP: deepspeed --num_gpus 8 inference.py (модель реплицируется).
   • TP: deepspeed --num_gpus 8 inference.py --tp 8.
   • PP: deepspeed --num_gpus 8 inference.py --pp 8 (на одном узле).
   • TP+PP: deepspeed --num_gpus 8 --num_nodes 2 inference.py --tp 4 --pp 4.
3. Замерить:
   • Latency: время генерации одного токена (медиана по 100 запросам).
   • Throughput: количество токенов в секунду при batch size = 1, 4, 16.
   • Memory: занятая память на GPU (nvidia-smi).
4. Построить таблицу и графики.

Ожидаемый результат:

• DP даст высокий throughput при batch size > 1, но не уменьшит память.
• TP уменьшит память, но добавит коммуникационную задержку.
• PP увеличит latency, но позволит работать с моделью, не влезающей в один узел.
• TP+PP — компромисс: умеренная latency, возможность масштабирования.

Связь с другими вопросами

Навигация

• Предыдущий: 202
• Следующий: 204
• Индекс: 00. Индекс разборов