En los últimos años, los grandes modelos de lenguaje (LLM) basados en la arquitectura Transformer han suscitado un gran interés debido a su rendimiento excepcional. Los LLM industriales deben procesar secuencias de entrada largas para proporcionar un servicio de calidad. Sin embargo, el consumo de memoria aumenta cuadráticamente con la longitud de la secuencia, lo que limita la capacidad de extensión del entrenamiento en secuencias largas. Los métodos paralelos actuales generan tensores redundantes durante la ejecución, lo que proporciona espacio para la optimización de la memoria; al mismo tiempo, el paralelismo tensorial (TP) no puede lograr una superposición eficaz de cálculos y comunicaciones. Para resolver estos problemas, este artículo propone un método paralelo general: paralelismo tensorial de alta eficiencia de memoria (METP), diseñado específicamente para la unidad de cálculo principal del entrenamiento Transformer (es decir, dos multiplicaciones matriciales consecutivas y posibles operaciones funcionales entre ellas O=f (AB)C). METP asigna las sub tareas de cálculo de O en varios dispositivos, utilizando una comunicación punto a punto (envío/recepción) en lugar de un intercambio de comunicaciones colectivas para completar el cálculo, evitando la generación de tensores redundantes. La técnica de doble búfer permite una superposición profunda de cálculos y comunicaciones, y se proponen condiciones teóricas de superposición completa para guiar el entrenamiento de Transformer en secuencias largas. El análisis teórico muestra: cuando la paralelización es p, el METP en términos de costo de memoria en caso de ausencia de cálculo de atención FlashAttention es O(1/p3); al utilizar FlashAttention para calcular la autoatención con múltiples cabezas, se puede ahorrar al menos un 41,7% de memoria en comparación con TP. Los experimentos muestran que, basándose en 8 GPU A100, METP puede aumentar la longitud de la secuencia de 2,38 a 2,99 veces en comparación con otros métodos.
Keywords
Aprendizaje distribuido; Gran modelo de lenguaje (LLM); Secuencia larga; Sistema de aprendizaje automático; Eficiencia de la memoria; Paralelismo tensorial