Как работает распределенное обучение?

Распределённое обучение (Distributed Training)

Распределённое обучение — это техника параллельной обработки данных на нескольких GPU/TPU/машинах для ускорения тренировки глубоких нейросетей.

Основные парадигмы

1. Data Parallelism (параллелизм данных)

Одна модель копируется на несколько устройств, каждое обрабатывает свой батч данных:

import torch
import torch.nn as nn
from torch.nn import DataParallel

# Простой способ на одной машине с несколькими GPU
model = YourModel()
model = DataParallel(model)  # Автоматически распределяет на все GPU

# Или более явно
model = torch.nn.parallel.DataParallel(
    model, device_ids=[0, 1, 2, 3]  # GPU индексы
)

output = model(input_data)  # Автоматически распределяет батч

Процесс:

1. Forward pass на разных GPU параллельно
2. Вычисление loss на каждом GPU
3. Backward pass параллельно
4. Градиенты собираются на master GPU
5. Обновление весов (синхронизирован)
6. Обновленные веса рассылаются на остальные GPU

2. Model Parallelism (параллелизм моделей)

Модель разделяется между несколькими устройствами (для очень больших моделей):

# Разделение слоёв моделей
class LargeModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1000, 500).cuda(0)
        self.fc2 = nn.Linear(500, 100).cuda(1)
        self.fc3 = nn.Linear(100, 10).cuda(2)
    
    def forward(self, x):
        x = x.cuda(0)
        x = self.fc1(x)
        x = x.cuda(1)
        x = self.fc2(x)
        x = x.cuda(2)
        x = self.fc3(x)
        return x

Распределённое обучение на нескольких машинах

DistributedDataParallel (DDP) — PyTorch

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

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group('nccl', rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

def train(rank, world_size):
    setup(rank, world_size)
    
    # Создаём модель
    model = YourModel().to(rank)
    model = DDP(model, device_ids=[rank])
    
    # DistributedSampler для правильного распределения данных
    train_sampler = DistributedSampler(
        dataset, num_replicas=world_size, rank=rank, shuffle=True
    )
    train_loader = DataLoader(dataset, sampler=train_sampler, batch_size=32)
    
    optimizer = torch.optim.Adam(model.parameters())
    
    for epoch in range(num_epochs):
        train_sampler.set_epoch(epoch)  # Шафлирование каждую эпоху
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(rank), target.to(rank)
            
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()  # Градиенты автоматически синхронизируются
    
    cleanup()

if __name__ == "__main__":
    world_size = 4  # 4 GPU
    mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)

TensorFlow Distributed Training

import tensorflow as tf

# Простой способ с несколькими GPU на одной машине
strategy = tf.distribute.MirroredStrategy()  # Использует все GPU

with strategy.scope():
    model = tf.keras.Sequential([...])
    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')

model.fit(x_train, y_train, epochs=10, batch_size=32)

# Для нескольких машин
strategy = tf.distribute.MultiWorkerMirroredStrategy()

Оптимизация общих параметров

Batch Size

# При распределении на K GPU, глобальный batch size увеличивается
local_batch_size = 32
num_gpus = 4
global_batch_size = local_batch_size * num_gpus  # 128

# Требует корректировки learning rate
lr = 0.001 * num_gpus  # Линейное масштабирование

Синхронизация градиентов

Время обучения одной эпохи (с коммуникацией):
- Forward: 0.5с на GPU
- Backward: 0.5с на GPU
- All-reduce градиентов: зависит от сети (0.1-1.0с)

С 4 GPU:
- Последовательно: 4 * (0.5 + 0.5) = 4с
- Параллельно: 0.5 + 0.5 + 0.1 = 1.1с (speedup ≈ 3.6x)

Практические советы

# 1. Используй TensorFlow/PyTorch встроенные инструменты
# Не пиши All-reduce вручную

# 2. Профилируй (найди узкие места)
from torch.profiler import profile
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    for data in train_loader:
        train_step(data)
print(prof.key_averages().table(sort_by='cuda_time_total'))

# 3. Монитори пропускную способность
# Каждый GPU должен быть >= 80% загруженным

# 4. Используй Apex (Automatic Mixed Precision)
from apex import amp
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')

# 5. Уменьшай размер моделей через прунинг/квантизацию

Когда использовать какую стратегию

Типичные ошибки

1. Забыли обернуть модель в DDP перед обучением
2. Не использовали DistributedSampler (некоторые образцы пропускают)
3. Не синхронизировали random seeds
4. Не масштабировали learning rate с batch size
5. Забыли вызвать cleanup() для освобождения ресурсов

Распределённое обучение критично для обучения моделей на миллиардах параметров и петабайтах данных.