Как работает tensor parallelism для LLM инференса? В чем отличие от pipeline parallelism?

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

parallelism|Tensor parallelism (TP) разрезает веса одного слоя между несколькими GPU, позволяя каждому GPU вычислять свою часть активаций, после чего результаты объединяются через AllReduce. Pipeline parallelism (PP) делит модель по слоям: каждый GPU отвечает за целый набор последовательных слоёв, и данные передаются между GPU последовательно. Для инференса LLM на одном узле (8 GPU) чаще используют TP, так как он даёт меньший pipeline bubble и более равномерную загрузку, но требует высокой пропускной способности межсоединений (NVLink). PP эффективен при большом числе GPU (межузловой) и позволяет уменьшить объём памяти на каждом устройстве.

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

LLM (Large Language Model) с десятками или сотнями миллиардов параметров не помещаются в память одного GPU (например, 70B модель требует ~140 ГБ в FP16). Даже если модель помещается, latency одного forward pass может быть слишком велик. Параллелизм позволяет:

• Распределить модель по нескольким GPU (model parallelism)
• Ускорить вычисления за счёт параллельной обработки разных частей данных (data parallelism) или разных частей модели (pipeline, tensor)

Для инференса (в отличие от обучения) важны latency (время ответа) и throughput (количество запросов в секунду). Разные виды параллелизма по-разному влияют на эти метрики.

2. Tensor parallelism (TP) — детальный разбор

2.1 Идея

В TP каждый слой модели (например, слой self-attention или feed-forward network) разрезается на части по определённой оси. Каждый GPU хранит только свою часть весов и вычисляет соответствующую часть выходных активаций. После вычисления необходимо синхронизировать результаты между GPU.

2.2 Разрезание attention

В multi-head attention головы внимания независимы. TP распределяет головы между GPU. Например, для модели с 32 heads и 4 GPU: каждый GPU получает 8 heads. После вычисления внимания каждый GPU имеет частичный output, который нужно объединить (concat) и затем применить output projection. Output projection тоже разрезается: его веса делятся по столбцам, и после умножения на частичный input требуется AllReduce для суммирования вкладов от всех GPU.

2.3 Разрезание feed-forward (FFN)

FFN состоит из двух линейных слоёв: W1 (расширение) и W2 (сжатие). TP разрезает W1 по столбцам (column-wise), а W2 — по строкам (row-wise). После первого слоя каждый GPU имеет часть промежуточных активаций, после второго — частичные выходы, которые суммируются AllReduce.

2.4 Коммуникация: AllReduce

После каждого разрезанного слоя требуется AllReduce — операция, которая суммирует значения со всех GPU и рассылает результат каждому. В современных GPU (NVIDIA) AllReduce оптимизирован через NVLink и NCCL. Затраты на коммуникацию пропорциональны размеру активаций, а не весов.

2.5 Пример кода (концептуальный)

# Псевдокод для tensor parallel linear layer
class TensorParallelLinear:
    def __init__(self, in_features, out_features, world_size, rank):
        # Разрезаем веса по столбцам (out_features)
        self.weight = weight[rank * chunk_size : (rank+1) * chunk_size, :]
        self.bias = bias[rank * chunk_size : (rank+1) * chunk_size] if bias else None

    def forward(self, x):
        # x — полный вход (batch, in_features)
        local_out = torch.matmul(x, self.weight.T)  # (batch, chunk_size)
        # AllReduce суммирует local_out со всех GPU
        out = all_reduce(local_out)
        if self.bias is not None:
            out += self.bias
        return out

2.6 Преимущества и недостатки TP

3. Pipeline parallelism (PP) — детальный разбор

3.1 Идея

PP делит модель по слоям: GPU 1 выполняет слои 1–4, GPU 2 — слои 5–8 и т.д. Данные (микробатчи) последовательно проходят через все GPU. Чтобы не простаивать, используют microbatching: один большой батч разбивается на несколько микробатчей, и GPU начинают обрабатывать следующий микробатч, не дожидаясь завершения всего батча.

3.2 Pipeline bubble

Пузырь (bubble) — время простоя GPU в начале и конце пайплайна. При идеальном заполнении (количество микробатчей >> числа GPU) доля пузыря мала. Для инференса с маленькими батчами (batch size = 1) пузырь может быть значительным.

3.3 Пример схемы

Пусть 4 GPU, 4 микробатча. Время выполнения одного микробатча на одном GPU — 1 единица. Идеальный пайплайн:

Пузырь: первые 3 шага GPU1–3 простаивают, последние 3 шага GPU0–2 простаивают. Всего 7 шагов вместо 4 (идеал). Эффективность = 4/7 ≈ 57%.

3.4 Преимущества и недостатки PP

4. Сравнение TP и PP

5. Комбинированный подход: 3D parallelism

На практике для больших моделей (например, GPT-3 175B) используют комбинацию:

• Data parallelism (DP) — реплика модели на нескольких узлах, каждый обрабатывает свой батч.
• Tensor parallelism (TP) — внутри узла (8 GPU).
• Pipeline parallelism (PP) — между узлами (например, 4 узла по 8 GPU).

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

6. Особенности для инференса LLM

• Latency-sensitive (чат-боты): предпочтительнее TP, так как нет pipeline bubble. Но TP требует дорогих GPU с NVLink.
• Throughput-oriented (пакетная обработка): можно использовать PP с большим batch size, чтобы минимизировать пузырь.
• Memory-bound: PP позволяет разместить модель на нескольких узлах с медленной связью, где TP неэффективен.
• KV cache: при TP каждый GPU хранит свою часть KV cache, что увеличивает общий объём, но уменьшает latency.

7. Практические рекомендации

• Если модель помещается в память одного узла (8 GPU) и есть NVLink → используйте TP.
• Если модель не помещается в один узел → используйте PP между узлами, а внутри узла TP.
• Для инференса с batch size = 1 (одиночные запросы) TP даёт наименьший latency.
• Для инференса с большим batch (offline) можно использовать PP, но лучше комбинировать с TP.

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

Задача Реализовать симуляцию tensor parallelism и pipeline parallelism для инференса небольшой модели (например, 2 слоя attention) на 4 виртуальных GPU.

Инструменты Python, PyTorch (можно без реальных GPU, только симуляция времени), библиотека torch.distributed (опционально), mpi4py или multiprocessing.

Шаги:

1. Создайте простую модель из двух линейных слоёв (имитация FFN).
2. Реализуйте TP: разрежьте веса по столбцам, добавьте AllReduce (можно симулировать через torch.distributed.all_reduce или просто сложить вручную).
3. Реализуйте PP: разделите слои между процессами, передавайте активации через очередь (multiprocessing.Queue).
4. Замерьте время выполнения одного forward pass для TP и PP при разном количестве микробатчей.
5. Постройте график зависимости latency от batch size.

Ожидаемый результат Вы увидите, что TP даёт меньшее время для batch=1, а PP начинает догонять при увеличении числа микробатчей. Также можно измерить pipeline bubble.

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

Навигация

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