PyTorch Entrenamiento Distribuido

Cuando los modelos y los datos crecen, una sola GPU ya no puede satisfacer las necesidades de entrenamiento.

El entrenamiento distribuido utiliza múltiples GPU o incluso múltiples máquinas para el cómputo en paralelo, lo que puede reducir significativamente el tiempo de entrenamiento.

Esta sección detalla las técnicas de entrenamiento distribuido en PyTorch, incluyendo DataParallel, DistributedDataParallel y entrenamiento distribuido de precisión mixta.


1. Fundamentos del Entrenamiento Distribuido

1.1 ¿Por qué necesitamos entrenamiento distribuido?

La escala de los modelos de aprendizaje profundo crece año tras año, y el tiempo de entrenamiento también aumenta. El objetivo principal del entrenamiento distribuido es:Distribuir las tareas de cómputo entre múltiples unidades de cómputo para completar el entrenamiento de modelos y datos a gran escala en un tiempo aceptable.。

El entrenamiento distribuido resuelve principalmente dos problemas:

  • Memoria insuficiente: Los parámetros de los modelos grandes, el estado del optimizador y los gradientes requieren mucha memoria de video (VRAM).
  • Tiempo de entrenamiento demasiado largo: Entrenar un modelo grande con una sola GPU puede llevar semanas o incluso meses.

1.2 Los dos modos principales de entrenamiento distribuido

El entrenamiento distribuido se divide principalmente en dos modos:

Modo Principio Escenario de aplicación
Paralelismo de datos (Data Parallel) Cada unidad de cómputo guarda el modelo completo, y diferentes unidades procesan diferentes datos. El modelo cabe en la memoria de una sola GPU; grandes volúmenes de datos requieren acelerar el entrenamiento; equipos pequeños con recursos limitados.
Paralelismo de modelos (Model Parallel) Dividir el modelo entre múltiples unidades de cómputo, cada una guardando solo una parte del modelo. Modelos extremadamente grandes que no caben en una sola GPU; clústeres dedicados, entrenamiento en múltiples máquinas.

En entornos de producción, el paralelismo de datos es el enfoque más común. PyTorch proporciona dos implementaciones: DataParallel y DistributedDataParallel.


2. Guía de uso de DataParallel

2.1 Múltiples GPU en una sola máquina: DataParallel

DataParallel(DP)Es la forma más sencilla de entrenamiento con múltiples GPU que ofrece PyTorch. Solo necesita modificar unas pocas líneas de código para aprovechar varias GPU en una sola máquina para entrenar.

Ejemplo

import torch
import torch.nn as nn
import torch.optim as optim

# ── Definición del modelo ──────────────────────────────────────
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, 10)
        )

    def forward(self, x):
        return self.net(x)


# Comprueba el número de GPU disponibles
print(f"Número de GPU disponibles: {torch.cuda.device_count()}")

# Crea el modelo y muévelo a la GPU
model = SimpleModel()

# ── Método 1: usar DataParallel (la forma más simple)───────
if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)

# Muévelo a la GPU (si se usa DataParallel, device_ids se maneja automáticamente)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# Función de pérdida y optimizador
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# ── Bucle de entrenamiento ──────────────────────────────────────
def train_epoch_dp(model, loader, criterion, optimizer, device):
    model.train()
    total_loss = 0
    correct = 0
    total = 0

    for inputs, labels in loader:
        inputs = inputs.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        total_loss += loss.item() * inputs.size(0)
        _, predicted = outputs.max(1)
        correct += predicted.eq(labels).sum().item()
        total += labels.size(0)

    return total_loss / total, correct / total


# Datos simulados
train_loader = [
    (torch.randn(32, 128), torch.randint(0, 10, (32,))) for _ in range(10)
]

# Iniciar el entrenamiento
for epoch in range(3):
    loss, acc = train_epoch_dp(model, train_loader, criterion, optimizer, device)
    print(f"Epoch {epoch+1}: Loss={loss:.4f}, Acc={acc:.4f}")

print("¡Entrenamiento con DataParallel completado!")

2.2 Cómo funciona DataParallel

El flujo de trabajo de DataParallel es el siguiente:

  1. Copia del modelo: Copiar el modelo a cada GPU.
  2. Distribución de datos: Distribuir uniformemente los datos del batch a cada GPU.
  3. Cálculo paralelo: Cada GPU calcula la propagación hacia adelante de forma independiente.
  4. Agregación de gradientes: Reunir los gradientes de todas las GPU en la GPU principal.
  5. Actualización de parámetros: La GPU principal actualiza los parámetros y los sincroniza con las otras GPU.

Ejemplo

# Descripción de los parámetros clave de DataParallel
model = nn.DataParallel(
    module,                    # Modelo a paralelizar (obligatorio)
    device_ids=[0, 1, 2, 3],   # IDs de los dispositivos GPU a usar; por defecto usa todas las tarjetas
    output_device=0,            # Dispositivo al que se envían los resultados; por defecto es la primera tarjeta
    dim=0                      # Dimensión para dividir los datos; por defecto es la dimensión del batch
)

# Ver el dispositivo donde está el modelo actual
print(f"Dispositivo del modelo: {next(model.parameters()).device}")

# Ver los dispositivos realmente utilizados
print(f"Número de GPU utilizadas: {model.device_ids if hasattr(model, 'device_ids') else 'N/A'}")

DataParallel es simple y fácil de usar, pero tiene dos problemas principales: 1) La GPU principal tiene mucha presión de memoria (necesita agregar gradientes); 2) La eficiencia de comunicación entre GPUs es baja. Por lo tanto, se recomienda oficialmente usar DistributedDataParallel.


3. Explicación detallada de DistributedDataParallel

3.1 Fundamentos de DDP

DistributedDataParallel(DDP)Es el método de entrenamiento distribuido recomendado por PyTorch. En comparación con DataParallel, DDP tiene las siguientes ventajas:

  • Cada GPU calcula los gradientes de forma independiente, sin necesidad de agregarlos en la GPU principal.
  • Utiliza un algoritmo eficiente de sincronización de gradientes (Ring AllReduce).
  • Admite entrenamiento con múltiples máquinas y múltiples GPU.
  • Entrenamiento más rápido y uso más eficiente de la memoria de video.

3.2 Entrenamiento con DDP en una sola máquina con múltiples GPU

Ejemplo

import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import os

# ── Inicializar el entorno distribuido ────────────────────────────
def setup(rank, world_size):
    """Configurar el entorno distribuido"""
    # Establecer la GPU que usa el proceso actual
    torch.cuda.set_device(rank)

    # Inicializar el grupo de procesos
    dist.init_process_group(
        backend="nccl",           # Usar el backend NCCL (recomendado para GPU)
        init_method="env://",     # Método de inicialización mediante variables de entorno
        world_size=world_size,    # Número total de procesos
        rank=rank                 # Rango del proceso actual
    )


def cleanup():
    """Limpiar el entorno distribuido"""
    dist.destroy_process_group()


# ── Cargador de datos (compatible con distribución)───────────────────────
def get_distributed_loader(batch_size, world_size):
    """Crear un cargador de datos compatible con distribución"""
    from torch.utils.data import DataLoader, DistributedSampler

    # Conjunto de datos simulado
    dataset = torch.utils.data.TensorDataset(
        torch.randn(100, 3, 32, 32),
        torch.randint(0, 10, (100,))
    )

    # DistributedSampler dividirá los datos automáticamente
    sampler = DistributedSampler(
        dataset,
        num_replicas=world_size,
        rank=rank,
        shuffle=True
    )

    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        sampler=sampler,
        num_workers=2,
        pin_memory=True
    )

    return loader


# ── Definición del modelo ──────────────────────────────────────
class ImageClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten()
        )
        self.classifier = nn.Linear(64, 10)

    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x


# ── Función de entrenamiento con DDP ──────────────────────────────────
def train_ddp(rank, world_size, epochs=3):
    """Función principal de entrenamiento distribuido"""
    # Inicialización
    setup(rank, world_size)

    # Crear el modelo y moverlo a la GPU correspondiente
    model = ImageClassifier().cuda(rank)

    # Envolver como modelo DDP
    ddp_model = DDP(
        model,
        device_ids=[rank],        # Especificar el dispositivo que usa el proceso actual
        output_device=rank,
        find_unused_parameters=False  # Si se detectan parámetros no utilizados
    )

    # Función de pérdida y optimizador
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(ddp_model.parameters(), lr=1e-3)

    # Obtener el cargador de datos
    loader = get_distributed_loader(batch_size=16, world_size=world_size)

    # Bucle de entrenamiento
    for epoch in range(epochs):
        # Al comienzo de cada epoch, establece el número de epoch (para barajar los datos)
        loader.sampler.set_epoch(epoch)

        for batch_idx, (inputs, labels) in enumerate(loader):
            inputs = inputs.cuda(rank, non_blocking=True)
            labels = labels.cuda(rank, non_blocking=True)

            optimizer.zero_grad()
            outputs = ddp_model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            if rank == 0 and batch_idx % 5 == 0:
                print(f"Epoch {epoch+1}, Batch {batch_idx}, Loss: {loss.item():.4f}")

    # Limpieza
    if rank == 0:
        print("¡Entrenamiento completado!")
    cleanup()


# Nota: en la ejecución real, es necesario especificar rank y world_size mediante parámetros de línea de comandos
# Este código debe ejecutarse de forma independiente en cada proceso

3.3 Iniciar entrenamiento distribuido

El entrenamiento distribuido de PyTorch debe iniciarse mediantetorchrunotorch.distributed.launchiniciar. Las siguientes son las formas de inicio más comunes:

Ejemplo

# Método 1: usar torchrun (recomendado, PyTorch 2.0+)
# Archivo: train_ddp.py

"""
Modo de uso:
torchrun --nproc_per_node=4 train_ddp.py
"
""

# Método 2: usar python -m (método tradicional)
"""
python -m torch.distributed.launch --nproc_per_node=4 train_ddp.py
"
""

# Método 3: inicio con múltiples máquinas y múltiples GPU
# Máquina 1 (nodo maestro)
torchrun --nnodes=2 --nproc_per_node=4 --node_rank=0 --master_addr=192.168.1.1 --master_port=29500 train_ddp.py

# Máquina 2 (nodo esclavo)
torchrun --nnodes=2 --nproc_per_node=4 --node_rank=1 --master_addr=192.168.1.1 --master_port=29500 train_ddp.py

# Descripción de parámetros:
# --nproc_per_node: número de GPU por nodo
# --nnodes: número de nodos
# --node_rank: número del nodo actual (comienza desde 0)
# --master_addr: dirección IP del nodo maestro
# --master_port: número de puerto del nodo maestro

3.4 Script completo de entrenamiento con DDP

Ejemplo

# Script completo de entrenamiento DDP (guardar como train_ddp.py)
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
import argparse
import os


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--local_rank", type=int, default=-1, help="Pasado automáticamente por torchrun")
    parser.add_argument("--epochs", type=int, default=10)
    parser.add_argument("--batch_size", type=int, default=32)
    parser.add_argument("--lr", type=float, default=1e-3)
    return parser.parse_args()


class TrainModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(128, 10)
        )

    def forward(self, x):
        return self.net(x)


def setup(rank, world_size):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355"
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)


def cleanup():
    dist.destroy_process_group()


def train(rank, world_size, args):
    setup(rank, world_size)

    # Crear el modelo
    model = TrainModel().cuda(rank)
    model = DDP(model, device_ids=[rank])

    # Optimizador y función de pérdida
    optimizer = optim.Adam(model.parameters(), lr=args.lr)
    criterion = nn.CrossEntropyLoss()

    # Conjunto de datos
    dataset = torch.utils.data.TensorDataset(
        torch.randn(1000, 3, 32, 32),
        torch.randint(0, 10, (1000,))
    )

    # DistributedSampler asegura que cada proceso vea datos diferentes
    sampler = DistributedSampler(
        dataset,
        num_replicas=world_size,
        rank=rank,
        shuffle=True
    )

    loader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        sampler=sampler,
        num_workers=2,
        pin_memory=True
    )

    # Bucle de entrenamiento
    for epoch in range(args.epochs):
        # Clave: en cada epoch se debe establecer el epoch para barajar la partición de datos
        sampler.set_epoch(epoch)
        model.train()

        epoch_loss = 0
        for inputs, labels in loader:
            inputs = inputs.cuda(rank, non_blocking=True)
            labels = labels.cuda(rank, non_blocking=True)

            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            epoch_loss += loss.item()

        # Sincronizar todos los procesos, asegurando que cada proceso complete el epoch actual
        dist.barrier()

        if rank == 0:
            avg_loss = epoch_loss / len(loader)
            print(f"Epoch {epoch+1}/{args.epochs}, Loss: {avg_loss:.4f}")

    cleanup()


if __name__ == "__main__":
    args = parse_args()

    # Obtener world_size de la variable de entorno
    world_size = int(os.environ["WORLD_SIZE"])

    # Obtener rank de la variable de entorno (establecido automáticamente por torchrun)
    rank = int(os.environ["RANK"])

    if torch.cuda.is_available():
        train(rank, world_size, args)
    else:
        print("Se necesita un entorno CUDA para ejecutar el entrenamiento distribuido")

Punto clave: en cada epoch es necesario llamar asampler.set_epoch(epoch), asegurando que la partición de datos se baraje; de lo contrario, múltiples procesos verán los mismos datos.


4. Precisión mixta en el entrenamiento distribuido

4.1 Combinación de DDP + AMP

La combinación del entrenamiento distribuido con el entrenamiento de precisión mixta puede mejorar aún más la velocidad de entrenamiento. El DDP de PyTorch 2.0+ ya está perfectamente integrado con AMP.

Ejemplo

import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.cuda.amp import autocast, GradScaler


class AMPModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(128, 10)
        )

    def forward(self, x):
        return self.net(x)


def train_ddp_amp(rank, world_size):
    """Entrenamiento distribuido + precisión mixta"""
    # Inicialización
    torch.cuda.set_device(rank)
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

    # Crear el modelo
    model = AMPModel().cuda(rank)
    model = DDP(model, device_ids=[rank])

    # Optimizador y GradScaler
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    scaler = GradScaler()

    criterion = nn.CrossEntropyLoss()

    # Bucle de entrenamiento
    model.train()
    for inputs, labels in loader:
        inputs = inputs.cuda(rank, non_blocking=True)
        labels = labels.cuda(rank, non_blocking=True)

        optimizer.zero_grad()

        # Usar autocast para cambiar automáticamente la precisión
        with autocast(device_type='cuda'):
            outputs = model(inputs)
            loss = criterion(outputs, labels)

        # Usar scaler para manejar el escalado de gradientes
        scaler.scale(loss).backward()

        # Recorte de gradientes (requiere unscale primero)
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

        scaler.step(optimizer)
        scaler.update()

    dist.destroy_process_group()


# Comando de inicio
# torchrun --nproc_per_node=4 train_ddp_amp.py

4.2 Sincronización de gradientes en el entrenamiento distribuido

DDP utilizaRing AllReducealgoritmo para la sincronización de gradientes, cuyo principio de funcionamiento es el siguiente:

  1. Cada GPU divide los gradientes en n partes (n es el número de GPU)
  2. Cada GPU intercambia una parte de gradientes con la GPU adyacente
  3. Después de repetir n veces, cada GPU tiene los gradientes completos

Ejemplo

# Configuración relacionada con la sincronización de gradientes de DDP

# 1. Establecer el tamaño del búfer del bucket de gradientes
model = DDP(
    model,
    device_ids=[rank],
    gradient_as_bucket_view=True,  # Ahorrar memoria de GPU (PyTorch 1.8+)
    broadcast_buffers=False,        # No transmitir buffers, reducir la comunicación
    bucket_cap_mb=25               # Tamaño del bucket de gradientes (MB)
)

# 2. Sincronizar manualmente los parámetros del modelo (si es necesario)
# En algunos escenarios, es necesario sincronizar manualmente los parámetros del modelo
def sync_params(model):
    for param in model.parameters():
        dist.broadcast(param, src=0)

# 3. Verificar el estado de sincronización de gradientes
# DDP maneja automáticamente la sincronización de gradientes, no se requiere intervención manual
# Pero se puede verificar de la siguiente manera
for name, param in model.named_parameters():
    if param.grad is not None:
        print(f"{name}: grad requires_grad={param.grad.requires_grad}")

5. Paralelismo de modelos y canalizaciones (Pipeline)

5.1 Paralelismo de modelos simple

Cuando el modelo es demasiado grande y una sola GPU no puede contenerlo, es necesario dividir el modelo en múltiples GPU.

Ejemplo

# Ejemplo de paralelismo de modelos: colocar diferentes capas del modelo en diferentes GPU

class ModelParallel(nn.Module):
    def __init__(self, num_gpus=2):
        super().__init__()
        self.num_gpus = num_gpus

        # Primera parte: GPU 0
        self.features1 = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        ).to(f"cuda:0")

        # Segunda parte: GPU 1
        self.features2 = nn.Sequential(
            nn.Conv2d(64, 128, 3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten()
        ).to(f"cuda:1")

        # Clasificador: GPU 0
        self.classifier = nn.Linear(128, 10).to(f"cuda:0")

    def forward(self, x):
        # Los datos se transfieren entre diferentes GPU
        x = x.to("cuda:0")
        x = self.features1(x)
        x = x.to("cuda:1")
        x = self.features2(x)
        x = x.to("cuda:0")
        x = self.classifier(x)
        return x


# Usar
model = ModelParallel(num_gpus=2)
output = model(torch.randn(1, 3, 32, 32))
print(f"Dispositivo de salida: {output.device}")

5.2 Paralelismo de canalizaciones (Pipeline Parallelism)

PyTorch proporcionatorch.distributed.pipeline.syncmódulo (paquete PP) para implementar el paralelismo de pipeline, que divide el modelo por capas formando un pipeline de cálculo.

Ejemplo

# Ejemplo de paralelismo de pipeline (requiere instalar la dependencia torch-pipeline)
# pip install torch-pipeline

"""
# Nota: el pipeline nativo de PyTorch requiere una versión específica o una biblioteca de terceros
# Aquí se muestra el concepto, con una implementación simplificada

from torch.distributed.pipeline.sync import Pipe

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Linear(128, 256)
        self.layer2 = nn.Linear(256, 256)
        self.layer3 = nn.Linear(256, 10)

    def forward(self, x):
        x = torch.relu(self.layer1(x))
        x = torch.relu(self.layer2(x))
        x = self.layer3(x)
        return x


# Dividir el modelo por capas en diferentes GPU
# layer1 -> GPU 0
# layer2 -> GPU 1
# layer3 -> GPU 2

model = SimpleModel()
model.layer1 = model.layer1.cuda(0)
model.layer2 = model.layer2.cuda(1)
model.layer3 = model.layer3.cuda(2)

# Usar Pipeline para conectar el modelo
# Nota: el uso real requiere instalar la versión correspondiente de la biblioteca pipeline
# from torch.distributed.pipeline.sync import Pipe
# model = Pipe(model, chunks=4)
"""

print("El paralelismo de pipeline requiere instalar adicionalmente la biblioteca torch-pipeline")
print("pip install torch-pipeline-sync")

6. Mejores prácticas para el entrenamiento distribuido

6.1 Consejos de optimización del rendimiento

Elemento de optimización Técnica Descripción
Optimización de comunicación Usar backend NCCL NCCL es la mejor opción para GPU; es más rápido que Gloo
Carga de datos Establecer num_workers y pin_memory num_workers debe establecerse al número de núcleos de CPU; pin_memory acelera la transferencia de CPU a GPU
Sincronización de gradientes Ajustar bucket_cap_mb Para modelos grandes se puede aumentar el bucket; para modelos pequeños se puede reducir para mejorar la respuesta
Precisión mixta DDP + AMP Reduce la cantidad de comunicación; aumenta la velocidad de entrenamiento

6.2 Problemas comunes y soluciones

Problema Causa Solución
Comunicación lenta Ancho de banda de red insuficiente Usar InfiniBand o redes de alta velocidad
Memoria de GPU insuficiente batch size demasiado grande Reducir el batch size o usar acumulación de gradientes
Carga desequilibrada División del modelo no razonable Ajustar la distribución de capas para que sea lo más uniforme posible
Error de sincronización Gradientes NaN Revisar la configuración de recorte y escalado de gradientes

6.3 Script completo para entrenamiento con múltiples máquinas y GPU

Ejemplo

# Plantilla de script completo de entrenamiento con múltiples máquinas y múltiples GPU

import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
import argparse
import os


class TrainingConfig:
    def __init__(self):
        self.local_rank = -1
        self.epochs = 30
        self.batch_size = 32
        self.lr = 1e-3
        self.num_workers = 4
        self.gradient_clip = 1.0
        self.use_amp = True


def init_distributed(config):
    """Inicializar el entorno distribuido"""
    # Obtener la configuración de las variables de entorno
    local_rank = config.local_rank

    # Configuración de NCCL
    os.environ["NCCL_DEBUG"] = "WARN"

    torch.cuda.set_device(local_rank)
    dist.init_process_group(backend="nccl")

    return local_rank


def create_dataloader(dataset, config, world_size, rank):
    """Crear el cargador de datos"""
    sampler = DistributedSampler(
        dataset,
        num_replicas=world_size,
        rank=rank,
        shuffle=True,
        drop_last=True
    )

    loader = DataLoader(
        dataset,
        batch_size=config.batch_size,
        sampler=sampler,
        num_workers=config.num_workers,
        pin_memory=True,
        persistent_workers=True
    )

    return loader, sampler


def train_epoch(model, loader, criterion, optimizer, scaler, config, rank):
    """Entrenar un epoch"""
    model.train()
    total_loss = 0
    correct = 0
    total = 0

    for inputs, labels in loader:
        inputs = inputs.cuda(rank, non_blocking=True)
        labels = labels.cuda(rank, non_blocking=True)

        optimizer.zero_grad()

        if config.use_amp and scaler is not None:
            # Entrenamiento de precisión mixta
            with torch.cuda.amp.autocast():
                outputs = model(inputs)
                loss = criterion(outputs, labels)

            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), config.gradient_clip)
            scaler.step(optimizer)
            scaler.update()
        else:
            # Entrenamiento normal
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), config.gradient_clip)
            optimizer.step()

        total_loss += loss.item() * inputs.size(0)
        _, predicted = outputs.max(1)
        correct += predicted.eq(labels).sum().item()
        total += labels.size(0)

    return total_loss / total, correct / total


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--local_rank", type=int, default=-1)
    args = parser.parse_args()

    config = TrainingConfig()
    config.local_rank = args.local_rank

    rank = init_distributed(config)
    world_size = dist.get_world_size()

    print(f"Proceso {rank}/{world_size} inicializado")

    # Crear el modelo
    model = YourModel().cuda(rank)
    model = DDP(model, device_ids=[rank])

    # Optimizador
    optimizer = optim.Adam(model.parameters(), lr=config.lr)

    # GradScaler (si se usa precisión mixta)
    scaler = GradScaler() if config.use_amp else None

    # Conjunto de datos (reemplázalo con tu conjunto de datos)
    dataset = YourDataset()
    loader, sampler = create_dataloader(dataset, config, world_size, rank)

    criterion = nn.CrossEntropyLoss()

    # Bucle de entrenamiento
    for epoch in range(config.epochs):
        sampler.set_epoch(epoch)

        loss, acc = train_epoch(model, loader, criterion, optimizer, scaler, config, rank)

        # Imprimir solo en rank 0
        if rank == 0:
            print(f"Epoch {epoch+1}/{config.epochs}: Loss={loss:.4f}, Acc={acc:.4f}")

    dist.destroy_process_group()


# Forma de inicio:
# torchrun --nnodes=2 --nproc_per_node=4 --node_rank=0 --master_addr=192.168.1.1 --master_port=29500 train.py
# torchrun --nnodes=2 --nproc_per_node=4 --node_rank=1 --master_addr=192.168.1.1 --master_port=29500 train.py

La clave del entrenamiento distribuido es: 1) elegir el modo de paralelismo adecuado (paralelismo de datos vs paralelismo de modelos); 2) usar DDP en lugar de DataParallel para obtener un mejor rendimiento; 3) combinar el entrenamiento de precisión mixta para acelerar aún más; 4) configurar adecuadamente el cargador de datos y los parámetros de comunicación.


7. Monitoreo y depuración del entrenamiento distribuido

7.1 Herramientas de monitoreo comunes

Ejemplo

# Código de monitoreo del entrenamiento distribuido

def monitor_training(rank, world_size):
    """Monitorear el estado del entrenamiento distribuido"""
    # 1. Verificar el estado del grupo de procesos
    print(f"Rank: {rank}, World Size: {world_size}")
    print(f"Backend: {dist.get_backend()}")
    print(f"Rank: {dist.get_rank()}")
    print(f"CUDA Device: {torch.cuda.current_device()}")

    # 2. Monitorear el uso de memoria de GPU
    if torch.cuda.is_available():
        for i in range(torch.cuda.device_count()):
            allocated = torch.cuda.memory_allocated(i) / 1024**2
            reserved = torch.cuda.memory_reserved(i) / 1024**2
            print(f"GPU {i}: asignado {allocated:.1f} MB, reservado {reserved:.1f} MB")

    # 3. Monitorear la sincronización de gradientes
    # DDP registra automáticamente el tiempo de sincronización
    # Se pueden habilitar registros detallados configurando variables de entorno
    # os.environ["NCCL_DEBUG"] = "INFO"


# Llamar periódicamente en el bucle de entrenamiento
monitor_training(rank, dist.get_world_size())

7.2 Consejos de depuración

  • Probar primero en una sola máquina, luego expandir a múltiples máquinas
  • Usaros.environ["NCCL_DEBUG"] = "INFO"Ver registros de comunicación detallados
  • Asegurar que todas las máquinas tengan el tiempo sincronizado (usando NTP)
  • Verificar si el firewall permite los puertos de comunicación NCCL

Aunque el entrenamiento distribuido aumenta la complejidad, la mejora de rendimiento que aporta es significativa. Para el entrenamiento de modelos grandes, es casi una técnica imprescindible.

Otras extensiones