Distributed Training

O treino distribuído reparte o trabalho de treinar um modelo por vários processadores, normalmente GPUs ou TPUs numa máquina, ou num cluster na nuvem ou local, para concluir o treino mais depressa ou permitir treinar um modelo cuja dimensão excede a memória de um único dispositivo. Os dispositivos trocam gradientes, ativações ou fragmentos de parâmetros através de interligações rápidas, usando bibliotecas de comunicação como NVIDIA NCCL. É assim que se treinam os modernos modelos de base e grandes modelos de linguagem, e também permite acelerar o treino quotidiano de deteção de objetos em estações de trabalho com várias GPUs.

Como funciona o treino distribuído#

O paralelismo de dados é a estratégia mais comum. Cada GPU mantém uma réplica completa do modelo, processa a sua própria parte de cada lote de dados de treino e participa numa operação all-reduce que calcula a média dos gradientes, para que cada réplica aplique a mesma atualização de descida do gradiente. O DistributedDataParallel (DDP) do PyTorch sobrepõe essa comunicação à passagem para trás. O tamanho global do lote corresponde ao lote por GPU multiplicado pelo número de GPUs, pelo que normalmente é necessário escalar a taxa de aprendizagem de acordo com esse valor: Goyal et al. treinaram o ResNet-50 com ImageNet em uma hora, usando 256 GPUs e um lote de 8,192 imagens, com uma regra de escalonamento linear e uma fase de aquecimento, alcançando cerca de 90% de eficiência de escalonamento entre 8 e 256 GPUs.

O paralelismo de dados com fragmentação reduz o custo de memória dessas réplicas. O ZeRO e o Fully Sharded Data Parallel (FSDP) do PyTorch dividem os parâmetros, os gradientes e os estados do otimizador entre GPUs e reúnem os pesos completos de cada camada apenas durante o respetivo cálculo, reduzindo a memória necessária por GPU à medida que se acrescentam GPUs.

O paralelismo de modelos divide a própria rede neural. O paralelismo de tensores divide multiplicações de matrizes individuais entre GPUs; o paralelismo de pipeline atribui camadas consecutivas a GPUs diferentes e encaminha micro-lotes por elas; o paralelismo de contexto divide sequências longas; e o paralelismo de especialistas coloca os especialistas de um modelo de mistura de especialistas em GPUs diferentes. Estruturas de trabalho como o Megatron-LM combinam estas técnicas com o paralelismo de dados para treinar modelos com centenas de milhares de milhões de parâmetros.

Treino com várias GPUs na prática#

A comunicação é o principal custo adicional: o tráfego all-reduce aumenta com o tamanho do modelo, por isso GPUs conectadas por NVLink dentro de um nó escalam melhor do que nós conectados por redes mais lentas. As estatísticas de normalização em lote são calculadas por GPU, a menos que as camadas sejam convertidas para normalização em lote sincronizada, o que importa quando a parte do lote de cada GPU é pequena. O Ultralytics YOLO26 executa o treinamento com várias GPUs por meio de DDP: passar vários IDs de dispositivos inicia um processo por GPU com torch.distributed.run, usando o backend NCCL no CUDA, e divide igualmente entre as GPUs o valor de batch.

from ultralytics import YOLO

# Carrega um modelo YOLO26 pré-treinado
model = YOLO("yolo26n.pt")

# Treina nas GPUs 0 e 1 com DDP; batch=32 dá 16 imagens por passo a cada GPU
results = model.train(data="coco8.yaml", epochs=100, imgsz=640, batch=32, device=[0, 1])

Aplicações no mundo real#

• Grandes modelos de linguagem: A Meta treinou o Llama 3 405B com até 16K GPUs H100, combinando paralelismo de tensores, pipeline, contexto e dados no que os autores chamam de paralelismo 4D.
• Veículos autónomos: As equipas de perceção treinam modelos de deteção e segmentação com grandes arquivos de gravações de câmaras, lidar e radar de frotas de veículos, algo que só é viável em grandes clusters de GPUs.
• Análise de imagens médicas: Os volumes de TC e RM 3D de alta resolução exigem demasiada memória para uma única GPU, e distribuir os lotes ou os fragmentos do modelo por várias GPUs permite treinar volumes e modelos maiores.

Treino distribuído vs. conceitos relacionados#

• Aprendizagem federada: A aprendizagem federada também treina em vários dispositivos, mas os dados permanecem nos telemóveis, hospitais ou outros clientes por motivos de privacidade, e apenas as atualizações do modelo são enviadas para um servidor central. O treino distribuído reúne os dados num único cluster para maximizar a velocidade.
• Aprendizagem dividida: A aprendizagem dividida separa um modelo entre um cliente e um servidor para que os dados em bruto nunca saiam do cliente, trocando privacidade por comunicação adicional.
• Acumulação de gradientes: A acumulação de gradientes soma os gradientes de vários lotes pequenos antes de uma atualização, reproduzindo a matemática de um lote grande numa única GPU, à custa de tempo; o paralelismo de dados, por sua vez, calcula esses lotes simultaneamente em várias GPUs.