Обучение больших моделей на множестве GPU

Ещё десять лет назад слова «обучить языковую модель» означали запуск одного Python-скрипта на одной видеокарте. Сегодня GPT-4, Llama 3, Gemini Ultra — это артефакты кластеров из тысяч GPU, работающих в унисон. Но как именно координируется такой оркестр? Что делать, когда модель просто не помещается в одну видеокарту?

В этой статье — подробный разбор всех ключевых техник распределённого обучения: от классического data parallelism до 3D-параллелизма и ZeRO-оптимизатора. Материал основан на знаковой публикации Lilian Weng и дополнен актуальными практиками 2024–2026 годов.

Почему одного GPU недостаточно

Главный узкий bottleneck при обучении сверхбольших нейросетей — колоссальный спрос на GPU-память, который далеко выходит за пределы одного устройства. Помимо весов модели (десятки миллиардов чисел с плавающей точкой), ещё дороже обходится хранение промежуточных вычислений: градиентов и состояний оптимизатора (например, моментумов и вариаций в Adam).

Для примера: BLOOM — 175-миллиардный трансформер от Hugging Face. Хранение весов в формате bfloat16 требует 350 ГБ, тогда как GPU, использованные при обучении, имели «лишь» 80 ГБ памяти каждый. В итоге финальное обучение было распределено между 384 GPU.

Существует несколько парадигм параллелизма, позволяющих вести обучение на множестве GPU, а также ряд архитектурных и memory-saving решений, делающих возможным тренировку очень больших нейросетей.

Параллелизм данных (Data Parallelism)

Data parallelism равномерно распределяет данные по нескольким GPU. Каждый GPU хранит полную копию модели и параллельно обрабатывает свою порцию данных. По завершении шага результаты синхронизируются и объединяются.

Этот подход масштабируется почти линейно с числом GPU для большинства задач, что делает его предпочтительной стратегией там, где модель помещается в одну карту. Ключевое преимущество — простота: существующий однопроцессорный код часто переводится на data parallelism с минимальными правками.

# PyTorch DistributedDataParallel — минимальный пример
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group(backend="nccl")
model = MyModel().to("cuda")
model = DDP(model)  # оборачиваем модель

# Дальше — обычный train-loop

Самый наивный вариант DP — скопировать одни и те же веса на несколько воркеров и раздать каждому фракцию данных для одновременной обработки. Однако наивный DP не работает хорошо, если размер модели превышает память одного узла.

Эмпирическое правило: data parallelism хорошо масштабируется до сотен GPU. После этого стоимость коммуникации становится слишком высокой.

Параллелизм модели: тензорный и конвейерный

Тензорный параллелизм (Tensor Parallelism)

Когда модель слишком велика для одного GPU (учитывая параметры, состояния оптимизатора, градиенты и активации), её можно разбить на несколько GPU двумя способами: разделив слои по разным GPU или разбив отдельные слои внутри.

Tensor parallelism делит вычисления отдельных слоёв — конкретно матрицы весов — между несколькими GPU. Например, большое матричное умножение внутри трансформерного слоя можно разбить так, что разные GPU вычисляют части результата, которые затем объединяются.

Megatron-LM реализует это, разделяя крупные матрицы весов по нескольким устройствам. Для трансформера это обычно касается блоков Multi-Head Attention (MHA) и Multi-Layer Perceptron (MLP).

Конвейерный параллелизм (Pipeline Parallelism)

Pipeline parallelism сочетает модельный параллелизм с параллелизмом данных, чтобы уменьшить неэффективные «пузыри» простоя. Основная идея: разбить один мини-батч на несколько микробатчей и позволить каждому stage-воркеру обрабатывать по одному микробатчу одновременно.

Межпроцессная коммуникация передаёт только активации (при прямом проходе) и градиенты (при обратном). Расписание проходов и агрегация градиентов варьируются в разных подходах.

«Pipeline parallelism makes it possible to train large models that don’t fit into a single GPU’s memory.» — Pipeline Parallelism Guide

GPipe — один из пионеров подхода:

GPipe достигает почти линейного ускорения пропускной способности при увеличении числа устройств, хотя это не гарантируется, если параметры модели распределены неравномерно между воркерами.

ZeRO и DeepSpeed: революция в управлении памятью

ZeRO — мощный набор техник оптимизации памяти, обеспечивающих эффективное обучение больших моделей с триллионами параметров.

ZeRO использует агрегированные вычислительные и память ресурсы data parallelism, чтобы снизить требования к памяти каждого GPU. ZeRO сокращает потребление памяти каждым GPU, разбивая различные состояния обучения модели: веса, градиенты и состояния оптимизатора.

ZeRO делится на три стадии:

Для моделей до 13 миллиардов параметров можно использовать ZeRO-powered data parallelism без необходимости прибегать к model parallelism, тогда как стандартный data parallelism исчерпает память уже при 1,4 миллиарда параметров.

Ключевое преимущество ZeRO перед альтернативными подходами model parallelism — не нужно модифицировать код модели. Чтобы использовать ZeRO в DeepSpeed, достаточно изменить несколько настроек в конфигурационном JSON-файле. Никаких изменений кода не требуется.

// deepspeed_config.json — включение ZeRO Stage 2
{
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true
  }
}

3D-параллелизм: объединяем всё вместе

Для действительно гигантских моделей ни одна из стратегий в одиночку не достаточна. Современный подход — 3D-параллелизм, сочетающий все три измерения:

graph TD
    A[3D-Параллелизм] --> B[Data Parallelism + ZeRO]
    A --> C[Tensor Parallelism]
    A --> D[Pipeline Parallelism]
    B --> E[Шардирование состояний между репликами]
    C --> F[Разделение весов слоёв внутри узла]
    D --> G[Разделение слоёв между узлами]
    E & F & G --> H[Обучение 100B+ моделей]

Именно так был обучен 176-миллиардный BLOOM: с использованием Megatron-DeepSpeed — комбинации двух технологий. DeepSpeed разработал реализацию 3D-параллелизма, сочетая ZeRO-шардирование и pipeline parallelism из DeepSpeed с tensor parallelism из Megatron-LM.

Популярный и эффективный паттерн — использовать Megatron-LM для TP и PP (эффективных реализаций тензорного и конвейерного параллелизма), в сочетании с DeepSpeed для продвинутых DP-оптимизаций через ZeRO.

Поиск оптимального баланса между размером микробатча, числом pipeline-стадий и степенью tensor parallelism часто требует значительного числа экспериментов.

Выбор стратегии по размеру модели

Архитектурные оптимизации и техники экономии памяти

Помимо стратегий параллелизма, существуют архитектурные приёмы, которые снижают потребление памяти и ускоряют обучение.

Gradient Checkpointing (Activation Recomputation)

Вместо хранения всех промежуточных активаций для обратного прохода — пересчитывать их на лету. Это экономит память за счёт времени (~30% замедление обратного прохода).

# В PyTorch:
from torch.utils.checkpoint import checkpoint

output = checkpoint(my_layer, input_tensor)

Mixed Precision Training (FP16 / BF16)

Обучение в формате fp16 или bfloat16 вместо fp32 вдвое сокращает потребность в памяти и ускоряет матричные вычисления на современных GPU (A100, H100).

Flash Attention

Оптимизированный алгоритм вычисления attention, который избегает материализации полной матрицы внимания в памяти. Критически важен для длинных контекстов.

Mixture of Experts (MoE)

Архитектуры GShard и Switch Transformer используют token choice — каждый токен выбирает одного или двух лучших экспертов. Вспомогательная функция потерь поощряет более сбалансированную нагрузку, но не гарантирует оптимальной производительности. Лимит вместимости эксперта может приводить к «выброшенным» токенам, если эксперт достиг своего предела.

Коммуникации и инфраструктура кластера

Стандартные GPU-кластеры для обучения состоят из нескольких вычислительных узлов, соединённых быстрым Ethernet или специализированным коммуникационным бэкендом типа InfiniBand. Каждый узел содержит несколько GPU. GPU взаимодействуют с CPU и RAM через PCIe. GPU внутри одного узла обычно соединены быстрым интерконнектом вроде Nvidia NVLink.

Model parallelism вертикально разрезает модель, разделяя вычисления и параметры в каждом слое между несколькими устройствами, что требует значительного межслоевого общения. Это хорошо работает внутри одного узла с высокой пропускной способностью GPU-интерконнекта, но эффективность резко падает при выходе за пределы одного узла.

Именно поэтому tensor parallelism обычно применяется внутри одного узла (по NVLink), а pipeline parallelism — между узлами (по InfiniBand).

Pipeline parallelism в DeepSpeed снижает объём коммуникаций при распределённом обучении, что позволяет обучать модели с миллиардами параметров в 2–7 раз быстрее на кластерах с ограниченной сетевой пропускной способностью.

Заключение: карта решений

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

1. Data Parallelism — отправная точка. Работает до ~100 GPU без особых усилий.
2. ZeRO / FSDP — следующий шаг. Дают «бесплатную» экономию памяти без переписывания модели.
3. Tensor Parallelism — для моделей, не помещающихся даже при ZeRO-3. Требует модификации кода, эффективен внутри узла.
4. Pipeline Parallelism — масштабирование между узлами; важно бороться с «пузырями».
5. 3D-параллелизм — комбинация всех трёх для моделей 70B+.
6. Архитектурные техники (Flash Attention, MoE, Gradient Checkpointing, Mixed Precision) — обязательный комплект поверх любой стратегии параллелизма.

Мир распределённого обучения меняется быстро: появляются ZeRO-Infinity (offload на CPU/NVMe), Ring Attention для длинных контекстов, sequence parallelism. Но фундамент — понимание трёх осей параллелизма и природы коммуникационных накладных расходов — останется актуальным ещё долго.