Capítulo 2: Treinamento Distribuído

Estratégias de Paralelismo

Exemplo de DDP com 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("nccl", rank=rank, world_size=world_size)

def train(rank, world_size):
    setup(rank, world_size)
    
    model = MyModel().to(rank)
    model = DDP(model, device_ids=[rank])
    
    optimizer = torch.optim.Adam(model.parameters())
    
    for epoch in range(num_epochs):
        for batch in dataloader:
            outputs = model(batch)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
    
    dist.destroy_process_group()

# Launch with torchrun
# torchrun --nproc_per_node=4 train.py

Principais Pontos

Summary

1. Paralelismo de dados funciona para a maioria dos casos de uso
2. Paralelismo de modelo é necessário para modelos muito grandes
3. DeepSpeed/FSDP permitem treinar modelos maiores que a memória da GPU

---

Próximo: Capítulo 3: Treinamento em Nuvem