Настройка распределённого обучения (Distributed Training) моделей

Настройка распределённого обучения (Distributed Training) моделей

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

Стратегии параллелизма

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

Model Parallelism (Tensor Parallelism) — модель разбивается по слоям или тензорам между GPU. Необходим, когда модель слишком большая для одного GPU. Используется в Megatron-LM, DeepSpeed.

Pipeline Parallelism — слои модели распределяются по GPU последовательно. Разные GPU обрабатывают разные micro-batches одновременно. Используется в GPipe, PipeDream.

3D Parallelism — комбинация всех трёх стратегий. Используется DeepSpeed и Megatron-LM для обучения LLM с сотнями миллиардов параметров.

Data Parallel с PyTorch DDP

DistributedDataParallel (DDP) — рекомендуемый подход для data parallelism в PyTorch:

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup(rank, world_size):
    dist.init_process_group(
        backend='nccl',  # nccl для GPU, gloo для CPU
        rank=rank,
        world_size=world_size
    )
    torch.cuda.set_device(rank)

def train(rank, world_size, model, dataset):
    setup(rank, world_size)

    model = model.to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    loader = DataLoader(dataset, sampler=sampler, batch_size=32)

    optimizer = torch.optim.AdamW(ddp_model.parameters(), lr=1e-4)

    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)  # Важно для перемешивания
        for batch in loader:
            optimizer.zero_grad()
            loss = ddp_model(batch)
            loss.backward()  # all-reduce автоматически
            optimizer.step()

Запуск на одном узле (8 GPU):

torchrun --nproc_per_node=8 train.py

Запуск на нескольких узлах:

# На узле 0 (master):
torchrun --nnodes=4 --nproc_per_node=8 \
  --node_rank=0 \
  --master_addr="10.0.0.1" --master_port=29500 \
  train.py

# На узлах 1-3 (worker):
torchrun --nnodes=4 --nproc_per_node=8 \
  --node_rank=1 \  # 2, 3 соответственно
  --master_addr="10.0.0.1" --master_port=29500 \
  train.py

Accelerate от Hugging Face

Для более простой настройки с поддержкой mixed precision, gradient accumulation и различных distributed backends:

from accelerate import Accelerator

accelerator = Accelerate(
    mixed_precision='bf16',
    gradient_accumulation_steps=4
)

model, optimizer, train_dataloader = accelerator.prepare(
    model, optimizer, train_dataloader
)

for batch in train_dataloader:
    with accelerator.accumulate(model):
        outputs = model(**batch)
        loss = outputs.loss
        accelerator.backward(loss)
        optimizer.step()
        optimizer.zero_grad()

Автоматический выбор стратегии

Scaling efficiency и оптимизация

Ключевые метрики: GPU utilization (цель > 85%), MFU (Model FLOPS Utilization). Частые bottleneck:

• IO-bound: чтение данных медленнее, чем GPU обрабатывает. Решение: prefetch, увеличение num_workers, NVMe хранилище.
• Communication-bound: all-reduce занимает слишком много времени. Решение: gradient compression, увеличение batch size.
• Memory-bound: gradient checkpointing, mixed precision (BF16/FP16), activation offloading.

На кластере 8x A100 80GB с NVLink при обучении модели 7B параметров (DDP + ZeRO-2) достигается MFU около 40-50% — типичный показатель для well-tuned setup.