Planifique una ejecución reproducible de PyTorch

Demuestre un paso completo antes de escalar

Fije las revisiones del modelo y del conjunto de datos, las versiones de dependencias, la configuración de entrenamiento, las semillas aleatorias y el resumen del contenedor. Ejecute la carga del modelo, un paso completo de entrenamiento, la evaluación, el guardado del punto de control y la restauración del punto de control en un acelerador.

Utilice la misma longitud de secuencia y configuración de micro-lotes prevista para la ejecución más larga. Un piloto diminuto que omita el estado del optimizador o use secuencias más cortas puede ocultar el problema de memoria que intenta medir.

Mida un paso de optimización representativo

Dentro de su programa de entrenamiento existente, rodee un paso representativo con una medición de memoria máxima. El ejemplo asume que el modelo, el lote y el optimizador ya se han inicializado en la GPU. Incluya el paso del optimizador porque parte del estado se crea de forma diferida.

torch.cuda.reset_peak_memory_stats()
optimizer.zero_grad(set_to_none=True)
loss = model(**batch).loss
loss.backward()
optimizer.step()
torch.cuda.synchronize()
peak_gib = torch.cuda.max_memory_allocated() / 2**30
print(f"Peak tensor memory: {peak_gib:.2f} GiB")

Esto informa de la asignación máxima de tensores, no de la huella completa del dispositivo. Inspeccione también la memoria reservada, la sobrecarga del tiempo de ejecución y el informe de memoria del dispositivo del sistema.

Elige DDP o fragmentación por el motivo correcto

DistributedDataParallel mantiene una réplica del modelo en cada proceso y sincroniza los gradientes. Es útil cuando el estado completo de entrenamiento cabe en cada GPU y se desea procesar más datos en paralelo; no fusiona la VRAM en un grupo contiguo.

Fully Sharded Data Parallel puede distribuir parámetros, gradientes y estado del optimizador. Úselo cuando la presión de memoria justifique el particionamiento y tenga en cuenta la comunicación adicional y la complejidad de los puntos de control distribuidos.

La precisión mixta y el checkpointing de activaciones son decisiones independientes. Valide el comportamiento numérico para la precisión más baja; mida el coste computacional de recalcular activaciones cuando el checkpointing ahorra memoria.

Lanzar una prueba piloto de un solo nodo

El script de entrenamiento debe inicializar la ejecución distribuida y asignar un proceso a cada GPU utilizando la información de rango local. Haga coincidir el número de procesos con las GPU visibles. Lo siguiente es un ejemplo de lanzamiento para un script con soporte distribuido, no una implementación de entrenamiento completa.

torchrun --standalone --nproc-per-node=4 \
  train.py --config configs/pilot.yaml

Ejecuta esto solo en un nodo aprovisionado. La demostración local de AnchorGPU no ejecuta entrenamiento ni asigna GPU remotas.

Seleccione el hardware según el resultado del piloto

Utilice A100 como base de 80 GB CUDA de menor precio, y pruebe H100 si la carga de trabajo tiene una ruta útil específica de Hopper. H200 ofrece una configuración de memoria CUDA mayor. MI300X ofrece más capacidad por acelerador en este catálogo, siempre que la aplicación completa esté cualificada en ROCm.

Para AMD, compruebe las extensiones personalizadas, los kernels, las versiones de los paquetes y la imagen PyTorch probada con ROCm como un conjunto completo. La compatibilidad a nivel de framework no valida todas las operaciones opcionales.

Pruebe la restauración, no solo el guardado

Registre el modelo, el optimizador, el programador, el paso de entrenamiento, la configuración y el estado del escalador de precisión mixta cuando corresponda. Guarde un punto de control, termine el proceso y luego restaure y reanude antes de comprometerse con la ejecución larga.

Utilice un método de punto de control apropiado para el estado distribuido. Pruebe que los artefactos guardados puedan leerse independientemente del proceso original y que la salida de evaluación permanezca dentro de la tolerancia esperada.

Presupueste todo el proceso

Incluye descargas, preprocesamiento, carga de datos, validación, puntos de control, exportación y copia de salidas fuera del nodo. Mide el rendimiento junto con la memoria, las pausas de carga de datos, el tiempo de comunicación, la duración de los puntos de control y la calidad de la validación.

Un plazo de 7 días es útil para una pasada acotada de compatibilidad y perfilado. Elige una reserva de 30 días solo después de poder estimar el trabajo útil y el margen operativo. Más GPUs ayudan únicamente cuando la computación adicional supera los límites de comunicación y del flujo de entrada.