Tensor Parallelism (también conocido como "intra-layer parallelism" o "data parallelism dentro de una capa") es una estrategia de paralelización de modelos de Machine Learning, especialmente en redes neuronales profundas (DNNs) y Large Language Models (LLMs), donde las operaciones de tensores dentro de una misma capa se distribuyen a través de múltiples dispositivos (GPUs, TPUs). En lugar de replicar el modelo completo o dividir los datos de entrada, Tensor Parallelism divide las matrices de pesos y/o activaciones de una capa en fragmentos, asignando cada fragmento a un dispositivo diferente. Cada dispositivo ejecuta una parte de la operación matricial, y los resultados parciales se combinan (por ejemplo, mediante una reducción o una concatenación) para producir el resultado final de la capa. Esto permite entrenar o inferir modelos que son demasiado grandes para caber en la memoria de un solo dispositivo.
En el mundo real, Tensor Parallelism es fundamental para el entrenamiento y la inferencia de LLMs masivos. Frameworks como DeepSpeed de Microsoft y Megatron-LM de NVIDIA son pioneros en la implementación de esta técnica. DeepSpeed, por ejemplo, ofrece varias estrategias de paralelismo, incluyendo Tensor Parallelism, para escalar el entrenamiento de modelos con miles de millones de parámetros. Megatron-LM de NVIDIA es un ejemplo prominente que utiliza Tensor Parallelism para distribuir las operaciones de multiplicación de matrices en capas de Transformers, permitiendo entrenar modelos como GPT-3 y sus sucesores. Otros sistemas como JAX/XLA de Google también proporcionan primitivas para implementar Tensor Parallelism de manera eficiente en TPUs y GPUs.
Para un Arquitecto de Sistemas, Tensor Parallelism es una técnica crucial para diseñar infraestructuras de Machine Learning escalables y eficientes. Permite entrenar y desplegar modelos que de otro modo serían inviables debido a las limitaciones de memoria de un solo dispositivo. Sin embargo, introduce complejidades significativas: requiere una coordinación de comunicación de alta velocidad entre los dispositivos (por ejemplo, a través de NVLink o InfiniBand) para la agregación de resultados parciales, lo que puede convertirse en un cuello de botella si no se gestiona correctamente. La elección de la estrategia de división (por filas, columnas, o una combinación) y la topología de comunicación son decisiones de diseño críticas que impactan directamente el rendimiento y la eficiencia. Un arquitecto debe sopesar el costo de la comunicación entre dispositivos frente al beneficio de la mayor capacidad de modelo, y a menudo combinar Tensor Parallelism con otras estrategias como Pipeline Parallelism o Data Parallelism para optimizar el uso de recursos y el throughput.