Pular para o conteúdo principal

Capítulo 3: Técnicas de Treinamento Distribuído

Introdução​

Treinar LLMs com bilhões de parâmetros exige distribuir o trabalho entre várias GPUs. Este capítulo aborda as principais estratégias de paralelismo.

Paralelismo de Dados​

A abordagem mais simples​

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel

# Cada GPU processa dados diferentes, mesmo modelo
model = DistributedDataParallel(model)

Limitação: O modelo deve caber em uma única GPU

Paralelismo de Modelo​

Paralelismo de Tensor​

Divide camadas individuais entre GPUs:

[Matriz de Pesos da Camada]
   ↓ Divide colunas
[GPU 0: W[:, :d/2]] [GPU 1: W[:, d/2:]]

Paralelismo de Pipeline​

Camadas diferentes em GPUs diferentes:

GPU 0: Camadas 0-11
GPU 1: Camadas 12-23
GPU 2: Camadas 24-35
GPU 3: Camadas 36-47

ZeRO: Zero Redundancy Optimizer​

Três Estágios​

EstágioO que é particionadoEconomia de memória
ZeRO-1Estados do otimizador4x
ZeRO-2+ Gradientes8x
ZeRO-3+ ParâmetrosLinear com GPUs

Implementação com DeepSpeed​

import deepspeed

model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    config={
        "zero_optimization": {
            "stage": 3,
            "offload_optimizer": {"device": "cpu"},
            "offload_param": {"device": "cpu"}
        },
        "bf16": {"enabled": True}
    }
)

Cálculo de Treinamento​

Chinchilla Scaling​

Treinamento ótimo: aproximadamente 20 tokens por parâmetro

# Tamanho ótimo do modelo para o orçamento de computação
def optimal_model_size(compute_budget):
    # C ≈ 6 * N * D (N=parâmetros, D=tokens)
    # Ótimo: D ≈ 20N
    return compute_budget / (6 * 20)

Principais Pontos​

Summary
  1. Paralelismo de dados para modelos pequenos
  2. Paralelismo de modelo quando o modelo não cabe em uma GPU
  3. ZeRO reduz drasticamente os requisitos de memória
  4. Precisão mista (bf16) acelera o treinamento