1. Introducción
tener un modelo. Tienes una sola GPU. La formación dura 72 horas. Solicita una segunda máquina con cuatro GPU más y ahora necesita su código para usarlas. Este es el momento exacto en el que la mayoría de los practicantes chocan contra una pared. No porque la capacitación distribuida sea conceptualmente difícil, sino porque la ingeniería necesaria para realizarla correctamente (grupos de procesos, registro de rango, siembra de muestras, barreras de puntos de control) se encuentra dispersa en docenas de tutoriales, cada uno de los cuales cubre una pieza del rompecabezas.
Este artículo es la guía que desearía tener cuando amplié el entrenamiento por primera vez más allá de un solo nodo. Construiremos un canal de capacitación de múltiples nodos completo y de nivel de producción desde cero utilizando DistributedDataParallel (DDP) de PyTorch. Cada archivo es modular, cada valor es configurable y cada concepto distribuido se hace explícito. Al final, tendrá una base de código que podrá colocar en cualquier clúster y comenzar a entrenar de inmediato.
Lo que cubriremos: el modelo mental detrás de DDP, una estructura de proyecto modular limpia, gestión distribuida del ciclo de vida, carga eficiente de datos entre rangos, un ciclo de entrenamiento con precisión mixta y acumulación de gradientes, registro y puntos de control con reconocimiento de rangos, scripts de lanzamiento de múltiples nodos y los problemas de rendimiento que hacen tropezar incluso a los ingenieros experimentados.
El código base completo está disponible en GitHub. Cada bloque de código de este artículo se extrae directamente de ese repositorio.
2. Cómo funciona DDP: el modelo mental
Antes de escribir cualquier código, necesitamos un modelo mental claro. DistributedDataParallel (DDP) no es mágico: es un patrón de comunicación bien definido construido sobre operaciones colectivas.
La configuración es sencilla. Inicia N procesos (uno por GPU, potencialmente en varias máquinas). Cada proceso inicializa un grupo de procesos: un canal de comunicación respaldado por NCCL (Biblioteca de comunicaciones colectivas de NVIDIA) para transferencias de GPU a GPU. Cada proceso obtiene tres números de identidad: su rango global (único en todas las máquinas), su rango local (único dentro de su máquina) y el tamaño mundial total.
Cada proceso contiene una copia idéntica del modelo. Los datos se dividen entre procesos mediante DistributedSampler: cada rango ve una porción diferente del conjunto de datos, pero los pesos del modelo comienzan (y permanecen) idénticos.
El mecanismo crítico es lo que sucede durante el proceso de retroceso(). DDP registra ganchos en cada parámetro. Cuando se calcula un gradiente para un parámetro, DDP lo agrupa con gradientes cercanos y activa una operación de reducción total en todo el grupo de procesos. Esta reducción total calcula el gradiente medio en todos los rangos. Debido a que ahora cada rango tiene el mismo gradiente promedio, el siguiente paso del optimizador produce actualizaciones de peso idénticas, manteniendo todas las réplicas sincronizadas, sin ningún código de sincronización explícito por nuestra parte.
Esta es la razón por la que DDP es estrictamente superior al antiguo DataParallel: no existe un único cuello de botella de GPU "maestro", ni pases hacia adelante redundantes, y la comunicación de gradiente se superpone con el cálculo hacia atrás.
Terminología clave
3. Descripción general de la arquitectura
Un proceso de capacitación en producción nunca debe ser un guión monolítico único. El nuestro se divide en seis módulos enfocados, cada uno con una única responsabilidad. El siguiente gráfico de dependencia muestra cómo se conectan; tenga en cuenta que config.py se encuentra en la parte inferior y actúa como la única fuente de verdad para cada hiperparámetro.
Aquí está la estructura del proyecto:
pytorch-multinode-ddp/ ├── train.py # Punto de entrada: bucle de entrenamiento ├── config.py # Configuración de clase de datos + argparse ├── ddp_utils.py # Configuración distribuida, desmontaje y puntos de control ├── model.py # MiniResNet (variante ligera de ResNet) ├── dataset.py # Conjunto de datos sintéticos + Cargador DistributedSampler ├── utils/ │ ├── logger.py # Registro estructurado con reconocimiento de rangos │ └── metrics.py # Promedios móviles + reducción total distribuida ├── scripts/ │ └── launch.sh # Envoltorio de ejecución de antorcha de múltiples nodos └── requisitos.txt
Esta separación significa que puede intercambiar un conjunto de datos real editando solo dataset.py o reemplazar el modelo editando solo model.py. El circuito de entrenamiento nunca necesita cambiar.
4. Configuración centralizada
Los hiperparámetros codificados son enemigos de la reproducibilidad. Usamos una clase de datos de Python como nuestra única fuente de configuración. Todos los demás módulos importan TrainingConfig y leen de él; nada está codificado.
La clase de datos también funciona como nuestro analizador CLI: el método de clase from_args() realiza una introspección de los nombres y tipos de campos, generando automáticamente indicadores argparse con valores predeterminados. Esto significa que obtienes –batch_size 128 y –no-use_amp gratis, sin escribir una sola línea del analizador a mano.
@dataclass class TrainingConfig: """Bolsa inmutable de cada parámetro que necesita el proceso de capacitación.""" # Modelo num_classes: int = 10 in_channels: int = 3 image_size: int = 32 # Datos de tamaño de lote: int = 64 # por GPU num_workers: int = 4 # Épocas del optimizador/programador: int = 10 lr: float = 0.01 impulso: float = 0.9 Weight_decay: float = 1e-4 # Backend distribuido: str = "nccl" # Precisión mixta use_amp: bool = True # Acumulación de gradiente grad_accum_steps: int = 1 # Puntos de control checkpoint_dir: str = "./checkpoints" save_every: int = 1 resume_from: Opcional[str] = Ninguno # Registro y creación de perfiles log_interval: int = 10 enable_profiling: bool = False semilla: int = 42 @classmethod def from_args(cls) -> "TrainingConfig": parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter) defaults = cls() para nombre, val en vars(defaults).items(): arg_type = type(val) if val no es Ninguno más str if isinstance(val, bool): parser.add_argument(f"–{name}", default=val, action=argparse.BooleanOptionalAction) else: parser.add_argument(f"–{name}", type=arg_type, default=val) return cls(**vars(parser.parse_args()))
¿Por qué una clase de datos en lugar de YAML o JSON? Tres razones: (1) el IDE y mypy aplican sugerencias de tipo, (2) no hay dependencia de bibliotecas de configuración de terceros y (3) cada parámetro tiene un valor predeterminado visible justo al lado de su declaración. Para sistemas de producción que necesitan configuraciones jerárquicas, siempre puede colocar Hydra u OmegaConf encima de este patrón.
5. Gestión distribuida del ciclo de vida
El ciclo de vida distribuido tiene tres fases: inicializar, ejecutar y desmantelar. Hacer algo incorrecto en cualquiera de estos casos puede producir bloqueos silenciosos, por lo que envolvemos todo en un manejo de errores explícito.
Inicialización del grupo de procesos
La función setup_distributed() lee las tres variables de entorno que torchrun establece automáticamente (RANK, LOCAL_RANK, WORLD_SIZE), fija la GPU correcta con torch.cuda.set_device() e inicializa el grupo de procesos NCCL. Devuelve una clase de datos congelada (DistributedContext) que el resto del código base pasa en lugar de volver a leer os.environ.
@dataclass(frozen=True) class DistributedContext: """Instantánea inmutable de la identidad distribuida del proceso actual.""" rango: int local_rank: int world_size: int dispositivo: torch.device def setup_distributed(config: TrainingConfig) -> DistributedContext: require_vars = ("RANK", "LOCAL_RANK", "WORLD_SIZE") desaparecido = [v para v en require_vars si v no en os.environ] si falta: aumente RuntimeError( f"Faltan variables de entorno: {faltantes}. " "Inicie con torchrun o configúrelas manualmente"). = int(os.environ["WORLD_SIZE"]) torch.cuda.set_device(local_rank) dispositivo = torch.device("cuda", local_rank) dist.init_process_group(backend=config.backend) return DistributedContext( ranking=rank, local_rank=local_rank, world_size=world_size, dispositivo=dispositivo) Puntos de control con Rank Guards
El error más común en los puntos de control distribuidos es que todos los rangos escriben en el mismo archivo simultáneamente. Protegemos el guardado detrás de is_main_process() y la carga detrás de dist.barrier(); esto garantiza que el rango 0 termine de escribir antes de que otros rangos intenten leer.
def save_checkpoint(ruta, época, modelo, optimizador, escalador=Ninguno, rango=0): """Persistir el estado de entrenamiento en el disco (solo rango-0).""" si no es_main_process(rank): return Ruta(ruta).parent.mkdir(parents=True, exist_ok=True) state = { "epoch": epoch, "model_state_dict": model.module.state_dict(), "optimizer_state_dict": optimizador.state_dict(), } si el escalador no es Ninguno: estado["scaler_state_dict"] = escalar.state_dict() torch.save(estado, ruta) def load_checkpoint(ruta, modelo, optimizador=None, escalador=None, dispositivo="cpu"): """Restaurar el estado de entrenamiento. Todos los rangos se cargan después de la barrera.""" dist.barrier() # esperar a que el rango 0 termine de escribir ckpt = torch.load(path, map_location=device,weights_only=False) model.load_state_dict(ckpt["model_state_dict"]) si optimizador y "optimizer_state_dict" en ckpt: optimizador.load_state_dict(ckpt["optimizer_state_dict"]) si escalador y "scaler_state_dict" en ckpt: escalar.load_state_dict(ckpt["scaler_state_dict"]) devuelve ckpt.get("época", 0)
6. Diseño de modelo para DDP
Usamos una variante liviana de ResNet llamada MiniResNet: tres etapas residuales con canales crecientes (64, 128, 256), dos bloques por etapa, agrupación promedio global y un cabezal completamente conectado. Es lo suficientemente complejo como para ser realista pero lo suficientemente liviano como para ejecutarse en cualquier hardware.
El requisito crítico de DDP: el modelo debe moverse a la GPU correcta antes de empaquetarlo. DDP no mueve modelos por usted.
def create_model(config: TrainingConfig, dispositivo: torch.device) -> nn.Module: """Crear una instancia de MiniResNet y moverlo al dispositivo.""" model = MiniResNet( in_channels=config.in_channels, num_classes=config.num_classes, ) return model.to(device) def wrap_ddp(model: nn.Module, local_rank: int) -> DDP: """Ajustar modelo con DistributedDataParallel.""" return DDP(modelo, device_ids=[local_rank])
Tenga en cuenta el patrón de dos pasos: create_model() → wrap_ddp(). Esta separación es intencional. Al cargar un punto de control, necesita el modelo desenvuelto (model.module) para cargar los dictados de estado y luego volver a empaquetar. Si fusiona la creación y el ajuste, la carga del punto de control se vuelve incómoda.
7. Carga de datos distribuidos
DistributedSampler es lo que garantiza que cada GPU vea una porción única de datos. Divide los índices en rangos de tamaño mundial y devuelve un subconjunto que no se superpone para cada uno. Sin él, cada GPU se entrenaría en lotes idénticos, quemando cómputo sin ningún beneficio.
Hay tres detalles que hacen tropezar a la gente:
Primero, se debe llamar a sampler.set_epoch(epoch) al comienzo de cada época. El muestreador utiliza el número de época como semilla aleatoria para barajar. Si olvida esto, cada época repetirá los datos en el mismo orden, lo que degrada la generalización.
En segundo lugar, pin_memory=True en DataLoader preasigna memoria de host bloqueada por página, lo que permite transferencias asincrónicas de CPU a GPU cuando llamas a tensor.to(device, non_blocking=True). De esta superposición es de donde provienen las ganancias reales de rendimiento.
En tercer lugar, persistente_workers=True evita la reaparición de procesos de trabajo en cada época, una reducción significativa de los gastos generales cuando num_workers > 0.
def create_distributed_dataloader(conjunto de datos, configuración, ctx): muestra = DistributedSampler( conjunto de datos, num_replicas=ctx.world_size, rango=ctx.rank, shuffle=True,) cargador = DataLoader( conjunto de datos, tamaño_por lotes=config.batch_size, muestra=sampler, num_workers=config.num_workers, pin_memory=True, drop_last=True, persistente_workers=config.num_workers > 0, ) cargador de retorno, muestra
8. El circuito de formación: donde todo confluye
Este es el corazón del oleoducto. El siguiente bucle integra todos los componentes que hemos creado hasta ahora: modelo envuelto en DDP, cargador de datos distribuido, precisión mixta, acumulación de gradiente, registro con reconocimiento de rango, programación de tasa de aprendizaje y puntos de control.
Precisión mixta (AMP)
La precisión mixta automática (AMP) mantiene los pesos maestros en FP32 pero ejecuta el cálculo de pérdida y pase hacia adelante en FP16. Esto reduce a la mitad los requisitos de ancho de banda de memoria y permite la aceleración de Tensor Core en las GPU NVIDIA modernas, lo que a menudo produce una mejora del rendimiento de 1,5 a 2 veces con un impacto insignificante en la precisión.
Usamos torch.autocast para el pase directo y torch.amp.GradScaler para escalar pérdidas. Una sutileza: creamos el GradScaler con enable=config.use_amp. Cuando está deshabilitado, el escalador deja de funcionar: la misma ruta de código, cero gastos generales y sin bifurcaciones.
Acumulación de gradiente
A veces necesita un tamaño de lote efectivo mayor que el que permite la memoria de su GPU. La acumulación de gradiente simula esto ejecutando múltiples pases hacia adelante y hacia atrás antes de pasar al optimizador. La clave es dividir la pérdida por grad_accum_steps antes de back(), para que el gradiente acumulado se promedie correctamente.
def train_one_epoch(modelo, cargador, criterio, optimizador, escalador, ctx, configuración, época, registrador): model.train() tracker = MetricTracker() total_steps = len(loader) use_amp = config.use_amp and ctx.device.type == "cuda" autocast_ctx = torch.autocast("cuda", dtype=torch.float16) if use_amp else nullcontext() optimizador.zero_grad(set_to_none=True) para el paso, (imágenes, etiquetas) en enumerar(cargador): imágenes = imágenes.to(ctx.device, non_blocking=True) etiquetas = etiquetas.to(ctx.device, non_blocking=True) con autocast_ctx: salidas = modelo(imágenes) pérdida = criterio(salidas, etiquetas) pérdida = pérdida / config.grad_accum_steps # escala para acumulación scaler.scale(loss).backward() if (step + 1) % config.grad_accum_steps == 0: scaler.step(optimizer) scaler.update() optimizador.zero_grad(set_to_none=True) # reinicio eficiente de la memoria # Seguimiento de la pérdida sin procesar (sin escalar) para el registro raw_loss = loss.item() * config.grad_accum_steps acc = Compute_accuracy(salidas, etiquetas) tracker.update("loss", raw_loss, n=images.size(0)) tracker.update("accuracy", acc, n=images.size(0)) if is_main_process(ctx.rank) y (paso + 1) % config.log_interval == 0: log_training_step(logger, época, paso + 1, total_steps, raw_loss, optimizador.param_groups[0]["lr"]) rastreador de retorno
Dos detalles dignos de destacar. Primero, zero_grad(set_to_none=True) desasigna tensores de gradiente en lugar de llenarlos con ceros, ahorrando memoria proporcional al tamaño del modelo. En segundo lugar, los datos se mueven a la GPU con non_blocking=True; esto permite que la CPU continúe llenando el siguiente lote mientras se transfiere el actual, explotando la superposición de pin_memory.
La función principal
La función main() organiza todo el proceso. Tenga en cuenta el patrón try/finally que garantiza que el grupo de procesos se desmantele incluso si ocurre una excepción; sin esto, una falla en un rango puede dejar otros rangos colgados indefinidamente.
def main(): config = TrainingConfig.from_args() ctx = setup_distributed(config) logger = setup_logger(ctx.rank) torch.manual_seed(config.seed + ctx.rank) model = create_model(config, ctx.device) model = wrap_ddp(model, ctx.local_rank) optimizador = torch.optim.SGD(model.parameters(), lr=config.lr, momento=config.momentum, peso_decay=config.weight_decay) planificador = CosineAnnealingLR(optimizador, T_max=config.epochs) escalador = torch.amp.GradScaler(enabled=config.use_amp) start_epoch = 1 if config.resume_from: start_epoch = load_checkpoint(config.resume_from, model.module, optimizador, escalador, ctx.device) + 1 conjunto de datos = SyntheticImageDataset(size=50000, image_size=config.image_size, num_classes=config.num_classes) cargador, sampler = create_distributed_dataloader(dataset, config, ctx) criterio = nn.CrossEntropyLoss() intente: para la época en el rango (start_epoch, config.epochs + 1): sampler.set_epoch(epoch) tracker = train_one_epoch(modelo, cargador, criterio, optimizador, escalador, ctx, config, epoch, logger) Scheduler.step() avg_loss = all_reduce_scalar(tracker.average("pérdida"), ctx.world_size, ctx.device) if is_main_process(ctx.rank): log_epoch_summary(logger, epoch, {"loss": avg_loss}) if epoch % config.save_every == 0: save_checkpoint(f"checkpoints/epoch_{epoch}.pt", epoch, model, optimizador, escalador, ctx.rank) finalmente: cleanup_distributed()
9. Lanzamiento entre nodos
torchrun de PyTorch (introducido en v1.10 como reemplazo de torch.distributed.launch) maneja la generación de un proceso por GPU y la configuración de las variables de entorno RANK, LOCAL_RANK y WORLD_SIZE. Para el entrenamiento de múltiples nodos, cada nodo debe especificar la dirección del nodo maestro para que todos los procesos puedan establecer la conexión NCCL.
Aquí está nuestro script de inicio, que lee todos los ajustes de las variables de entorno:
#!/usr/bin/env bash set -euo pipefail NNODES="${NNODES:-2}" NPROC_PER_NODE="${NPROC_PER_NODE:-4}" NODE_RANK="${NODE_RANK:-0}" MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}" MASTER_PORT="${MASTER_PORT:-12355}" torchrun –nnodes="${NNODES}" –nproc_per_node="${NPROC_PER_NODE}" –node_rank="${NODE_RANK}" –master_addr="${MASTER_ADDR}" –master_port="${MASTER_PORT}" train.py "$@"
Para una prueba rápida de un solo nodo en una GPU:
torchrun –standalone –nproc_per_node=1 train.py –epochs 2
Para el entrenamiento de dos nodos con cuatro GPU cada uno, ejecútelo en el nodo 0:
MASTER_ADDR=10.0.0.1 NODE_RANK=0 NNODES=2 NPROC_PER_NODE=4 scripts bash/launch.sh
Y en el nodo 1:
MASTER_ADDR=10.0.0.1 NODE_RANK=1 NNODES=2 NPROC_PER_NODE=4 scripts bash/launch.sh
10. Errores y consejos de rendimiento
Después de crear cientos de trabajos de capacitación distribuidos, estos son los errores que veo con más frecuencia:
Olvidando sampler.set_epoch(). Sin él, el orden de los datos es idéntico en cada época. Este es el error de DDP más común y silenciosamente perjudica la convergencia.
Cuello de botella en la transferencia CPU-GPU. Utilice siempre pin_memory=True en su DataLoader y non_blocking=True en sus llamadas .to(). Sin estos, la CPU se bloquea en cada transferencia por lotes.
Registro de todos los rangos. Si se imprimen todos los rangos, la salida es basura intercalada. Proteja todos los registros detrás del rango == 0 controles.
zero_grad() sin set_to_none=True. El valor predeterminado zero_grad() llena los tensores de gradiente con ceros. set_to_none=True los desasigna, reduciendo el pico de memoria.
Guardando puntos de control de todos los rangos. Múltiples rangos que escriben el mismo archivo provocan corrupción. Solo se debe guardar el rango 0 y todos los rangos deben hacer una barrera antes de cargar.
No sembrar con compensación de rango. torch.manual_seed(seed + ranking) garantiza que el aumento de datos de cada rango sea diferente. Sin la compensación, los aumentos son idénticos en todas las GPU.
Cuándo NO usar DDP
DDP replica el modelo completo en cada GPU. Si su modelo no cabe en la memoria de una sola GPU, el DDP por sí solo no ayudará. Para tales casos, busque en Fully Sharded Data Parallel (FSDP), que fragmenta parámetros, gradientes y estados del optimizador en todos los rangos, o marcos como DeepSpeed ZeRO.
11. Conclusión
Hemos pasado de una mentalidad de capacitación de una sola GPU a una canalización de nivel de producción totalmente distribuida capaz de escalar entre máquinas, sin sacrificar la claridad ni la capacidad de mantenimiento.
Pero lo más importante es que no se trataba sólo de hacer que DDP funcionara. Se trataba de construirlo correctamente.
Resumamos las conclusiones más importantes:
Conclusiones clave
DDP es ingeniería determinista, no magia
Una vez que comprenda los grupos de procesos, los rangos y la reducción total, la capacitación distribuida se vuelve predecible y depurable. La estructura importa más que la escala
Una base de código limpia y modular (configuración → datos → modelo → entrenamiento → utilidades) es lo que hace factible el escalado de 1 GPU a 100 GPU. La fragmentación correcta de datos no es negociable
DistributedSampler + set_epoch() es la diferencia entre el escalado verdadero y el cálculo desperdiciado. El rendimiento proviene de pequeños detalles
pin_memory, non_blocking, set_to_none=True y AMP en conjunto ofrecen ganancias masivas de rendimiento. La conciencia de rango es esencial
El registro, los puntos de control y la aleatoriedad deben respetar el rango; de lo contrario, se producirá el caos. DDP escala la computación, no la memoria
Si su modelo no cabe en una GPU, necesita FSDP o ZeRO, no más GPU.
El panorama más amplio
Lo que ha creado aquí no es solo un guión de capacitación: es una plantilla para sistemas de aprendizaje automático del mundo real.
Este patrón exacto se utiliza en:
Producción Canales de ML Laboratorios de investigación que entrenan modelos grandes Startups que escalan desde el prototipo hasta la infraestructura
¿Y la mejor parte?
Ahora puedes:
Conecte un conjunto de datos real Intercambie un transformador o una arquitectura personalizada Escale entre nodos sin cambios de código
Qué explorar a continuación
Una vez que se sienta cómodo con esta configuración, la próxima frontera es el entrenamiento a gran escala y con eficiencia de memoria:
Paralelo de datos totalmente fragmentados (FSDP) → modelo de fragmentos + gradientes DeepSpeed ZeRO → estados del optimizador de fragmentos Paralelismo de canalización → dividir modelos entre GPU Paralelismo tensorial → dividir las capas mismas
Estas técnicas impulsan los modelos más grandes de la actualidad, pero todas se basan en la base exacta de DDP que ahora comprende.
La capacitación distribuida a menudo resulta intimidante, no porque sea intrínsecamente compleja, sino porque rara vez se presenta como un sistema completo.
Ahora has visto la imagen completa.
Y una vez que lo veas de principio a fin…
El escalamiento se convierte en una decisión de ingeniería, no en un problema de investigación.
¿Qué sigue?
Este canal maneja el entrenamiento de datos en paralelo, el patrón distribuido más común. Cuando sus modelos superen la memoria de una sola GPU, explore Fully Sharded Data Parallel (FSDP) para fragmentar parámetros o DeepSpeed ZeRO para particionar en estado de optimizador. Para modelos verdaderamente masivos, el paralelismo de canalización (dividir el modelo en GPU capa por capa) y el paralelismo tensorial (dividir capas individuales) se vuelven necesarios.
Pero para la gran mayoría de las cargas de trabajo de capacitación (desde ResNets hasta Transformers de mediana escala), el canal DDP que creamos aquí es exactamente lo que usan los equipos de producción. Escale agregando nodos y GPU; el código se encarga del resto.
El código base completo y listo para producción para este proyecto está disponible aquí: pytorch-multinode-ddp
Referencias
[1]Descripción general distribuida de PyTorch, documentación de PyTorch (2024), https://pytorch.org/tutorials/beginner/dist_overview.html
[2]S. Li et al., PyTorch Distributed: Experiencias sobre la aceleración de la capacitación paralela de datos (2020), VLDB Endowment
[3]API PyTorch DistributedDataParallel, https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html
[4]NCCL: primitivas optimizadas para comunicación colectiva multi-GPU, NVIDIA, https://developer.nvidia.com/nccl
[5]PyTorch AMP: Precisión mixta automática, https://pytorch.org/docs/stable/amp.html