Capítulo 2: Entrenamiento Distribuido

Estrategias de Paralelismo

Ejemplo de DDP en 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

Puntos Clave

Summary

1. El paralelismo de datos funciona para la mayoría de los casos de uso
2. El paralelismo de modelo es necesario para modelos muy grandes
3. DeepSpeed/FSDP permiten entrenar modelos más grandes que la memoria de la GPU

---

Siguiente: Capítulo 3: Entrenamiento en la Nube