Налаштування розподіленого навчання (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.