Antecedentes y motivación
Entrenar redes neuronales grandes, especialmente modelos de lenguaje de gran tamaño y otros modelos de aprendizaje profundo, requiere enormes recursos computacionales. A medida que los tamaños de los modelos crecen hasta miles de millones o billones de parámetros, el entrenamiento en un solo dispositivo se vuelve inviable debido a los límites de memoria y cómputo. El paralelismo de datos replica el modelo completo en cada dispositivo y divide los datos, pero falla cuando el modelo excede la memoria de un solo dispositivo. El paralelismo de modelos divide el propio modelo entre dispositivos, pero la colocación ingenua por capas conduce a una subutilización severa: solo un dispositivo computa a la vez mientras los demás esperan, lo que resulta en un bajo rendimiento y un alto tiempo de inactividad.
El paralelismo de canalizaciones aborda esto dividiendo el modelo en etapas (típicamente grupos contiguos de capas) y procesando múltiples micro-lotes de forma concurrente. La idea clave es superponer el cómputo entre etapas: mientras la etapa 1 computa en el micro-lote 2, la etapa 2 puede computar en el micro-lote 1, y así sucesivamente. Esto se asemeja a una línea de ensamblaje, donde cada etapa procesa continuamente los datos entrantes y pasa los resultados a la siguiente etapa.
Conceptos fundamentales
La unidad fundamental de trabajo en el paralelismo de canalizaciones es el micro-lote. Un lote de entrenamiento grande se divide en micro-lotes más pequeños que fluyen secuencialmente a través de las etapas de la canalización. Cada etapa realiza pasadas hacia adelante y hacia atrás en sus capas asignadas. La canalización se "llena" alimentando micro-lotes uno tras otro, y después de un período inicial de calentamiento, todas las etapas pueden estar ocupadas simultáneamente, logrando un alto rendimiento.
La eficiencia de una canalización a menudo se mide por la relación de burbuja, que es la fracción de tiempo que las etapas están inactivas. Para una canalización con p etapas y m micro-lotes, el rendimiento ideal se logra cuando m es mucho mayor que p, minimizando la sobrecarga relativa de la burbuja. Sin embargo, aumentar m también aumenta el uso de memoria para las activaciones, creando una compensación.
Esquemas principales
Se han propuesto varios esquemas de paralelismo de canalizaciones, que difieren en cómo programan las pasadas hacia adelante y hacia atrás y cómo manejan las actualizaciones de gradientes.
GPipe
Google introdujo GPipe en 2019. Divide el modelo en etapas y utiliza un programa simple: cada micro-lote realiza pasadas hacia adelante a través de todas las etapas, luego pasadas hacia atrás a través de todas las etapas en orden inverso. GPipe utiliza actualizaciones de gradientes síncronas, lo que significa que todas las etapas esperan a que todos los micro-lotes se completen antes de actualizar los pesos. Esto garantiza gradientes consistentes pero introduce una burbuja que crece con el número de etapas. GPipe también requiere o bien la recomputación de activaciones o bien su almacenamiento, intercambiando cómputo por memoria.
PipeDream
PipeDream, de Microsoft Research (2019), utiliza un programa asíncrono donde cada etapa procesa micro-lotes en un orden round-robin, realizando pasadas hacia adelante y hacia atrás tan pronto como sea posible. Esto reduce la burbuja pero introduce desactualización de pesos: diferentes etapas pueden usar diferentes versiones de los pesos, lo que puede perjudicar la convergencia. PipeDream también utiliza una técnica llamada "weight stashing" para mantener múltiples versiones de los pesos para diferentes micro-lotes, aumentando la sobrecarga de memoria.
PipeDream-2BW y variantes
PipeDream-2BW (2020) mejora PipeDream utilizando dos buffers de pesos (uno para la pasada hacia adelante, otro para la hacia atrás) para reducir la memoria y mejorar la convergencia. Variantes como V-Pipe y PipeMare exploran diferentes estrategias de sincronización.
Megatron-LM y programación intercalada
NVIDIA's Megatron-LM (2019) combinó el paralelismo tensorial (dividiendo capas individuales) con el paralelismo de canalizaciones. Más tarde, Megatron-2 (2020) introdujo un programa intercalado que divide el modelo en más etapas que dispositivos, permitiendo que cada dispositivo maneje múltiples etapas. Esto reduce el tamaño de la burbuja aumentando el número de micro-lotes en vuelo, a costa de más comunicación.
Quimera y otros esquemas
Quimera (2021) utiliza una canalización bidireccional para reducir aún más las burbujas. Otros esquemas como PTD-P (2021) y ZeroBubble (2023) exploran técnicas más sofisticadas para eliminar o reducir las burbujas de la canalización.
Consideraciones matemáticas y prácticas
El paralelismo de canalizaciones a menudo se combina con otras estrategias de paralelismo. El paralelismo tensorial divide capas individuales entre dispositivos, mientras que el paralelismo de canalizaciones divide el modelo verticalmente. El paralelismo de datos replica la canalización entre grupos de dispositivos. Este enfoque híbrido es estándar en el entrenamiento de modelos grandes, como se ve en sistemas como Megatron-Turing NLG y OpenAI's GPT-4.
La elección de la profundidad de la canalización (número de etapas) y el tamaño del micro-lote afecta el rendimiento. Canalizaciones más profundas reducen la memoria por etapa pero aumentan la comunicación y la sobrecarga de burbujas. Micro-lotes más grandes mejoran la utilización pero consumen más memoria. La configuración óptima depende del tamaño del modelo, la memoria del dispositivo y el ancho de banda de interconexión.
Aplicaciones e impacto
El paralelismo de canalizaciones ha sido fundamental para entrenar algunos de los modelos de IA más grandes. Por ejemplo, GPT-3 (175 mil millones de parámetros) se entrenó utilizando una combinación de paralelismo de modelos y de canalizaciones. Los modelos Claude de Anthropic y PaLM de Google también dependen de técnicas similares. El enfoque es compatible con los principales marcos de trabajo, incluidos PyTorch (a través del módulo torch.distributed.pipeline), TensorFlow (vía tf.distribute) y JAX (con pjit).
Desafíos y limitaciones
El paralelismo de canalizaciones introduce varios desafíos. La sobrecarga de comunicación surge al enviar activaciones y gradientes entre etapas, lo que puede convertirse en un cuello de botella en interconexiones lentas. El desequilibrio de carga ocurre cuando las etapas tienen tiempos de cómputo desiguales, reduciendo la eficiencia. La presión de memoria por almacenar activaciones para la pasada hacia atrás puede mitigarse mediante la recomputación de activaciones, pero a costa de cómputo adicional. Además, los esquemas asíncronos pueden sufrir problemas de convergencia debido a pesos desactualizados, mientras que los esquemas síncronos incurren en burbujas.
Direcciones futuras
La investigación continúa mejorando el paralelismo de canalizaciones, centrándose en reducir las burbujas, minimizar la memoria y adaptarse a hardware heterogéneo. Técnicas como la partición automática de etapas, el equilibrio dinámico de carga y la integración con modelos mezcla de expertos son áreas activas. El auge de modelos extremadamente grandes y la necesidad de un entrenamiento eficiente en clústeres en la nube mantendrán probablemente el paralelismo de canalizaciones como un componente clave de los sistemas de entrenamiento distribuido.