Paralelismo de modelo

El paralelismo de modelos es una técnica de computación distribuida utilizada en aprendizaje automático y aprendizaje profundo para entrenar o realizar inferencia en redes neuronales grandes que no caben en la memoria de un solo dispositivo, como una GPU o TPU. A diferencia del paralelismo de datos, donde cada dispositivo tiene una copia completa del modelo y procesa diferentes lotes de datos, el paralelismo de modelos particiona el propio modelo, típicamente asignando diferentes capas o subconjuntos de capas a diferentes dispositivos. Este enfoque es esencial para escalar modelos como los Transformer (architecture)-basados en grandes modelos de lenguaje, que pueden tener miles de millones o billones de parámetros.

Descripción general

En el paralelismo de modelos, la red neuronal se divide a lo largo de su profundidad (por capas) o, en formas más avanzadas, a lo largo de otras dimensiones como el ancho o incluso operaciones tensoriales individuales. La forma más directa es la partición por capas, donde capas consecutivas se colocan en diferentes dispositivos. Durante la propagación hacia adelante, las activaciones fluyen de un dispositivo al siguiente; durante la retropropagación, los gradientes fluyen en dirección inversa. Esto crea una ejecución tipo pipeline, que puede introducir tiempo de inactividad (burbujas) si no se programa cuidadosamente. El paralelismo de modelos a menudo se contrasta con el paralelismo de datos, pero ambos pueden combinarse en enfoques híbridos, como el paralelismo 3D utilizado en el entrenamiento de modelos grandes.

Motivación

Las redes neuronales modernas, particularmente los grandes modelos de lenguaje, han crecido exponencialmente en tamaño. Por ejemplo, GPT-3 (2020) tiene 175 mil millones de parámetros, requiriendo más de 350 GB de memoria solo para los pesos en precisión de 32 bits, superando con creces la capacidad de 80 GB de GPUs de gama alta como la A100. Incluso con optimizaciones de memoria como precisión mixta y checkpointing de gradientes, el entrenamiento en un solo dispositivo es inviable. El paralelismo de modelos permite a los investigadores entrenar modelos que son órdenes de magnitud más grandes de lo que un solo dispositivo puede contener, distribuyendo el modelo a través de un clúster de dispositivos conectados por interconexiones de alta velocidad.

Tipos de paralelismo de modelos

Paralelismo por capas (pipeline)

Esta es la forma más simple, donde el modelo se divide en etapas secuenciales, cada una asignada a un dispositivo diferente. Por ejemplo, una red de 100 capas podría dividirse en 10 etapas de 10 capas cada una, con cada etapa en una GPU separada. Durante el pase hacia adelante, los datos fluyen a través de la etapa 1, luego la etapa 2, y así sucesivamente. El principal desafío es el equilibrio de carga y la reducción de burbujas en el pipeline. Técnicas como GPipe y PipeDream introducen micro-lotes y programación para mejorar la utilización.

Paralelismo tensorial

El paralelismo tensorial divide operaciones individuales (por ejemplo, multiplicaciones de matrices) entre múltiples dispositivos. Por ejemplo, en el mecanismo de atención de un transformer, las proyecciones de consulta, clave y valor pueden particionarse entre GPUs, con resultados combinados mediante operaciones de all-reduce. Este enfoque se utiliza en Megatron-LM y es efectivo para reducir la memoria por dispositivo mientras mantiene una alta eficiencia computacional, pero requiere comunicación de alto ancho de banda.

Paralelismo de expertos

Utilizado en modelos de mezcla de expertos (MoE), el paralelismo de expertos coloca diferentes redes expertas en diferentes dispositivos, mientras que el enrutador (red de compuerta) distribuye tokens a los expertos apropiados. Esta es una forma de paralelismo de modelos que permite recuentos masivos de parámetros sin un aumento proporcional en el cómputo, ya que solo un subconjunto de expertos se activa por token.

Desafíos de implementación

El paralelismo de modelos introduce varios desafíos:

• Sobrecarga de comunicación: Los dispositivos deben intercambiar activaciones y gradientes, lo que puede convertirse en un cuello de botella, especialmente con interconexiones lentas.
• Desequilibrio de carga: Tamaños de capas o requisitos de cómputo desiguales pueden causar que algunos dispositivos estén inactivos mientras otros trabajan.
• Burbujas en el pipeline: En el paralelismo por pipeline, los dispositivos pueden esperar datos de etapas anteriores, reduciendo la utilización.
• Fragmentación de memoria: Particionar modelos puede llevar a un uso desigual de la memoria, requiriendo una colocación y programación cuidadosas.
• Tolerancia a fallos: Con muchos dispositivos, la probabilidad de fallos aumenta, lo que requiere mecanismos de checkpointing y recuperación.

Comparación con el paralelismo de datos

En el paralelismo de datos, cada dispositivo tiene una copia completa del modelo y procesa un mini-lote diferente de datos. Los gradientes se promedian entre dispositivos después de cada paso. Esto es simple y escala bien para modelos que caben en la memoria de un solo dispositivo. El paralelismo de modelos, por otro lado, es necesario cuando el modelo en sí es demasiado grande. Sin embargo, el paralelismo de modelos a menudo tiene mayores costos de comunicación y menor eficiencia debido a dependencias seriales. En la práctica, el entrenamiento a gran escala utiliza ambos: paralelismo de datos entre nodos y paralelismo de modelos dentro de un nodo.

Aplicaciones

El paralelismo de modelos es crítico para entrenar grandes modelos de lenguaje como GPT-4, PaLM y LLaMA, así como para servirlos en producción. Empresas como OpenAI, Google (a través de Google Cloud) y Anthropic dependen del paralelismo de modelos para entrenar y desplegar sus modelos. También se utiliza en otros dominios, como visión por computadora (por ejemplo, CNNs 3D) y computación científica, donde los modelos son demasiado grandes para un solo acelerador.

Soporte de software

Varios marcos de trabajo proporcionan soporte integrado para el paralelismo de modelos:

• PyTorch: Ofrece torch.distributed con paralelismo por pipeline (por ejemplo, torch.distributed.pipeline.sync.Pipe) y paralelismo tensorial a través de bibliotecas tensor_parallel.
• TensorFlow: Proporciona tf.distribute con estrategias como TPUStrategy y MultiWorkerMirroredStrategy que pueden incorporar paralelismo de modelos.
• Megatron-LM: Una biblioteca especializada de NVIDIA para paralelismo tensorial y por pipeline.
• DeepSpeed: De Microsoft, ofrece ZeRO (Optimizador de Redundancia Cero), que es una forma de paralelismo de modelos que particiona estados de optimizador, gradientes y parámetros.
• JAX: Con jax.sharding, permite la colocación explícita de dispositivos y la ejecución paralela.

Consideraciones de hardware

El paralelismo de modelos depende en gran medida de la comunicación entre dispositivos de alta velocidad. Dentro de un solo nodo, NVLink y PCIe proporcionan conexiones rápidas; entre nodos, se utilizan InfiniBand o Ethernet de alta velocidad. Hardware especializado como GPUs de NVIDIA y TPUs de Google a menudo incluyen interconexiones dedicadas (por ejemplo, NVSwitch, enlaces TPU) para facilitar el paralelismo de modelos. Empresas como Cerebras y SambaNova han diseñado sistemas con gran memoria en chip para reducir la necesidad de paralelismo de modelos, pero para escala extrema, los enfoques distribuidos siguen siendo necesarios.

Avances recientes

La investigación reciente se ha centrado en mejorar la eficiencia del paralelismo de modelos. Se han desarrollado técnicas como el paralelismo de secuencias (dividir la dimensión de secuencia) y el paralelismo de contexto (para secuencias largas). Algoritmos automáticos de particionado de modelos, como los de Alpa y FlexFlow, optimizan la colocación de operaciones entre dispositivos. Además, el paralelismo híbrido que combina paralelismo de datos, tensorial y por pipeline es ahora estándar en el entrenamiento a gran escala, como ejemplifica el modelo Megatron-Turing NLG (530 mil millones de parámetros) entrenado por NVIDIA y Microsoft.

Limitaciones y direcciones futuras

A pesar de su poder, el paralelismo de modelos tiene limitaciones. La sobrecarga de comunicación puede dominar para modelos pequeños o redes lentas. La complejidad de particionar y programar aumenta con el tamaño del modelo y la heterogeneidad del clúster. Las direcciones futuras incluyen un auto-paralelismo más inteligente, mejor compresión de comunicación y arquitecturas novedosas que sean inherentemente más paralelizables. Con el auge de la IA generativa y modelos cada vez más grandes, el paralelismo de modelos seguirá siendo una piedra angular del aprendizaje profundo escalable.

Véase también

• Paralelismo de datos
• Paralelismo por pipeline
• Paralelismo tensorial
• Mezcla de expertos
• Entrenamiento distribuido
• Gran modelo de lenguaje
• Clúster de GPUs

Referencias

• Narayanan, D., et al. (2019). PipeDream: Generalized Pipeline Parallelism for DNN Training.
• Shoeybi, M., et al. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.
• Rajbhandari, S., et al. (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.
• Huang, Y., et al. (2019). GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism.