Distributed Training/es

Article
Topic area	Deep Learning
Prerequisites	Stochastic Gradient Descent, Backpropagation

Resumen

El entrenamiento distribuido es la práctica de dividir un único trabajo de entrenamiento de una red neuronal entre múltiples procesadores, aceleradores o máquinas de modo que la carga de trabajo se ejecute en paralelo. Aborda dos límites del entrenamiento en un solo dispositivo: el modelo puede ser demasiado grande para caber en la memoria de un dispositivo, y el tiempo de reloj de pared para converger sobre un conjunto de datos grande puede resultar inaceptable. Los modelos fundacionales modernos, con cientos de miles de millones de parámetros, no pueden entrenarse en absoluto sin él, e incluso los modelos más pequeños lo emplean rutinariamente para acortar el tiempo de entrenamiento de semanas a horas. El campo combina ideas de la optimización numérica, la computación de alto rendimiento y la ingeniería de sistemas: cómo coordinar la optimización estocástica entre los trabajadores, cómo mover tensores de manera eficiente entre los dispositivos y cómo programar el cálculo para que los aceleradores se mantengan ocupados.

Los dos ejes fundamentales son el paralelismo de datos, en el que cada trabajador mantiene una réplica del modelo y procesa una partición distinta del lote, y el paralelismo de modelo, en el que el modelo mismo se divide entre los trabajadores. Los sistemas del mundo real combinan ambos con el paralelismo de canalización y la fragmentación del estado del optimizador para escalar a miles de aceleradores manteniendo alta intensidad aritmética y comunicación tolerable.

Motivación

Tres presiones empujan al entrenamiento fuera de un único dispositivo. Primero, el número de parámetros: un modelo con $ N $ parámetros en precisión de 16 bits necesita al menos $ 2N $ bytes para los pesos, con cantidades similares para los gradientes y el estado del optimizador, de modo que un modelo de 70 mil millones de parámetros en precisión mixta supera fácilmente 1 TB de memoria una vez incluidos los momentos de Adam. Segundo, el volumen de datos: entrenar un modelo de lenguaje grande sobre billones de tokens al rendimiento de un solo dispositivo llevaría años. Tercero, la eficiencia muestral: tamaños de lote efectivos mayores, made viables por el paralelismo de datos, pueden estabilizar los gradientes y mejorar la generalización hasta un tamaño de lote crítico más allá del cual los rendimientos disminuyen.

El escalado, sin embargo, no es gratuito. La comunicación entre los trabajadores compite con el cálculo, y las barreras de sincronización pueden dejar dispositivos inactivos. La pregunta central de ingeniería del entrenamiento distribuido es cómo solapar, comprimir o eliminar esa comunicación preservando al mismo tiempo la dinámica de optimización del algoritmo original.

Paralelismo de datos

En el entrenamiento con paralelismo de datos síncrono, cada uno de los $ K $ trabajadores mantiene una copia idéntica de los parámetros $ \theta $ y procesa un microlote distinto de tamaño $ B $. Cada trabajador calcula un gradiente local $ g_k = \nabla_\theta L_k(\theta) $, los trabajadores intercambian gradientes mediante una operación all-reduce de modo que cada trabajador disponga del promedio $ \bar{g} = \frac{1}{K}\sum_{k=1}^{K} g_k $, y cada trabajador aplica localmente la misma actualización del optimizador:

$ {\displaystyle \theta \leftarrow \theta - \eta \, \bar{g}.} $

Como todas las réplicas comienzan idénticas y aplican la misma actualización al mismo gradiente promediado, permanecen idénticas bit a bit (salvo por el no determinismo en las reducciones). El tamaño de lote efectivo es $ K B $, lo que normalmente requiere un calentamiento de la tasa de aprendizaje y una regla de escalado lineal para mantener la estabilidad.

La primitiva all-reduce se implementa con un algoritmo de anillo o de árbol; el ring all-reduce mueve $ 2(K-1)/K $ veces el tamaño del gradiente por trabajador y es óptimo en ancho de banda sobre interconexiones uniformes. NCCL, Gloo y MPI proporcionan implementaciones afinadas.

Una alternativa más antigua, el servidor de parámetros, hace que un proceso central mantenga los parámetros canónicos y agregue las actualizaciones de los trabajadores. Las variantes asíncronas permiten que los trabajadores envíen gradientes obsoletos sin esperar, sacrificando calidad de optimización por rendimiento; este enfoque es mayormente histórico en el aprendizaje profundo denso, pero persiste en algunas cargas de recomendadores y de incrustaciones en las que la dispersión y el desequilibrio de carga lo favorecen.

Paralelismo de modelo

Cuando los parámetros o las activaciones de una sola capa exceden la memoria del dispositivo, la propia capa debe particionarse. En el paralelismo de tensores, una multiplicación de matrices $ Y = X W $ se divide a lo largo de la dimensión oculta: cada trabajador mantiene una losa de columnas $ W_k $, calcula $ Y_k = X W_k $, y las salidas parciales se concatenan o se suman según la capa. Megatron-LM popularizó este esquema para los bloques transformer al fragmentar las proyecciones de atención y de MLP de modo que cada bloque solo requiere dos all-reduces en los pasos hacia adelante y hacia atrás.

El paralelismo de tensores escala bien dentro de un nodo donde enlaces de alto ancho de banda como NVLink conectan un grupo pequeño de GPUs (típicamente de 4 a 8). Entre nodos, el coste de comunicación por paso crece rápidamente y suele combinarse con otras formas de paralelismo en lugar de utilizarse en solitario.

Paralelismo de canalización

El paralelismo de canalización particiona la red entre los trabajadores por profundidad: el trabajador $ k $ mantiene un bloque contiguo de capas y reenvía las activaciones al trabajador $ k+1 $. Ejecutar ingenuamente un microlote a la vez deja a la mayoría de los trabajadores inactivos (la "burbuja de canalización"). GPipe aborda esto dividiendo un minilote en $ M $ microlotes y solapando sus pasos hacia adelante; la fracción de burbuja es aproximadamente $ (P-1)/(M+P-1) $ para $ P $ etapas. Los planificadores PipeDream-Flush y 1F1B intercalan los pasos hacia adelante y hacia atrás de manera más agresiva para reducir la memoria de activaciones.

El paralelismo de canalización es sensible al equilibrio de carga: una sola etapa lenta detiene la canalización. Los profesionales perfilan el tiempo por etapa y ajustan la asignación de capas, a veces desplazando trabajo entre etapas o duplicando capas baratas.

Paralelismo híbrido y 3D

A escala de clúster, ningún eje único basta. La disposición común, a veces llamada paralelismo 3D, compone paralelismo de datos, de tensores y de canalización. El paralelismo de tensores maneja las capas más grandes dentro de un nodo, el paralelismo de canalización particiona el modelo entre nodos, y el paralelismo de datos replica toda esta canalización entre grupos de réplicas para obtener rendimiento. El paralelismo de secuencia, que fragmenta las activaciones a lo largo de la dimensión de secuencia, es un cuarto eje utilizado para contextos muy largos.

Elegir una partición es un problema de optimización discreta determinado por los presupuestos de memoria, la topología de interconexión y las formas de atención y de la red de propagación hacia adelante del modelo. Herramientas como Alpa, GSPMD y las rutinas de partición de DeepSpeed y Megatron exploran este espacio automáticamente.

ZeRO y estado de optimizador fragmentado

El Zero Redundancy Optimizer (ZeRO) reduce el coste de memoria del paralelismo de datos fragmentando lo que cada réplica almacena. ZeRO-1 particiona el estado del optimizador entre el grupo de paralelismo de datos; ZeRO-2 también particiona los gradientes; ZeRO-3 (también conocido como Fully Sharded Data Parallel, FSDP) particiona los propios parámetros, recolectándolos a demanda para cada capa. Como las primitivas gather y reduce-scatter se componen con el cálculo hacia atrás, la comunicación puede solaparse en gran medida con la aritmética, lo que produce un consumo de memoria cercano al del paralelismo de modelo con el modelo de programación más simple del paralelismo de datos.

Combinado con almacenamiento en precisión mixta y descarga a CPU o NVMe, ZeRO ha permitido entrenar modelos de un billón de parámetros en clústeres de GPU de escala comercial.

Comunicación y sincronización

El coste dominante en el entrenamiento distribuido es mover tensores. Varias técnicas lo reducen u ocultan:

• La compresión de gradientes, que incluye el SGD de 1 bit, signSGD y PowerSGD, reduce los bytes por all-reduce a costa de introducir sesgo o varianza que deben controlarse.
• El SGD local realiza $ H $ pasos en cada trabajador antes de promediar, reduciendo la frecuencia de comunicación al precio de una ligera deriva de optimización.
• El solapamiento entre cálculo y comunicación aprovecha que los gradientes de las capas iniciales están listos mientras las capas posteriores aún se están calculando; frameworks como PyTorch DDP y Horovod programan all-reduces asíncronos en buckets de modo que la red esté en uso durante el paso hacia atrás.
• Las colectivas conscientes de la topología prefieren los enlaces intra-nodo de alto ancho de banda para la reducción temprana y solo cruzan los enlaces lentos una vez.

Consideraciones prácticas y limitaciones

El entrenamiento distribuido amplifica cada pequeño fallo. Los rezagados, los errores de enlace y la corrupción silenciosa de datos se vuelven rutinarios con miles de aceleradores, por lo que la frecuencia de los puntos de control, los reinicios deterministas y la refragmentación elástica reciben una atención considerable de ingeniería. Los problemas numéricos como las inestabilidades por picos de pérdida son más frecuentes con tamaños de lote grandes y pueden requerir recorte de gradientes, tasas de aprendizaje más bajas o la recuperación desde un punto de control reciente.

La elección de la estrategia de paralelismo rara vez es portátil: una partición ajustada para una topología de clúster puede estar lejos de ser óptima en otra, y cambiar la forma del modelo puede requerir refragmentar el estado del optimizador. Las bibliotecas de comunicación, las sobrecargas de lanzamiento de kernels y la fragmentación de memoria por dispositivo se convierten todas en preocupaciones de primer orden. A pesar de la carga de ingeniería, el entrenamiento distribuido es la única ruta conocida hacia los modelos a escala de frontera actuales, y las mejoras en algoritmos colectivos, estrategias de fragmentación e interconexiones de aceleradores siguen ampliando la envolvente práctica.