LRMoE recorta el paso hacia atrás de MoE y acelera el entrenamiento entre un 21 y un 45%

Un equipo de autores del ámbito académico ha propuesto LRMoE, un sistema para entrenar modelos MoE en nodos donde la memoria de la GPU no alcanza. MoE (Mixture-of-Experts, «mezcla de expertos») es una arquitectura en la que cada capa se compone de decenas o cientos de pequeñas subredes, los expertos, y un enrutador activa solo unas pocas para cada token. Los autores llaman «par» a cada una de estas elecciones «token–experto».

La cifra principal del trabajo: en los modelos Qwen3-30B-A3B y GLM4.5-106B-A12B, el tiempo acumulado de entrenamiento de las capas MoE hasta un valor de la función de pérdida (loss) fijado de antemano se redujo entre un 21,2 y un 44,8% respecto a SYMI. Se trata de un sistema anterior que separa el almacenamiento de los parámetros del modelo y de los estados del optimizador para colocar a los expertos con más flexibilidad.

Por qué copiar expertos no cura el desequilibrio de carga

Los expertos se reparten entre las GPU y los tokens viajan hasta las tarjetas correspondientes mediante all-to-all, un intercambio de «todos con todos». Los expertos populares reciben más tokens, sus GPU se convierten en el cuello de botella y las demás esperan. El paso síncrono avanza a la velocidad del participante más cargado.

Las respuestas habituales son reubicar a los expertos o hacer copias de los más populares. Así funcionan FasterMoE, FlexMoE, SmartMoE, LAER-MoE, UltraEP y TopoEP. Pero las copias tienen un coste:

• cada réplica ocupa memoria que se necesita para los búferes de tokens y las activaciones guardadas;
• las copias hay que sincronizarlas y sus gradientes, sumarlos;
• el volumen total de trabajo no disminuye, solo se traslada a otras tarjetas.

Los autores muestran que, con un límite estricto de memoria por GPU, incluso un reparto perfectamente uniforme puede dejar la carga total por encima de lo que la memoria y la red son capaces de atender. El desequilibrio y la sobrecarga en conjunto son problemas distintos, y la igualación de carga solo resuelve el primero.

Qué hace LRMoE: paso hacia delante completo, paso hacia atrás recortado

La idea es tratar la carga del paso hacia atrás no como una consecuencia inevitable del enrutamiento, sino como una magnitud regulable. El trabajo se desarrolla en dos etapas.

• El paso hacia delante sigue siendo completo. Todos los pares token–experto se calculan como de costumbre, porque en esta fase aún no hay gradientes y no hay forma de evaluar el valor de cada par.
• Cuotas antes de calcular. Un controlador ve cuántos tokens han llegado a cada experto y, teniendo en cuenta la memoria y la red, decide cuántos pares de cada experto conservar en el paso hacia atrás. También planifica la ubicación de las réplicas. La cuota se fija para toda la llamada a la capa.
• Selección con la métrica LRM. Cuando los gradientes llegan a la capa, se estima la contribución de cada par a la actualización de los pesos del experto. La contribución de un par es el producto externo del gradiente entrante y la entrada guardada de la proyección. Por eso el cuadrado de su norma se descompone en el producto de tres factores: el cuadrado del peso de enrutamiento, el cuadrado de la norma del gradiente y el cuadrado de la norma de la entrada. Para ello no hace falta calcular las matrices completas de gradientes. Dentro de la cuota se conservan los pares con mayor puntuación.
• Redistribución. Un algoritmo entero reparte los pares restantes entre las copias disponibles de su experto. El objetivo es minimizar la carga pico en una sola GPU. En el paso hacia atrás no aparecen réplicas nuevas.
• Compactación física. Las filas restantes de gradientes y activaciones se agrupan en bloques densos por destinatario y por experto. El intercambio all-to-all se vuelve más compacto y las multiplicaciones de matrices se ejecutan como grouped GEMM de tamaño variable, sin relleno con ceros. Después, el resultado se devuelve a la disposición original.

Un detalle importante: limitarse a enmascarar las filas sobrantes en tensores de tamaño completo no aporta nada, ni en tiempo ni en memoria. La ganancia solo aparece cuando se reducen los tamaños reales de los envíos y de las matrices.

Qué mostraron los experimentos y dónde está la trampa

Además de la cifra principal del 21,2–44,8%, los autores hicieron aparte una reproducción en GPU y obtuvieron una reducción del 57,8% en la memoria total destinada a la ubicación de expertos, las réplicas y los búferes de tokens, en comparación con la asignación fija de expertos. Es una medición de memoria, no del tiempo de entrenamiento.

La trampa es que el paso hacia atrás pasa a ser aproximado: una parte de las contribuciones al gradiente simplemente se descarta. La métrica LRM es una estimación a nivel de una sola proyección y no ofrece garantías de convergencia. Por eso los autores comparan los sistemas por el tiempo hasta un loss fijo y no por la velocidad de cada paso: la aceleración solo cuenta si el modelo alcanza la misma calidad. Los costes adicionales de calcular las puntuaciones, empaquetar y enrutar están incluidos en las mediciones.

La conclusión para los ingenieros es esta. Si el entrenamiento de MoE se topa con la memoria de la GPU y la red, y ya no queda margen para añadir réplicas, conviene considerar la reducción del propio trabajo del paso hacia atrás y no solo su traslado a otras tarjetas. Las salvedades también importan. Los resultados se obtuvieron con dos modelos y se compararon con un único sistema de referencia, así que su extrapolación a otras arquitecturas y a entrenamientos largos debe comprobarse por separado. El coste de los gradientes aproximados puede aflorar con objetivos de calidad más exigentes.