Что такое 3D parallelism (data + tensor + pipeline)?

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

parallelism|3D parallelism — это комбинация трёх уровней параллелизма, используемая для распределённого обучения больших моделей (100B+ параметров): data parallelism (разделение батча по репликам модели), tensor parallelism (разделение одного слоя между GPU внутри узла) и pipeline parallelism (разделение слоёв модели между узлами). Такой подход позволяет преодолеть ограничения памяти и пропускной способности коммуникации, масштабируя обучение на тысячи GPU. Ключевой вызов — балансировка вычислений и коммуникаций, а также эффективное расписание микро-батчей для минимизации bubble (простоя в pipeline).

1. Термин «3D Parallelism»

parallelism|3D parallelism (parallelism|трёхмерный параллелизм) — это стратегия распределённого обучения, при которой модель и данные одновременно разбиваются по трём ортогональным осям:

• Data Parallel|Data parallelism (DP) – по батчам данных.
• parallelism|Tensor parallelism (TP) – по частям одного тензора/слоя.
• Pipeline parallelism (PP) – по группам слоёв (стадиям).

Каждый уровень решает свою проблему: DP — утилизация большого количества GPU при недостаточном батч-сайзе, TP — память для гигантских слоёв, PP — последовательное размещение слоёв на нескольких устройствах.

Вместе они образуют 3D-решётку, где каждое устройство выполняет часть forward/backward для своего фрагмента данных, части слоя и фрагмента последовательности слоёв.

Пример типичной конфигурации для обучения GPT-3 175B:

• 4 узла (node), в каждом 8 GPU (A100 80GB)
• DP=8 (реплик модели)
• TP=8 (внутри узла)
• PP=4 (между узлами)
• Всего GPU: 8 × 4 × 8 = 256? Нет: количество моделей = DP × PP? На самом деле PP и TP образуют одну «модельную параллельную группу»: количество GPU в такой группе = TP × PP. DP — количество таких групп. Итого: DP × TP × PP = 8 × 8 × 4 = 256 GPU.

2. Data Parallelism (DP)

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

Формула времени шага:

T_DP = compute_time + comm_time(DP)

где comm_time(DP) = время all-reduce, растущее с числом реплик и размером модели.

Градиентная аккумуляция — псевдо-DP, когда микро-батчи выполняются последовательно на одной GPU, но градиенты суммируются. Настоящий DP использует несколько GPU.

Недостаток: при малом батч-сайзе утилизация падает; при большом — сходимость ухудшается. Также DP не экономит память (каждая реплика хранит всю модель), поэтому для моделей > 80B на A100 80GB одного DP недостаточно.

3. Tensor Parallelism (TP)

Tensor parallelism (внутрислойный параллелизм) — разделение вычислений одного слоя, например, матричного умножения A @ B, по столбцам или строкам между несколькими GPU. Каждая GPU держит только часть весов слоя, выполняет частичное умножение, затем результаты объединяются через all-reduce (для линейного слоя) или all-gather (для слоя attention).

Пример с self-attention в Transformer:

• Разделяем головы внимания: каждая GPU вычисляет subset голов.
• После завершения all-reduce по всему слою.

TP требует очень быстрой внутриузловой связи (NVLink, NVSwitch) и обычно ограничен одним узлом (2–8 GPU). Гранулярность — уровень тензора, поэтому количество операций коммуникации высокое.

Формально: W = [W1, W2] (по столбцам), Y = X @ W = [X@W1, X@W2], затем коммуникация. Для backward — аналогично.

Преимущество: радикальное снижение памяти на одну GPU (в TP_size раз). Недостаток: высокий overhead на связь (почти каждый слой требует синхронизации).

4. Pipeline Parallelism (PP)

Pipeline parallelism (конвейерный параллелизм) — разделение модели по глубине: каждый узел (или группа GPU) отвечает за последовательный набор слоёв (stage). Обработка батча разбивается на микро-батчи; конвейер запускает их один за другим с задержкой (pipeline depth).

Ключевая метрика — pipeline bubble: доля времени, когда часть GPU простаивает в начале и конце конвейера. Для идеального конвейера из P стадий и M микро-батчей bubble = (P-1)/(M+P-1).

Существует несколько схем:

• GPipe (1F1B): последовательное заполнение, затем слив.
• 1F1B с interleaving (чередование): разбиение стадий на меньшие блоки для уменьшения bubble.
• Chimera: более сложное расписание.

PP требует передачи промежуточных активаций между стадиями (через p2p коммуникацию: send/recv). Коммуникация не масштабируется по всем GPU, а только между соседями, что делает её дешевле, чем all-reduce.

Преимущество: малое количество коммуникаций между узлами, хорошо для медленных сетей. Недостаток: сложность балансировки вычислительной нагрузки стадий и bubble.

5. Комбинация трёх уровней (3D)

В 3D параллелизме каждая GPU имеет три идентификатора: (dp_rank, tp_rank, pp_rank). Коммуникация происходит в разных группах:

• Data-parallel group: all-reduce градиентов между репликами (одинаковые tp_rank, pp_rank).
• Tensor-parallel group: all-reduce промежуточных активаций (одинаковые dp_rank, pp_rank).
• Pipeline group: p2p send/recv активаций между соседними стадиями (одинаковые dp_rank, tp_rank).

Выравнивание батчей: сначала батч делится на DP×PP микро-батчей (т.е. каждая реплика DP получает свой микро-батч, каждая стадия PP — тоже). TP не влияет на батч.

(в тексте описана).

Пример распределения GPU 64

• DP=4, TP=4, PP=4 → всего 4×4×4=64 GPU.
• Внутри одного узла (8 GPU): TP=4, PP=2 (т.е. узел содержит две стадии по 4 GPU каждая).

6. Где применяется

3D parallelism — стандарт для обучения всех современных больших языковых моделей: GPT-3 (175B), Llama 2/3 (70B, 405B), PaLM (540B), Megatron-Turing NLG (530B). Например:

• GPT-3: DP=16, TP=8, PP=4, всего 512 GPU A100.
• Llama 3 405B: используется 3D parallelism с разбивкой на 8 узлов по 8 GPU + FSDP.

Без 3D parallelism обучение модели размером 175B на одном GPU (даже с 80GB) невозможно: только веса занимают 700GB+ (в параметрах float32, а с optimizer states — более 2TB).

7. Оптимальная конфигурация

Выбор соотношения DP / TP / PP зависит от трёх факторов:

Правило: TP≤8 (ограничен внутри узла), PP≤16 (иначе большой bubble), DP — остальное. Например, для 1000 GPU: TP=8, PP=8 → DP=1000/64≈15.6 → выбираем DP=16, тогда общее GPU=8×8×16=1024.

Стандартный тюнинг: запустить профилирование с разными комбинациями, измерить throughput и memory per GPU.

8. Преимущества и недостатки

9. Альтернативы и родственные техники

• FSDP (Fully Sharded Data Parallelism) — Zero-3 в PyTorch: шардирует веса, градиенты и оптимизатор между данными, но без tensor parallelism (внутри слоя). Проще, но менее эффективен для гигантских слоёв.
• Sequence Parallelism — разбиение по длине последовательности (внимание), часто дополняет TP.
• Expert Parallelism (MoE) — для моделей со смесью экспертов, где разные эксперты на разных GPU.
• ZeRO-оптимизация — уменьшает память не параллелизмом, а шардированием состояний.

3D parallelism часто комбинируется с mixed precision training (FP16/BF16) и activation checkpointing (рекомпьютинг активаций вместо хранения).

10. Инструменты и фреймворки

• Megatron-LM (NVIDIA) — эталонная реализация 3D parallelism, поддерживает TP, PP, DP, sequence parallelism.
• DeepSpeed (Microsoft) — Zero-2/3, 1D/2D/3D parallelism, автоматическое разбиение.
• Colossal-AI — более гибкие комбинации.
• PyTorch Distributed (DTensor, FSDP) — частичная поддержка TP через DTensor и FSDP для DP.
• Horovod — только DP.

На практике для современных LLM чаще используется Megatron-LM + DeepSpeed (Zero-3 + TP + PP).

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

Задача: Обучить небольшую модель (например, GPT-2 1.5B) на 4 GPU, используя комбинацию DP=2, TP=2 (PP=1) и сравнить с обучением только на 1 GPU (без параллелизма) и с DP=4.

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

• PyTorch + torch.distributed
• megatron-lm (для 3D parallelism) или transformers + deepspeed (для DP и частичного TP)
• torch.cuda.Event для замера времени шага

Шаги:

1. Установить deepspeed и transformers.
2. Написать конфиг DeepSpeed: zero_optimization.stage=2, tensor_parallel.size=2 (если поддерживается) или через megatron.core.
3. Запустить обучение с разными конфигурациями: 1 GPU, DP=4, DP=2+TP=2.
4. Измерить: throughput (tokens/sec), memory per GPU, время шага, коммуникационный оверхед.
5. Визуализировать график зависимости производительности от числа GPU.

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

• Увидеть, что DP=4 даёт прирост скорости ~3.8x (ускорение не идеальное из-за all-reduce).
• Конфигурация DP=2+TP=2 может показать меньшее ускорение, но позволит обучать модель, не влезающую в память одной GPU.
• Построить сравнительную таблицу.

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

Навигация

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