Как обучать LLM на нескольких GPU: All-Reduce, ZeRO и гибридный параллелизм

Одна GPU — даже H100 или MI300X — едва вмещает модель на 70B параметров в VRAM. Про обучение такой модели за человеческую жизнь речи вообще нет. Современные LLM обучают на сотнях и тысячах GPU. Значит, нужно распределённое обучение: разбить данные и вычисления между устройствами, синхронизировать градиенты и не умереть от простоев.

В предыдущих частях мы разобрали матричное умножение, трансформер-блок, backpropagation, AdamW и mixed-precision. Теперь убираем границы одного устройства. Задача — разобраться, как работают коллективные коммуникации, что такое All-Reduce и ZeRO, и когда хватает Data Parallelism, а когда нужен гибрид.

NCCL и RCCL: общий интерфейс для NVIDIA и AMD

NVIDIA использует NCCL, AMD — RCCL. Хорошая новость: API одинаковый, отличается только префикс: nccl вместо rccl. Обычно это закрывают макросами препроцессора, чтобы писать переносимый код под обе платформы.

Инициализация коммуникатора начинается с уникального ID. Rank 0 генерирует его, остальные ранги получают и создают коммуникатор. В кластере ID раздают через MPI, для одного узла достаточно переменных окружения. Затем запускается All-Reduce.

Data Parallelism и операция All-Reduce

При Data Parallelism каждая GPU хранит полную копию модели. На каждом устройстве считается свой градиент, а затем градиенты усредняются:

dW_global = (1 / world_size) * Σ dW_local_i

Это и есть All-Reduce с операцией SUM. Наивная реализация через центральный сервер создаёт узкое место. Поэтому используют Ring-AllReduce: это оптимальный по пропускной способности алгоритм.

All-Reduce с SUM — это сумма, а не среднее. Делите на world_size.

Ring-AllReduce: два прохода по кольцу

1. Scatter-Reduce — N-1 шаг. Каждая GPU отправляет свою уникальную порцию градиента следующей и получает порцию от предыдущей, накапливая сумму.
2. All-Gather — N-1 шаг. Уже агрегированные порции передаются по кольцу, пока у всех не окажется полная сумма.

В коде это одна функция: ncclAllReduce для CUDA или rcclAllReduce для ROCm. Передаёте указатель на буфер градиентов, размер, тип данных, операцию SUM — и градиенты синхронизированы на всех рангах. Остаётся разделить на world_size — либо отдельным ядром, либо операцией AVG, если библиотека поддерживает.

ZeRO: шардирование памяти

Data Parallelism не спасает память: каждая GPU всё ещё хранит полную модель, градиенты и состояние оптимизатора. Для 175B параметров только веса в FP32 займут около 3.2 TB. ZeRO от Microsoft шардирует эти компоненты между GPU.

ZeRO Stage 1: шардируем состояние оптимизатора

В AdamW у каждого параметра есть момент m и дисперсия v — это два тензора того же размера, что и модель. В Stage 1 они разбиваются по GPU: Rank 0 отвечает за первую половину параметров, Rank 1 — за вторую.

Меняется только ядро AdamW: теперь оно обновляет локальную порцию параметров. Но после обновления нужно сделать All-Gather весов, чтобы в следующем forward-проходе каждая GPU снова имела полную модель. Это дополнительная коммуникация, которая окупается экономией памяти.

ZeRO Stage 2: шардируем градиенты

Stage 2 идёт дальше — шардируются и градиенты. Вместо полного All-Reduce выполняется Reduce-Scatter: каждая GPU суммирует только свою порцию градиентов. Объём пересылок сокращается до половины.

Псевдокод шага обучения с ZeRO-2 выглядит так: forward на полной модели, backward, Reduce-Scatter градиентов, локальное обновление AdamW, All-Gather весов. Именно так работают DeepSpeed и FairScale.

Гибридный параллелизм: Tensor и Pipeline

Когда модель переваливает за триллион параметров, даже ZeRO-3 не спасает. Приходится резать саму архитектуру.

• Tensor Parallelism — веса линейных слоёв разбиваются по строкам или столбцам. Например, матрица W делится по столбцам, и после каждого слоя нужен All-Reduce.
• Pipeline Parallelism — слои разбиваются по GPU: первая держит слои 1–10, вторая — 11–20. Между устройствами передаются активации в forward и градиенты в backward через P2P-коммуникации.

P2P в NCCL/RCCL — это вызовы send/recv между рангами. Один ранг отправляет активации, другой принимает. Главное — не перепутать порядок и согласовать потоки.

Что делать на практике

Если у вас 2–4 GPU и модель влезает в одну — начните с обычного Data Parallelism. Реализуйте All-Reduce, проверьте усреднение, добейтесь линейного ускорения. Дальше включайте ZeRO-1 и ZeRO-2 по мере роста модели.

Типичные грабли:

• Забыли разделить сумму градиентов на world_size.
• Сделали All-Gather весов не после каждого шага — модель рассинхронизируется.
• Используете один поток и для вычислений, и для коммуникаций — просадка производительности.
• Пытаетесь шардировать оптимизатор, но забыли, что локальное обновление нужно только для своей части параметров.

Гибридный параллелизм нужен только тогда, когда модель не влезает в GPU даже после ZeRO. Это уровень больших кластеров, и там свои сложности: балансировка нагрузки между слоями, пузыри в pipeline, пересылки через сеть.

Вывод

Распределённое обучение LLM — это не магия, а три кита: коллективные коммуникации, шардирование памяти и параллелизм вычислений. Начните с All-Reduce, затем добавьте ZeRO, и только если модель всё ещё не помещается — беритесь за гибрид.