Когда tensor parallelism хуже pipeline parallelism?

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

parallelism|Tensor parallelism (TP) и pipeline parallelism (PP) — две основные стратегии распределения вычислений при обучении и инференсе больших моделей. TP разрезает слои по GPU, требуя частых all-reduce на каждом шаге, что делает его неэффективным при очень маленьком batch size (<4) или медленных межсоединениях (PCIe вместо NVLink). PP разрезает модель по слоям, передавая активации только на границах стадий, что даёт меньший объём коммуникаций, но страдает от pipeline bubbles. TP хуже PP, когда overhead explosion|коммуникационные накладные расходы доминируют над временем вычислений.

1. Термин: Tensor Parallelism (TP)

parallelism|Tensor Parallelism — стратегия, при которой один слой нейронной сети (обычно линейный слой или слой внимания) распределяется между несколькими GPU. Каждый GPU хранит часть весов и выполняет вычисления только над своей частью. Для получения полного результата требуется операция all-reduce (суммирование тензоров со всех GPU и рассылка обратно).

Пример: линейный слой Y = X·W, где X — входной тензор размера (batch, hidden), W — матрица весов (hidden, hidden). При TP=2:

• GPU0: W[:, :hidden/2]; GPU1: W[:, hidden/2:].
• Каждый GPU умножает X на свою часть W → частичные выходы.
• All-reduce суммирует их → полный Y.

Термин AllReduce: коллективная коммуникационная операция, объединяющая данные со всех участников и отправляющая результат каждому участнику. В TP используется часто: на каждый прямой и обратный проход каждого слоя.

TP эффективен для моделей с очень широкими слоями (например, hidden=16384), так как позволяет распределить память и вычисления. Однако он требует высокопроизводительной сети с малой задержкой (NVLink, 600 GB/s). При слабой сети overhead коммуникации может превысить выигрыш от параллелизма.

2. Термин: Pipeline Parallelism (PP)

Pipeline Parallelism — стратегия, при которой различные части модели (стадии, состоящие из нескольких последовательных слоёв) размещаются на разных GPU. Входные данные обрабатываются последовательно: сначала на GPU0, затем результат передаётся на GPU1 и т. д. Для увеличения загрузки GPU используется микро-батчинг (разбиение batch на микро-пакеты).

Пример: модель из 8 слоёв, 2 стадии:

• GPU0: слои 1-4; GPU1: слои 5-8.
• Batch делится на 4 микро-батча.
• GPU0 обрабатывает микро-батч 1, передаёт на GPU1, одновременно начинает микро-батч 2, и т. д.

Термин Pipeline bubble (пузырь) — время простоя GPU в начале и конце пайплайна. Доля пузыря при оптимальной настройке: (K-1)/(K+M-1), где K — число стадий, M — число микро-батчей. Чем больше M, тем меньше пузырь.

Подход PP требует лишь point-to-point коммуникации (send/recv) между соседними стадиями — один раз в прямом проходе и один раз в обратном. Это значительно снижает требования к сети по сравнению с TP.

3. Сравнение затрат на коммуникацию

4. Когда TP хуже PP: детальный разбор

4.1 Очень маленький batch size (batch < 4)

При малом batch время вычисления на GPU мало (пропорционально batch), а время all-reduce фиксировано (зависит от размера тензора и пропускной способности сети). Соотношение compute / communication падает.

Расчёт для примера:

• Скрытый слой hidden=4096, A100 80GB: FP16 TFLOPS ~312, NVLink 600 GB/s.
• Размер all-reduce = 2×hidden² = 2×4096² = 33.5 MB (с учётом градиентов).
• Время all-reduce (при идеальном кольце) ≈ 33.5 MB / 600 GB/s = 0.056 ms.
• Время вычисления для batch=1: 1 × 4096 × 4096 × 2 FLOP / 312e12 = 0.108 ms.
• Overhead коммуникации ~52% от вычисления.
• При batch=4: время вычисления 0.432 ms, overhead ~13%.
• При batch=1 TP может быть медленнее PP, где коммуникация редкая (один раз на стадию) и не требует синхронизации.

4.2 Медленная сеть (PCIe вместо NVLink)

PCIe Gen4 x16 обеспечивает ~32 GB/s и имеет бóльшую задержку (~1 мкс vs <0.1 мкс у NVLink). All-reduce по PCIe будет примерно в 20 раз медленнее, чем по NVLink. Для batch=1 overhead коммуникации превысит 1000%, что делает TP практически неприменимым. PP использует point-to-point, которые менее чувствительны к задержке, и может скрывать передачу за вычислениями.

4.3 Модели с тонкими слоями (мало параметров на слой)

Если каждый слой имеет небольшую размерность (например, hidden=512), то выигрыш от распараллеливания внутри слоя минимален (скорость не увеличится в 2 раза из-за закона Амдаля), а накладные расходы на all-reduce (заголовки, синхронизация) съедают выгоду. PP разбивает модель по глубине, обеспечивая более грубое, но эффективное распараллеливание без дополнительной коммуникации на каждый слой.

4.4 Случай: инференс с малым batch (batch=1)

При инференсе (например, обслуживание чат-бота) batch часто равен 1. TP с all-reduce на каждый слой даёт большую задержку. PP может быть реализован с конвейером, где каждый микро-батч загружает все стадии, но при batch=1 пузырь максимален. Однако существуют техники pipeline parallelism for inference (например, разделение по слоям без микро-батчей, с пересылкой активаций), которые дают меньшую задержку, чем TP.

5. Когда PP хуже TP (для контраста)

• Неглубокие модели: если слоёв мало (<10), число стадий ограничено, пузырь становится большим, и PP проигрывает TP.
• Несбалансированные слои: если один слой (например, attention) в 10 раз тяжелее другого, стадия с этим слоем станет узким местом, вызывая простой остальных.
• Очень широкие слои: если слой не помещается в память одного GPU, PP не может его разделить — требуется TP.
• Приложение с жёсткими требованиями к latency: PP увеличивает время первого токена (TTFT), так как нужно последовательно пройти по стадиям; TP может выдать результат быстрее за счёт параллелизма внутри слоя.

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

На практике для крупномасштабного обучения (GPT-3, LLaMA 3) используется комбинация методов:

• TP внутри узла (NVLink) — для распараллеливания широких слоёв.
• PP между узлами (RDMA) — для уменьшения коммуникации по медленным каналам.
• Data Parallelism (DP) — для репликации модели на несколько групп.

Пример: Megatron-LM для модели 175B:

• TP=8 (8 GPU внутри узла, NVLink)
• PP=64 (64 стадии, межузловые соединения)
• DP=2 (2 копии модели по данным)

Таким образом, недостатки TP компенсируются за счёт ограничения области его применения (внутри узла), а межузловая коммуникация ложится на PP.

7. Пример кода (концептуальная симуляция)

import numpy as np

def simulate_tp(layers, hidden, batch, gpus, flops, bw_bytes):
    """
    Приблизительное время TP (игнорируем overhead синхронизации)
    flops: FLOP/s на GPU, bw_bytes: пропускная способность all-reduce (B/s)
    """
    compute_per_layer = 2 * batch * hidden * hidden / flops  # forward+backward
    comm_per_layer = 4 * hidden * hidden / bw_bytes  # float32 all-reduce (4 bytes)
    return layers * (compute_per_layer / gpus + comm_per_layer)

def simulate_pp(layers, stages, microbatch, flops, bw_bytes, hidden):
    """
    Приблизительное время PP (без учета градиентов назад)
    """
    layers_per_stage = layers / stages
    compute_per_stage = 2 * microbatch * layers_per_stage * hidden * hidden / flops
    bubble = (stages - 1) / (stages + (batch/microbatch) - 1)
    comm_per_stage = microbatch * hidden * 4 / bw_bytes  # send activations
    return layers_per_stage * compute_per_stage + bubble * compute_per_stage + comm_per_stage

# Параметры
hidden = 4096
layers = 96
flops = 312e12
bw_nvlink = 600e9  # B/s
bw_pcie = 32e9

for batch in [1, 2, 4, 8, 16]:
    t_tp = simulate_tp(layers, hidden, batch, 8, flops, bw_nvlink)
    t_pp = simulate_pp(layers, 8, min(batch, 4), flops, bw_nvlink, hidden)
    print(f"batch={batch}: TP={t_tp:.4f}s, PP={t_pp:.4f}s, winner={'TP' if t_tp < t_pp else 'PP'}")

При batch=1, PP может оказаться быстрее; при batch=16 — TP.

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

Задача Реализовать симулятор времени выполнения для TP и PP и визуализировать точку пересечения в зависимости от batch size и пропускной способности сети.

Инструменты Python, numpy, matplotlib.

Шаги:

1. Определить параметры: hidden=4096, layers=96, flops=312e12.
2. Реализовать функции time_tp(gpus, batch, bandwidth) и time_pp(stages, batch, bandwidth) с учётом микро-батчей.
3. Для двух сценариев: NVLink (600 GB/s) и PCIe (32 GB/s).
4. Построить графики зависимости времени от batch size (1..32) для TP (gpus=8) и PP (stages=8).
5. Отметить точку, в которой TP становится быстрее PP.
6. Добавить параметр «количество слоёв» и провести анализ чувствительности.

Ожидаемый результат График, демонстрирующий, что при batch < 4 PP выигрывает, при batch > 8 TP выигрывает (на NVLink), а на PCIe TP почти всегда хуже.

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

10. Навигация

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

Навигация

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