Распределённое обучение: DP, TP, PP и ZeRO

Что делать, когда модель или данные не помещаются на одну GPU: data, tensor и pipeline parallelism, шардирование состояний ZeRO/FSDP и как из них собирают 3D-параллелизм для обучения больших LLM.

Интуиция: два разных «не тянет»

Одна GPU перестаёт справляться по двум очень разным причинам, и лечатся они по-разному. Первая: модель влезает в память, но данных так много, что одна карта будет учиться месяцами — нужно распараллелить данные. Вторая: модель сама не помещается в память карты — нужно резать уже модель.

Насколько легко упереться во второй случай, показывает арифметика из тем про fine-tuning и VRAM: обучение с Adam в смешанной точности требует ~16 байт на параметр — веса fp16 (2) + градиенты (2) + состояния Adam m, v (8) + fp32-мастер-веса (4). Модель 7B — это уже ~112 ГБ, больше любой одиночной карты; 70B — свыше терабайта.

Data parallelism: копируем модель, режем данные

Параллелизм по данным (data parallelism, DP) — самый простой уровень: на каждой GPU лежит полная копия модели, но батчи всем достаются разные. Каждая карта считает свой forward/backward, после чего градиенты усредняются операцией all-reduce, и все копии делают одинаковый шаг — модели остаются синхронными. Пропускная способность обучения растёт почти линейно по числу карт.

Узких мест два. Во-первых, коммуникация: синхронизировать градиенты (по объёму — как сами веса) нужно на каждом шаге, и на медленной сети обучение начинает ждать обменов, а не считать. Во-вторых, память: DP не экономит её вообще — каждая карта тащит все 16 байт на параметр.

Режем модель: tensor и pipeline parallelism

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

• Tensor parallelism (TP) — режем поперёк: каждая матрица слоя делится между GPU (например, по колонкам), и карты считают один и тот же слой сообща, обмениваясь активациями внутри каждого слоя. Обменов много и они синхронные, поэтому TP живёт только на быстром интерконнекте NVLink — то есть внутри одного узла (обычно до 8 GPU).
• Pipeline parallelism (PP) — режем вдоль: первые слои на GPU 1, следующие на GPU 2 и так далее — модель превращается в конвейер, между стадиями летят только активации точка-точка. Плата — пузырь конвейера (pipeline bubble): в начале шага последние стадии ждут, пока данные до них дойдут, а в конце первые уже простаивают. Лечение — резать батч на много микробатчей (micro-batches), чтобы стадии работали внахлёст (схемы GPipe, 1F1B); чем больше микробатчей, тем меньше доля простоя.

ZeRO и FSDP: убираем избыточность DP

Посмотрим на DP ещё раз: N карт хранят N одинаковых копий оптимизатора, градиентов и весов. Это чистая избыточность, и ZeRO (Zero Redundancy Optimizer; в PyTorch — FSDP) её убирает, шардируя состояния по картам. Стадии наращивают агрессивность: ZeRO-1 шардирует только состояния оптимизатора, ZeRO-2 — ещё и градиенты, ZeRO-3 — даже сами веса: каждая GPU постоянно хранит лишь 1/N всего, а полные веса слоя собирает по all-gather непосредственно перед его вычислением и тут же отбрасывает. Память на карту падает почти в N раз — ценой дополнительных коммуникаций каждый шаг.

Обрати внимание: даже скромная 7B в режиме DP требует 112 ГБ на каждую карту — обучение не влезает уже на старте. Переключись на ZeRO-3 — те же 112 ГБ нарезаются на четыре шарда по 28 ГБ, и всё помещается. Теперь потяни размер модели вверх: около 20B кончается и это (4 × 80 ГБ — потолок для полного файнтюна), дальше нужны либо десятки GPU, либо PEFT. TP и PP дают ту же память на карту, но другой ценой: TP непрерывно гоняет активации внутри каждого слоя, а у PP появляется пузырь простоя на таймлайне внизу.

3D-параллелизм: как учат фронтир-модели

Стратегии не исключают, а дополняют друг друга. Большие LLM учат 3D-параллелизмом: TP режет слои между 8 GPU внутри узла (там есть NVLink), PP раскладывает группы слоёв по узлам (между узлами связь медленнее, а PP общается редко), а поверх всего DP гонит через получившиеся «копии» разные части данных. Плюс ZeRO по оси DP и рекомпьютация активаций (gradient checkpointing), чтобы не хранить весь forward.

Практическое правило выбора на собеседовании: сначала DP (+ZeRO) — самый простой и хорошо масштабируемый уровень; если модель всё ещё не влезает — TP, но только внутри узла; если узла мало — PP между узлами. И всегда держать в голове альтернативу: возможно, полный файнтюн вообще не нужен, и LoRA решит задачу на одной карте.