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.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:
- Copia del modelo: Copiar el modelo a cada GPU.
- Distribución de datos: Distribuir uniformemente los datos del batch a cada GPU.
- Cálculo paralelo: Cada GPU calcula la propagación hacia adelante de forma independiente.
- Agregación de gradientes: Reunir los gradientes de todas las GPU en la GPU principal.
- Actualización de parámetros: La GPU principal actualiza los parámetros y los sincroniza con las otras GPU.
Ejemplo
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.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
# 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
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 a
sampler.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.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:
- Cada GPU divide los gradientes en n partes (n es el número de GPU)
- Cada GPU intercambia una parte de gradientes con la GPU adyacente
- Después de repetir n veces, cada GPU tiene los gradientes completos
Ejemplo
# 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
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
# 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
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
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
- Usar
os.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
Otras extensionesAunque 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.