En los últimos años, los grandes modelos de lenguaje (LLM) basados en la arquitectura Transformer han generado un amplio interés gracias a su rendimiento sobresaliente. Los LLM industriales deben procesar secuencias largas para ofrecer servicios de alta calidad. Sin embargo, el consumo de memoria crece de forma cuadrática con la longitud de la secuencia, lo que limita la capacidad de escalado del entrenamiento con secuencias largas. Los métodos paralelos existentes generan tensores redundantes durante la ejecución, dejando espacio para optimizar la memoria; además, el paralelismo de tensor (TP) no puede lograr una superposición efectiva del cálculo y la comunicación. Para abordar estos problemas, este trabajo propone un método paralelo universal: paralelismo de tensor con eficiencia de memoria (METP), diseñado específicamente para las unidades de cálculo centrales en el entrenamiento de Transformers (es decir, dos multiplicaciones de matrices consecutivas y las posibles funciones entre ellas O=f(AB)C). METP distribuye las subtareas del cálculo de O en múltiples dispositivos, usando comunicación punto a punto (send/recv) en lugar de comunicación colectiva para intercambiar submatrices y completar los cálculos, evitando la generación de tensores redundantes. Mediante la técnica de doble buffer, se logra una superposición profunda entre cálculo y comunicación, además de proponer condiciones teóricas para la superposición completa que guían el entrenamiento de Transformers en secuencias largas. El análisis teórico muestra que, para un grado de paralelismo p, el costo de memoria de METP sin usar FlashAttention para calcular la atención es O(1/p³); al usar FlashAttention para calcular la autoatención multi-cabeza, METP ahorra al menos un 41.7% de memoria en comparación con TP. Los experimentos demuestran que, bajo una configuración con 8 GPUs A100, METP puede aumentar la longitud de la secuencia de 2.38 a 2.99 veces en comparación con otros métodos.

aprendizaje distribuido;modelos de lenguaje a gran escala;secuencias largas;sistemas de aprendizaje automático;eficiencia de memoria;paralelismo de tensor