Distributed Training (Multi-GPU)

Distributed Training: как обучить модель, которая не влезает в одну GPU

Одна GPU — это потолок для маленьких моделей. 7B параметров в FP16 — это ~14 GB только на веса. Добавь Adam optimizer (хранит m и v — ещё ×2 от весов, итого ×3) — уже ~42 GB. Плюс активации при forward pass: ещё 20-50 GB в зависимости от batch size и sequence length. Итого для 7B модели нужно 70-100 GB. A100 80 GB — впритык, и это для самой маленькой из серьёзных моделей.

• 7B (LLaMA 2 7B): ~14 GB веса (FP16) + ~42 GB optimizer + ~30 GB активации ≈ 86 GB. Одна A100 80 GB — на пределе
• 70B (LLaMA 2 70B): ~140 GB только веса. Одна GPU? Забудь
• 405B (LLaMA 3.1 405B): ~810 GB веса. Нужен кластер из сотен GPU
• DeepSeek-V3 (671B MoE): тренировали на 2048 GPU H800 в течение ~2 месяцев

Вывод: если хочешь обучить модель крупнее 7B — нужно распределять вычисления по нескольким GPU. И это не просто "разрезать модель пополам" — есть целый зоопарк стратегий, каждая со своими trade-offs. Разбираемся.

Data Parallelism (DP): каждый GPU считает свой батч

Самый простой подход. Каждый GPU хранит полную копию модели и получает свой кусок данных (micro-batch). Все GPU одновременно делают forward, считают loss, делают backward. Потом синхронизируют градиенты и обновляют веса. По сути: N GPU = N× throughput.

• Каждый GPU: полная копия модели + свой micro-batch данных
• Forward → loss → backward → All-Reduce градиентов → optimizer step
• Эффективный batch size = micro_batch × N_gpu. Throughput растёт линейно (почти)
• Проблема: каждый GPU хранит полную модель + полный optimizer state. Для 70B — невозможно

All-Reduce: как GPU синхронизируют градиенты

All-Reduce — это операция, после которой каждый GPU получает сумму градиентов со всех GPU. Наивный подход: все отправляют всё на один GPU, он суммирует, рассылает обратно. Плохо — bottleneck на одной GPU. Решение — Ring All-Reduce.

Ring All-Reduce: N GPU выстроены в кольцо. Каждый GPU разбивает свои градиенты на N chunks. За N-1 шагов (Reduce-Scatter) каждый GPU получает полную сумму одного chunk. Ещё за N-1 шагов (All-Gather) все получают все суммы. Итого: 2(N-1) шагов, и каждый GPU отправляет ровно 2 × (data_size / N) данных. Bandwidth-optimal — не зависит от количества GPU!

DeepSpeed ZeRO: шардим всё, что можно

Проблема Data Parallelism: каждый GPU хранит полную копию модели, градиентов и optimizer states. Это колоссальное дублирование. ZeRO (Zero Redundancy Optimizer) от Microsoft решает это прогрессивно, в три стадии:

• Stage 1: шардим optimizer states (Adam хранит m и v для каждого параметра). Каждый GPU хранит 1/N optimizer states. Экономия ~×4 памяти
• Stage 2: + шардим градиенты. Каждый GPU хранит градиенты только для "своих" параметров. Экономия ~×8 памяти
• Stage 3: + шардим сами параметры модели. Каждый GPU хранит 1/N модели. Перед compute — All-Gather нужных параметров, после backward — Reduce-Scatter градиентов. Экономия ~×N памяти

Пример: 7B модель с Adam в FP16. Без ZeRO: ~86 GB на каждом GPU. С ZeRO Stage 3 на 8 GPU: ~11 GB на GPU. Разница — как день и ночь.

// DeepSpeed ZeRO Stage 2 config     fine-tuning
{
  "bf16": { "enabled": true },
  "zero_optimization": {
    "stage": 2,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "reduce_bucket_size": 5e8,
    "allgather_bucket_size": 5e8
  },
  "gradient_accumulation_steps": 4,
  "gradient_clipping": 1.0,
  "train_micro_batch_size_per_gpu": 2,
  "wall_clock_breakdown": false
}

// Stage 3    ,    
{
  "zero_optimization": {
    "stage": 3,
    "overlap_comm": true,
    "reduce_bucket_size": 5e8,
    "param_persistence_threshold": 1e5,
    "offload_optimizer": { "device": "cpu" },
    "offload_param": { "device": "cpu" }
  }
}

FSDP: PyTorch-native аналог ZeRO-3

FSDP (Fully Sharded Data Parallel) — это ответ PyTorch на DeepSpeed ZeRO Stage 3. Идея та же: шардим параметры, градиенты и optimizer states по GPU. Но реализация встроена в PyTorch, без внешних зависимостей.

Как это работает: модель разбивается на "единицы шардинга" (обычно — один TransformerBlock). В обычном состоянии каждый GPU хранит только свой шард параметров. Когда нужен forward или backward через конкретный блок:

• Forward: All-Gather параметров блока (собираем полные веса со всех GPU) → compute → освобождаем полные веса
• Backward: All-Gather параметров → compute градиентов → Reduce-Scatter градиентов (каждый GPU получает градиенты для своего шарда) → освобождаем полные веса
• Optimizer step: каждый GPU обновляет только свой шард параметров
• Результат: в любой момент в памяти полные веса только одного блока — экономия колоссальная

import torch
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    MixedPrecision,
    ShardingStrategy,
)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from functools import partial

# Определяем auto-wrap policy: каждый TransformerBlock шардится отдельно
auto_wrap_policy = partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={TransformerBlock},  # ваш класс блока
)

# Mixed precision: compute в bf16, градиенты в fp32
mixed_precision = MixedPrecision(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.float32,
    buffer_dtype=torch.bfloat16,
)

# Оборачиваем модель в FSDP
model = FSDP(
    model,
    auto_wrap_policy=auto_wrap_policy,
    mixed_precision=mixed_precision,
    sharding_strategy=ShardingStrategy.FULL_SHARD,  # = ZeRO-3
    device_id=torch.cuda.current_device(),
)

# Дальше — обычный training loop
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
for batch in dataloader:
    loss = model(batch)       # FSDP сам делает All-Gather/Reduce-Scatter
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

Tensor Parallelism (TP): режем один слой на части

Data Parallelism и ZeRO/FSDP дублируют или шардят модель целиком. Tensor Parallelism идёт глубже: разрезает отдельные слои. Один nn.Linear размером [4096, 16384] разбивается на 8 GPU — каждый считает свою часть матричного умножения.

Megatron-LM style TP для MLP: Linear разбивается на column-parallel (A) и row-parallel (B). Первый Linear (up-projection) режется по столбцам: каждый GPU получает A₁, A₂, ..., Aₙ. Второй Linear (down-projection) режется по строкам: B₁, B₂, ..., Bₙ. Между ними — GeLU (поэлементный, параллелится тривиально). После row-parallel Linear — один All-Reduce, чтобы собрать результат.

• MLP: Y = GeLU(X @ A₁) @ B₁ на GPU₁, ... GeLU(X @ Aₙ) @ Bₙ на GPUₙ → All-Reduce → сумма
• Attention: Q, K, V heads распределяются по GPU. 32 heads, 8 GPU → 4 heads на GPU
• Коммуникация: All-Reduce после каждого TP-слоя. Это дорого → нужен NVLink (900 GB/s), не сетевые карты
• Правило: TP только внутри одной ноды (одного сервера), где GPU связаны NVLink

Pipeline Parallelism (PP): разные слои на разных GPU

TP режет один слой по GPU. PP режет модель по слоям: Layer 0-7 на GPU₀, Layer 8-15 на GPU₁, Layer 16-23 на GPU₂, Layer 24-31 на GPU₃. Каждый GPU считает только свои слои и передаёт activations следующему.

Проблема: если просто гнать один batch через pipeline, в каждый момент работает только одна GPU, остальные простаивают. Это "bubble" — wasted compute. Решение — micro-batching: разбиваем batch на маленькие micro-batches и пускаем их через pipeline конвейером.

• GPipe: все micro-batches forward, потом все backward. Простой, но большой bubble в начале/конце
• 1F1B (one forward, one backward): чередуем forward и backward. GPU быстрее начинают backward → меньше bubble
• Interleaved: каждый GPU хранит несколько непоследовательных chunk-ов слоёв → ещё меньше bubble, но сложнее
• Bubble fraction ≈ (PP_stages - 1) / total_micro_batches. Чем больше micro-batches, тем меньше простой
• PP коммуникация: только activations между соседними stage-ами. Tolerant к медленной сети → идеален между нодами

3D Parallelism: комбинируем всё

Frontier модели используют все три вида параллелизма одновременно. Каждый решает свою задачу:

• TP (Tensor Parallelism): внутри одной ноды, NVLink. Режем слои. 2-8 GPU
• PP (Pipeline Parallelism): между нодами, сетевая карта. Режем модель по слоям. 2-64 stages
• FSDP/ZeRO (Data Parallelism): по всем GPU. Шардим optimizer states и параметры

Пример: LLaMA 3.1 405B обучалась на 16K GPU H100. TP=8 (внутри ноды, 8 GPU с NVLink) × PP=16 (16 pipeline stages между нодами) × FSDP (шардинг optimizer states). Итого: 8 × 16 = 128 GPU на одну "копию" модели, и ~125 таких копий для data parallelism.

Таблица: размер модели → стратегия

• ≤7B: 1-8 GPU. FSDP или ZeRO Stage 2. TP не нужен, PP не нужен
• 7B-13B: 4-16 GPU. ZeRO Stage 2-3 или FSDP. TP=2-4 опционально
• 30B-70B: 16-64 GPU. TP=4-8 (внутри ноды) + FSDP. PP если несколько нод
• 70B-200B: 64-256 GPU. TP=8 + PP=4-8 + FSDP. Обязателен multi-node
• 400B+: 1000+ GPU. TP=8 + PP=16+ + FSDP. Месяцы обучения, миллионы долларов

Фреймворки

• PyTorch FSDP: встроен в PyTorch. Data parallelism + sharding. Простой, хорош для ≤70B
• DeepSpeed (Microsoft): ZeRO-1/2/3, offload на CPU/NVMe. Богатая экосистема, интеграция с HuggingFace
• Megatron-LM (NVIDIA): TP + PP + DP. Максимальная производительность для pre-training
• NVIDIA NeMo: поверх Megatron-LM, более user-friendly. TP + PP + FSDP
• torchtitan (Meta): новый фреймворк, чистый PyTorch. TP + PP + FSDP2. Используется для LLaMA 3