Distributed Training

El entrenamiento distribuido reparte el trabajo de entrenar un modelo entre varios procesadores, normalmente GPU o TPU en una máquina o en un clúster en la nube o local, para terminar antes el entrenamiento o poder entrenar un modelo que no cabe en la memoria de un único dispositivo. Los dispositivos intercambian gradientes, activaciones o fragmentos de parámetros a través de interconexiones rápidas mediante bibliotecas de comunicación como NVIDIA NCCL. Así se entrenan los modernos modelos fundacionales y modelos de lenguaje grandes, y también se acorta el entrenamiento cotidiano de detección de objetos en estaciones de trabajo con varias GPU.

Cómo funciona el entrenamiento distribuido#

El paralelismo de datos es la estrategia más habitual. Cada GPU contiene una réplica completa del modelo, procesa su propia parte de cada lote de datos de entrenamiento y participa en una operación all-reduce que promedia los gradientes para que cada réplica aplique la misma actualización de descenso del gradiente. DistributedDataParallel (DDP) de PyTorch solapa esa comunicación con la pasada hacia atrás. El tamaño global del lote es el tamaño del lote por GPU multiplicado por el número de GPU, así que normalmente hay que escalar la tasa de aprendizaje en consecuencia: Goyal et al. entrenaron ResNet-50 con ImageNet en una hora y en 256 GPU, con lotes de 8,192 imágenes, usando una regla de escalado lineal y una fase de calentamiento, con una eficiencia de escalado de alrededor del 90 % entre 8 y 256 GPU.

El paralelismo de datos con fragmentación reduce el coste de memoria de esas réplicas. ZeRO y Fully Sharded Data Parallel (FSDP) de PyTorch distribuyen los parámetros, los gradientes y los estados del optimizador entre las GPU, y reúnen todos los pesos de cada capa solo mientras se calcula esa capa; así, la memoria necesaria en cada GPU disminuye a medida que se añaden GPU.

El paralelismo de modelos divide la propia red neuronal. El paralelismo de tensores distribuye las multiplicaciones de matrices individuales entre las GPU; el paralelismo de tuberías asigna capas consecutivas a distintas GPU y transmite micro lotes a través de ellas; el paralelismo de contexto divide secuencias largas, y el paralelismo de expertos coloca los expertos de un modelo de mezcla de expertos en distintas GPU. Frameworks como Megatron-LM combinan estas técnicas con el paralelismo de datos para entrenar modelos con cientos de miles de millones de parámetros.

El entrenamiento con varias GPU en la práctica#

La comunicación es el principal coste adicional: el tráfico de all-reduce aumenta con el tamaño del modelo, por lo que las GPU conectadas mediante NVLink dentro de un nodo escalan mejor que los nodos conectados a través de redes más lentas. Las estadísticas de normalización por lotes se calculan en cada GPU, salvo que las capas se conviertan a normalización por lotes sincronizada; esto importa cuando el fragmento de lote de cada GPU es pequeño. Ultralytics YOLO26 ejecuta el entrenamiento con varias GPU mediante DDP: al pasar varios identificadores de dispositivo, se inicia un proceso por GPU con torch.distributed.run, se utiliza el backend NCCL en CUDA y se reparte por igual entre las GPU el valor de batch.

from ultralytics import YOLO

# Carga un modelo YOLO26 preentrenado
model = YOLO("yolo26n.pt")

# Entrena en las GPU 0 y 1 con DDP; batch=32 proporciona 16 imágenes a cada GPU en cada paso
results = model.train(data="coco8.yaml", epochs=100, imgsz=640, batch=32, device=[0, 1])

Aplicaciones en el mundo real#

• Modelos de lenguaje grandes: Meta entrenó Llama 3 405B con hasta 16K GPU H100, combinando el paralelismo de tensores, de tuberías, de contexto y de datos en lo que sus autores denominan paralelismo 4D.
• Vehículos autónomos: Los equipos de percepción entrenan modelos de detección y segmentación con grandes archivos de grabaciones de cámara, lidar y radar de flotas de vehículos, algo que solo resulta práctico en grandes clústeres de GPU.
• Análisis de imágenes médicas: Los volúmenes de TC y RM 3D de alta resolución saturan la memoria de una sola GPU, y distribuir los lotes o los fragmentos del modelo entre varias GPU permite entrenar con volúmenes y modelos más grandes.

Entrenamiento distribuido y conceptos relacionados#

• Aprendizaje federado: El aprendizaje federado también entrena en muchos dispositivos, pero, por privacidad, los datos permanecen en los teléfonos, los hospitales u otros clientes, y solo las actualizaciones del modelo viajan a un servidor central. El entrenamiento distribuido reúne los datos en un clúster para maximizar la velocidad.
• Aprendizaje dividido: El aprendizaje dividido reparte un modelo entre un cliente y un servidor para que los datos sin procesar nunca salgan del cliente, a cambio de más comunicación y mayor privacidad.
• Acumulación de gradientes: La acumulación de gradientes suma los gradientes de varios lotes pequeños antes de aplicar una actualización, para reproducir los cálculos de un lote grande en una sola GPU a costa de más tiempo; el paralelismo de datos, en cambio, procesa esos lotes simultáneamente en varias GPU.