Что такое pipeline parallelism и проблема pipeline bubbles?

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

Pipeline parallelism — это техника распределённого обучения моделей (особенно больших языковых), при которой слои нейросети разрезаются на последовательные блоки (стадии) и размещаются на разных GPU. Прямой и обратный проходы выполняются конвейером, но при наивной реализации возникает pipeline bubble — время, когда часть GPU простаивает в ожидании данных или освобождения памяти. Основной способ борьбы — разбиение батча на микробатчи (microbatches) и использование специальных схем планирования (1F1B, interleaved), что резко повышает загрузку оборудования.

1. Определение Pipeline Parallelism

Pipeline parallelism (parallelism|конвейерный параллелизм) — одна из стратегий параллельного обучения моделей, не умещающихся на одном GPU. Модель рассекается по глубине: разные наборы слоёв (стадии) назначаются разным устройствам. Например, в трансформере:

• GPU1: embedding + первые 4 слоя transformer.
• GPU2: следующие 4 слоя.
• GPU3: последние 4 слоя + выходной слой.

Термин стадия (stage) — часть модели, целиком выполняемая на одном устройстве.

Данные проходят через стадии последовательно: сначала обрабатываются на GPU1, результат передаётся GPU2, затем GPU3, и только после этого можно начать обратный проход (градиенты идут в обратном порядке). Такой подход отличается от data parallelism (копия всей модели на каждом GPU, данные делятся) и tensor parallelism (разрезание матричных операций внутри одного слоя). Pipeline parallelism — частный случай model parallelism, когда разделение идёт по вертикали.

2. Наивная (G-Pipe) реализация и возникновение pipeline bubbles

В наивном конвейере (впервые формализован в PipeDream, позже в G-Pipe) используется один батч размером B. Он разбивается на M микробатчей (microbatches) размером B/M. Последовательность действий для одного шага обучения:

1. Forward для всех стадий: сначала все M микробатчей проходят через GPU1, затем все — через GPU2, затем — через GPU3. На каждой стадии микробатчи обрабатываются один за другим.
2. После окончания forward на последней стадии начинается backward (обратный проход): GPU3 обрабатывает градиенты для всех M микробатчей (в порядке, обратном forward), затем GPU2, затем GPU1.
3. После завершения backward на GPU1 — обновление весов.

Проблема: в момент, когда GPU1 обрабатывает первые микробатчи, GPU2 простаивает, ожидая данных от GPU1. Аналогично в конце, когда GPU3 считает градиенты, GPU1 уже простаивает. Временные интервалы простоя называются pipeline bubbles («пузыри»).

Термин «bubble» — период времени, когда устройство не выполняет ни forward, ни backward, а ждёт данные от соседнего устройства или завершения прохода. В наивной реализации bubble возникает как на старте конвейера (пока первая стадия не отправит данные на вторую), так и на финише (последняя стадия завершает backward, а первая уже ничего не делает).

3. Диаграмма pipeline bubbles (иллюстрация)

Представим P = 2 стадии, M = 1 микробатч. Пусть время одного forward на стадии – t_f, одного backward – t_b (обычно t_f ≈ t_b). Временная линия:

• GPU1: forward (t_f) → backward (t_b)
• GPU2: простаивает (ожидает forward от GPU1) → forward (t_f) → backward (t_b)

Общее время шага: 2 * (t_f + t_b). Полезное время на GPU2 = t_f + t_b (50%). Bubble = 50% от времени работы GPU2. Если стадий больше и микробатч один, bubble растёт линейно с числом стадий.

Формула эффективности конвейера при наивной реализации:

Efficiency = (P * M) / (P + M - 1)

Вывод: при фиксированном P увеличивая M (число микробатчей), мы приближаем эффективность к 1. При M = 1 эффективность = 1/P (очень низкая).

4. Решение: микробатчи (microbatches)

Разбиение батча на M микробатчей — ключевой приём. Теперь forward и backward перемежаются: GPU1 обрабатывает микробатч 1, передаёт GPU2, тут же берется за микробатч 2, не дожидаясь завершения backward. Таким образом, большинство стадий заняты почти всё время, за исключением коротких периодов в начале и конце (т.н. remainder bubbles).

Формула остаточного bubble при идеальном планировании:

Bubble_time ≈ (P - 1) * (micropatch_time)

где micropatch_time — время обработки одного микробатча на одной стадии (forward + backward). Общее время шага приблизительно (P + M - 1) * micropatch_time. Доля пузырей:

Bubble_ratio = (P - 1) / (P + M - 1)

Пример: P=4, M=16 => bubble ratio = 3/19 ≈ 15.8% (хорошо). P=4, M=4 => 3/7 ≈ 42.9% (плохо). На практике M выбирают в несколько раз больше числа стадий (типично M = 4..16 для 4 стадий).

5. Сравнение схем планирования

Наивный конвейер (G-Pipe) сначала все forward, потом все backward — даёт большие bubbles и высокое потребление памяти (нужно хранить активации всех микробатчей для backward). Более современные схемы:

1F1B (One Forward One Backward) — стандарт в современных фреймворках (Megatron-LM, DeepSpeed). Память экономится, потому что активации микробатча освобождаются сразу после его backward.

6. Код-симуляция pipeline bubbles (Python)

Ниже упрощённая модель для расчета времени шага и загрузки:

def pipeline_simulator(P, M, t_f, t_b):
    """
    P - число стадий
    M - число микробатчей
    t_f, t_b - время forward и backward на одной стадии
    Возвращает общее время шага и среднюю загрузку GPU
    """
    # В наивном G-Pipe:
    # Forward: P * M * t_f
    # Backward: P * M * t_b
    # Но конвейерное наложение даёт формулу:
    # total_time = (P + M - 1) * (t_f + t_b)
    total_time = (P + M - 1) * (t_f + t_b)
    useful_time = P * M * (t_f + t_b)  # суммарное полезное время всех GPU
    efficiency = useful_time / (P * total_time)
    return total_time, efficiency

# Примеры
for P, M in [(2,1), (2,4), (4,1), (4,8), (4,16)]:
    total, eff = pipeline_simulator(P, M, t_f=10, t_b=10)
    print(f"P={P}, M={M}: total_time = {total}ms, efficiency = {eff:.2%}")

Вывод:

P=2, M=1: total_time = 40ms, efficiency = 50.00%
P=2, M=4: total_time = 100ms, efficiency = 80.00%
P=4, M=1: total_time = 80ms, efficiency = 25.00%
P=4, M=8: total_time = 220ms, efficiency = 72.73%
P=4, M=16: total_time = 380ms, efficiency = 84.21%

7. Практические аспекты в обучении больших моделей

Pipeline parallelism — неотъемлемая часть 3D параллелизма (data + tensor + pipeline), используемая в тренировке GPT-3, BLOOM, LLaMa и др. При выборе количества стадий (P) и числа микробатчей (M) руководствуются:

• Memory budget: микробатчи увеличивают градиентный аккумулятор, но требуют памяти для активаций. 1F1B снижает требования.
• Bubble ratio: стремятся к (P-1)/(P+M-1) ≤ 0.2 (20% простой). Для P=4 нужно M ≥ 16.
• Communication latency: передача данных между стадиями — по NVLink или InfiniBand. Её стоимость включается в t_f + t_b.
• Imbalance: разные стадии могут иметь разную вычислительную нагрузку (embedding быстрее, чем attention). Профилирование и balance-aware partition (например, с помощью torch.distributed.pipeline или DeepSpeed) обязательны.

8. Проблема дисбаланса стадий и её решение

Если стадии имеют разное время выполнения, bubble растёт из-за того, что быстрая стадия вынуждена ждать медленную. Способы борьбы:

• Перераспределение слоёв: torch.distributed.pipeline.sync.Pipeline в PyTorch использует динамическое профилирование.
• Разбивка одного слоя на несколько стадий (finer-grained pipeline).
• Использование multi-stream или асинхронной передачи данных.

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

Задача: Написать симулятор обучения модели с pipeline parallelism, который визуализирует загрузку GPU во времени.

• Инструменты: Python, numpy, matplotlib, scipy (опционально). Можно использовать torch.distributed.pipeline для реального обучения, но для понимания достаточно симуляции.
• Шаги:
  1. Определить P, M, t_f, t_b (задать константы или симулировать случайные задержки).
  2. Реализовать две стратегии: G-Pipe и 1F1B.
  3. Для каждой стратегии построить Gantt-диаграмму (занятость каждого GPU по времени).
  4. Посчитать долю bubbles и сравнить с теоретической формулой.
• Ожидаемый результат: график, показывающий, как с ростом M уменьшается время простоя; наглядное отличие naive от 1F1B.

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

Навигация

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