Fast MoE Inference via Predictive Prefetching and Expert Replication

Resumen Técnico: Inferencia Rápida de MoE Mediante Prefetching Predictivo y Replicación de Expertos

Enunciado del Problema
La arquitectura de Mezcla de Expertos (MoE) se ha convertido en un pilar fundamental para escalar los Modelos de Lenguaje Grandes (LLM), permitiendo conteos masivos de parámetros sin aumentos proporcionales en la sobrecarga computacional. Sin embargo, la inferencia de MoE enfrenta ineficiencias críticas que obstaculizan su implementación práctica:

1. Utilización Subóptima de la GPU: Debido a la dispersión inherente de la activación de expertos, solo un pequeño subconjunto de expertos está activo en cualquier momento dado, dejando la mayoría de los recursos de la GPU inactivos.
2. Desequilibrio de Carga y Latencia: El enrutamiento dinámico a menudo resulta en una distribución de carga desigual, donde múltiples tokens hacen cola para el mismo experto físico, causando estancamiento de tokens y aumentando la latencia.
3. Sobrecarga de Enrutamiento: Los mecanismos de enrutamiento tradicionales basados en softmax introducen complejidad computacional y sobrecarga de comunicación (enrutamiento de todos a todos) que no escalan favorablemente con el aumento del número de expertos.

Las técnicas de optimización existentes, como la destilación o las mejoras a nivel de kernel, a menudo ofrecen mejoras limitadas o comprometen la precisión del modelo.

Metodología: MoE-MPMC
Los autores proponen MoE-MPMC (Mezcla de Expertos con Múltiples Productores y Múltiples Consumidores), un marco que aborda estas ineficiencias mediante prefetching predictivo y replicación dinámica de expertos. El sistema opera utilizando dos hilos concurrentes:

1. Hilo de Inferencia:

   • Recupera una tabla hash precalculada que contiene asignaciones de expertos y conteos de replicación para el lote actual.
   • Carga dinámicamente en la memoria de la GPU solo los expertos predichos como activos, descargando los expertos no utilizados hacia la memoria de la CPU.
   • Ejecuta la inferencia utilizando una capa de MoE modificada donde los expertos activos están físicamente replicados en la GPU. Esto permite que múltiples tokens que requieren el mismo experto sean procesados concurrentemente por réplicas físicas distintas, eliminando los retrasos de colas.

2. Hilo de Construcción de Hash:

   • Se ejecuta en paralelo con el hilo de inferencia, procesando el siguiente lote de tokens.
   • Utiliza una Unidad Recurrente Simple (SRU) con atención dispersa para predecir qué expertos se activarán para el lote próximo.
   • Reemplaza el predictor basado en LSTM utilizado en trabajos anteriores (SiDA-MoE) con SRU para lograr velocidades de entrenamiento e inferencia más rápidas debido a la reducción de interdependencias y mecanismos de puerta más simples.
   • Calcula los conteos de replicación para cada experto predicho basándose en la demanda de tokens y actualiza la tabla hash, la cual luego se pone en cola para el hilo de inferencia.

Contribuciones Clave
El artículo describe tres contribuciones principales:

1. Predicción Basada en SRU: La introducción de la Unidad Recurrente Simple (SRU) para predecir rápidamente las selecciones de expertos para cada token. Esto reemplaza a los predictores LSTM más lentos, acelerando significativamente la fase de predicción mientras se mantiene la precisión.
2. Replicación Física Profunda: Una estrategia para replicar físicamente los expertos predichos en la GPU basándose en la demanda de tokens. Esto asegura que casi ningún token experimente retrasos de espera durante el procesamiento por lotes, maximizando así el paralelismo y la utilización de la GPU.
3. Límite de Capacidad: Un mecanismo práctico para limitar los conteos de replicación de expertos según las restricciones de memoria de la GPU predefinidas. Esto asegura que la replicación se optimice dentro de los límites del hardware para lograr una utilización de recursos casi óptima sin causar desbordamiento.

Resultados Experimentales
Los autores evaluaron MoE-MPMC en modelos SwitchTransformers a gran escala (Switch-base-8, 64, 128 y 256) utilizando conjuntos de datos como SST2 y MRPC. Los hallazgos clave incluyen:

• Utilización de la GPU: El método logra una utilización casi completa de la GPU (hasta ~100% en algunas configuraciones), una mejora significativa sobre las líneas base de SwitchTransformers (que mostraron una utilización tan baja como 1.3% para 128 expertos) y el enfoque SiDA-MoE.
• Velocidad de Inferencia: El enfoque produce hasta una mejora de 3x en la velocidad de inferencia en comparación con las arquitecturas base.
• Precisión: El método preserva aproximadamente el 90-95% del rendimiento de las arquitecturas base, manteniendo una precisión mejor que SiDA-MoE debido a las capacidades de predicción superiores de SRU sobre LSTM.
• Rendimiento (Throughput): El rendimiento aumenta significativamente a medida que aumenta el número de expertos, impulsado por la capacidad de procesar múltiples tokens en paralelo mediante la replicación de expertos.

Importancia y Afirmaciones
El artículo afirma que MoE-MPMC representa un paso significativo hacia la implementación práctica y económicamente viable de modelos de lenguaje a gran escala basados en MoE. Al cambiar del equilibrio de carga reactivo a una estrategia proactiva de prefetching predictivo y replicación, el método:

• Mitiga los cuellos de botella en el acceso a la memoria y el estancamiento de tokens.
• Optimiza la utilización de recursos de la GPU sin alterar el gráfico de computación subyacente de MoE.
• Demuestra que predecir selecciones de expertos es factible y efectivo cuando el ajuste fino y el preentrenamiento ocurren concurrentemente con tareas de predicción de tokens.

Los autores posicionan este trabajo como un avance sobre los métodos reactivos (como SiDA-MoE) al anticipar demandas futuras y ajustar dinámicamente la asignación de recursos, abordando así los errores típicos de memoria de GPU inactiva y desequilibrios de carga en arquitecturas dispersas.