Saltar al contenido principal

Capítulo 3: Técnicas de entrenamiento distribuido

Introducción​

Entrenar LLMs con miles de millones de parámetros requiere distribuir el trabajo entre muchas GPUs. Este capítulo cubre las principales estrategias de paralelismo.

Paralelismo de datos​

El enfoque más sencillo​

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

# Cada GPU procesa diferentes datos, mismo modelo
model = DistributedDataParallel(model)

Limitación: El modelo debe caber en una sola GPU

Paralelismo de modelo​

Paralelismo de tensores​

Divida capas individuales entre GPUs:

[Matriz de pesos de la capa]
   ↓ Dividir columnas
[GPU 0: W[:, :d/2]] [GPU 1: W[:, d/2:]]

Paralelismo de pipeline​

Capas diferentes en diferentes GPUs:

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

ZeRO: Optimizador de redundancia cero​

Tres etapas​

EtapaQué se particionaAhorro de memoria
ZeRO-1Estados del optimizador4x
ZeRO-2+ Gradientes8x
ZeRO-3+ ParámetrosLineal con GPUs

Implementación con 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ómputo de entrenamiento​

Escalado Chinchilla​

Entrenamiento óptimo: aproximadamente 20 tokens por parámetro

# Tamaño óptimo del modelo según el cómputo
def optimal_model_size(compute_budget):
    # C ≈ 6 * N * D (N=parámetros, D=tokens)
    # Óptimo: D ≈ 20N
    return compute_budget / (6 * 20)

Puntos clave​

Summary
  1. Paralelismo de datos para modelos pequeños
  2. Paralelismo de modelo cuando el modelo no cabe en una sola GPU
  3. ZeRO reduce drásticamente los requisitos de memoria
  4. Precisión mixta (bf16) acelera el entrenamiento