Zeppelin: Balance de Cargas de Trabajo de Longitud Variable en el Entrenamiento de Modelos Grandes en Paralelo de Datos

El entrenamiento de modelos de lenguaje grandes (LLMs) con longitudes de secuencia cada vez más largas y variables introduce seriosos desafíos de desequilibrio de carga en el entrenamiento a gran escala en paralelo de datos. Los marcos recientes intentan mitigar estos problemas a través de la reorganización de datos o estrategias paralelas híbridas. Sin embargo, a menudo pasan por alto cómo los costos computacionales y de comunicación escalan con la longitud de la secuencia, lo que resulta en un rendimiento subóptimo. Identificamos tres desafíos críticos: (1) relaciones de computación a comunicación variables a través de secuencias de diferentes longitudes en atención distribuida, (2) desajuste entre la afinidad estática NIC-GPU y las cargas de trabajo paralelas dinámicas, y (3) distintas estrategias de partición óptimas requeridas para atención cuadrática frente a componentes lineales. Para abordar estos desafíos, presentamos Zeppelin, un nuevo sistema de entrenamiento que integra tres técnicas clave: (1) un método jerárquico de partición de secuencias para el módulo de atención que reduce la sobrecarga de comunicación y equilibra la computación, apoyado por un motor de atención eficiente que aplica estrategias paralelas divergentes; (2) una capa de enrutamiento que orquesta transferencias entre nodos para utilizar al máximo el ancho de banda de la NIC; y (3) una capa de reconfiguración que transforma los diseños de secuencia entre los módulos de atención y lineales, asegurando alta eficiencia computacional en ambos. Evaluaciones completas a través de diversas configuraciones muestran que Zeppelin ofrece un promedio de 2.80x de aceleración sobre los métodos de vanguardia.