LocMoE: Un MoE de bajo costo para el entrenamiento de modelos de lenguaje grandes

El modelo Mixtures-of-Experts (MoE) es un método de aprendizaje distribuido e integrado ampliamente utilizado para modelos de lenguaje grandes (LLM), valorado por su capacidad para esparcir y expandir modelos de manera eficiente. Sin embargo, el rendimiento de MoE está limitado por el desequilibrio de carga y la alta latencia de la comunicación All-to-All, junto con un cálculo relativamente redundante debido a la gran capacidad de expertos. El desequilibrio de carga puede resultar de políticas de enrutamiento existentes que tienden a seleccionar consistentemente ciertos expertos. La comunicación frecuente entre nodos en el procedimiento All-to-All también prolonga significativamente el tiempo de entrenamiento. Para aliviar estos problemas de rendimiento, proponemos una estrategia de enrutamiento novedosa que combina equilibrio de carga y localidad al convertir parte de la comunicación inter-nodo en comunicación intra-nodo. Cabe destacar que elucidamos que existe un umbral mínimo para la capacidad del experto, calculado mediante la desviación angular máxima entre los pesos de activación de los expertos y los tokens asignados. Implementamos estas modificaciones en el modelo PanGu-Σ basado en el framework MindSpore con enrutamiento multinivel y realizamos experimentos en clusters Ascend. Los resultados experimentales demuestran que el LocMoE propuesto reduce el tiempo de entrenamiento por época entre un 12.68% y 22.24% en comparación con enrutadores clásicos, como hash router y switch router, sin afectar la precisión del modelo.