El paralelismo tensorial es una técnica para distribuir el cálculo de una sola capa de red neuronal entre múltiples dispositivos de hardware. En este enfoque, las matrices de pesos de una capa se dividen en fragmentos, y cada dispositivo contiene una porción de los parámetros y realiza la porción correspondiente de la multiplicación de matrices. Los resultados parciales se combinan luego mediante operaciones de comunicación colectiva, como all-reduce, para producir la salida de la capa. Este método es una forma de paralelismo de modelo, distinto del paralelismo de datos, donde cada dispositivo tiene una copia completa del modelo y procesa diferentes muestras de datos. El paralelismo tensorial es esencial para entrenar y ejecutar inferencia en modelos muy grandes, como los modelos de lenguaje grandes, que superan la capacidad de memoria de un solo acelerador.
La motivación principal del paralelismo tensorial son las demandas de memoria y cálculo de los modelos modernos de aprendizaje profundo. Por ejemplo, un modelo con cientos de miles de millones de parámetros no cabe en la memoria de una sola GPU o TPU, que típicamente varía de 16 a 80 gigabytes. Al fragmentar las matrices de pesos, el paralelismo tensorial permite distribuir el modelo entre muchos dispositivos, habilitando el uso de modelos más grandes y tiempos de entrenamiento más rápidos. También reduce la huella de memoria por dispositivo y permite tamaños de lote más grandes. Sin embargo, el paralelismo tensorial introduce una sobrecarga de comunicación, ya que los dispositivos deben sincronizar los resultados parciales, lo que puede convertirse en un cuello de botella, especialmente en interconexiones más lentas.
El paralelismo tensorial se usa comúnmente junto con otras estrategias de paralelismo, como el paralelismo de pipeline y el paralelismo de datos, para escalar el entrenamiento a través de miles de dispositivos. Es un componente clave de la infraestructura de entrenamiento de muchos modelos de última generación, incluidos los desarrollados por OpenAI, Anthropic y Google DeepMind. La técnica también es relevante para la inferencia, donde permite servir modelos grandes con baja latencia al distribuir el cálculo entre múltiples aceleradores.
Antecedentes Históricos
El concepto de dividir operaciones de matrices entre múltiples procesadores se remonta a la investigación temprana en computación paralela. En las décadas de 1980 y 1990, los investigadores exploraron implementaciones paralelas de redes neuronales, pero las herramientas de hardware y software eran limitadas. La forma moderna del paralelismo tensorial surgió con el auge del aprendizaje profundo basado en GPU en la década de 2010. Uno de los primeros usos notables fue en el entrenamiento de grandes redes convolucionales, donde investigadores de la Universidad de Toronto y otras instituciones experimentaron con paralelismo de modelo. Sin embargo, fue la llegada de los modelos basados en transformadores, como la arquitectura Transformer introducida en 2017, la que hizo del paralelismo tensorial una técnica dominante. Las capas de atención y las redes feed-forward del Transformer están compuestas por grandes multiplicaciones de matrices, que son candidatas naturales para la fragmentación.
En 2019, investigadores de NVIDIA y otras organizaciones publicaron descripciones detalladas del paralelismo tensorial para entrenar grandes modelos de lenguaje. El marco Megatron-LM, desarrollado por NVIDIA, demostró un paralelismo tensorial eficiente para modelos transformadores, logrando un escalado casi lineal en múltiples GPUs. Este trabajo sentó las bases para sistemas de entrenamiento a gran escala posteriores, como GShard de Google y DeepSpeed de Microsoft, que incorporan el paralelismo tensorial como una característica central.
Cómo Funciona el Paralelismo Tensorial
En una capa típica de red neuronal, el paso hacia adelante calcula una multiplicación de matrices: Y = XW, donde X son las activaciones de entrada, W es la matriz de pesos y Y es la salida. En el paralelismo tensorial, la matriz de pesos W se particiona a lo largo de una o más dimensiones. Para una capa totalmente conectada, un enfoque común es dividir W por columnas (es decir, a lo largo de la dimensión de salida). Cada dispositivo contiene un subconjunto de las columnas de W y calcula una salida parcial. Las salidas parciales se combinan luego usando una operación all-reduce para producir la Y final. Esto se conoce como particionado paralelo por columnas. Alternativamente, W se puede dividir por filas (a lo largo de la dimensión de entrada), lo que requiere una operación all-gather en la entrada y un all-reduce en la salida.
Para modelos transformadores, el paralelismo tensorial se aplica tanto a la atención multi-cabeza como a las redes feed-forward. En la atención multi-cabeza, las matrices de pesos de consulta, clave y valor se fragmentan entre dispositivos, y las cabezas de atención se distribuyen. La proyección de salida se combina luego. En la red feed-forward, las dos capas lineales se particionan típicamente: la primera capa se divide por columnas, y la segunda capa se divide por filas, de modo que las activaciones intermedias se fragmentan y la salida se reduce. Este diseño minimiza la comunicación al requerir solo operaciones all-reduce en puntos específicos.
El patrón de comunicación es crítico. Cada operación all-reduce requiere intercambiar datos entre todos los dispositivos en el grupo de paralelismo tensorial. El volumen de comunicación es proporcional al tamaño de las activaciones, no al de los pesos, que a menudo es más pequeño. Sin embargo, para modelos muy grandes, la comunicación aún puede ser significativa. Para mitigar esto, el paralelismo tensorial se usa típicamente dentro de un solo nodo con interconexiones de alta velocidad, como NVLink de NVIDIA o Infinity Fabric de AMD, mientras que el paralelismo de datos se usa entre nodos.
Comparación con Otras Técnicas de Paralelismo
El paralelismo tensorial es una de varias estrategias de paralelismo de modelo. El paralelismo de pipeline, por ejemplo, divide el modelo por capas, con cada dispositivo responsable de un conjunto contiguo de capas. Esto reduce la comunicación en comparación con el paralelismo tensorial porque solo se intercambian las activaciones en los límites de las capas. Sin embargo, el paralelismo de pipeline puede sufrir tiempo inactivo debido a burbujas de pipeline. El paralelismo de datos, por otro lado, replica el modelo en cada dispositivo y procesa diferentes muestras de datos, requiriendo sincronización de gradientes después de cada paso. El paralelismo tensorial es complementario a ambos: se puede combinar con el paralelismo de pipeline para reducir la memoria por dispositivo y con el paralelismo de datos para aumentar el rendimiento.
Otra técnica relacionada es el paralelismo de secuencia, que fragmenta la dimensión de secuencia en modelos transformadores. Esto se usa a menudo junto con el paralelismo tensorial para reducir aún más el uso de memoria. El paralelismo experto, utilizado en modelos de mezcla de expertos, fragmenta los módulos expertos entre dispositivos, lo que es una forma de paralelismo de modelo pero no estrictamente paralelismo tensorial.
Aplicaciones y Casos de Uso
El paralelismo tensorial se usa ampliamente en el entrenamiento e inferencia de grandes modelos de lenguaje. Por ejemplo, el modelo GPT-3, desarrollado por OpenAI, se entrenó usando una combinación de paralelismo de modelo, incluido el paralelismo tensorial, a través de miles de GPUs de NVIDIA. De manera similar, los modelos Claude de Anthropic y los modelos Gemini de Google DeepMind dependen del paralelismo tensorial para escalar a cientos de miles de millones de parámetros. En inferencia, el paralelismo tensorial permite servir modelos como GPT-4 con baja latencia al distribuir el cálculo entre múltiples GPUs, como lo hacen el servidor de inferencia Triton de NVIDIA y la generación de texto de inferencia de Hugging Face.
Más allá de los modelos de lenguaje, el paralelismo tensorial también se usa en otros dominios, como la visión por computadora y la computación científica, donde los modelos grandes o los tamaños de lote grandes requieren cálculo distribuido. Por ejemplo, entrenar grandes transformadores de visión (ViTs) en imágenes de alta resolución puede beneficiarse del paralelismo tensorial.
Desafíos y Limitaciones
A pesar de sus beneficios, el paralelismo tensorial tiene varios desafíos. El problema principal es la sobrecarga de comunicación. Las operaciones all-reduce requieren interconexiones de alta velocidad y baja latencia. En sistemas con redes más lentas, como Ethernet, el paralelismo tensorial puede volverse ineficiente. Por lo tanto, se usa típicamente dentro de un solo servidor o un clúster estrechamente acoplado. Otro desafío es el equilibrio de carga: una fragmentación desigual puede llevar a que algunos dispositivos estén subutilizados. Además, el paralelismo tensorial requiere una gestión cuidadosa de la memoria, ya que cada dispositivo debe almacenar su fragmento de los pesos y las activaciones intermedias. Finalmente, implementar el paralelismo tensorial manualmente es complejo; requiere modificar el código del modelo y manejar primitivas de comunicación. Afortunadamente, marcos como PyTorch y TensorFlow proporcionan soporte integrado para el paralelismo tensorial, abstrayendo gran parte de la complejidad.
Soporte de Software y Marcos
Varios marcos y bibliotecas de aprendizaje profundo admiten el paralelismo tensorial. PyTorch ofrece el módulo torch.distributed y la API tensor_parallel para fragmentar tensores. TensorFlow proporciona tf.distribute con soporte para paralelismo de modelo. Bibliotecas especializadas como Megatron-LM, DeepSpeed y Transformers de Hugging Face han implementado paralelismo tensorial para modelos transformadores. Estas herramientas permiten a investigadores e ingenieros aplicar paralelismo tensorial con cambios mínimos en el código, a menudo simplemente especificando el número de dispositivos y la estrategia de fragmentación.
Direcciones Futuras
A medida que los modelos continúan creciendo, el paralelismo tensorial seguirá siendo una técnica crítica. Los desarrollos futuros pueden incluir algoritmos de comunicación más eficientes, como all-reduce jerárquico, y una mejor integración con hardware heterogéneo, incluidas GPUs de AMD, aceleradores de Intel y chips personalizados como AWS Trainium y las LPUs de Groq. Además, el descubrimiento automatizado de paralelismo, donde el sistema determina la estrategia de fragmentación óptima, es un área activa de investigación. Con la creciente demanda de IA a gran escala, el paralelismo tensorial continuará evolucionando para enfrentar los desafíos de escala y eficiencia.