Top‑1 Routing y Poda de Expertos Rebajan el Cómputo de MoE FFN en un 50%

Un simple cambio de enrutamiento de top‑2 a top‑1 puede reducir casi a la mitad las FLOPs de expertos en las transformaciones con mezcla de expertos (MoE), pero muchos equipos no logran ver el incremento de velocidad en producción debido a rezagados, relleno y programación desequilibrada que recuperan la ganancia teórica. Con la matemática de enrutamiento adecuada, ajuste de capacidad y runtimes conscientes de MoE, la dispersión arquitectónica se convierte en tokens/seg estables.

Por qué ahora: modelos de última generación como Mixtral y DeepSeek‑MoE han hecho que MoE sea común, y guías prácticas de DeepSpeed‑MoE y MegaBlocks muestran cómo enrutar y programar expertos sin derretir el rendimiento. Este artículo demuestra cómo la reducción de top‑k y la poda de expertos reducen el cómputo, qué se rompe al cambiar el enrutamiento y cómo los kernels y planificadores restauran el equilibrio.

Revisaremos: a dónde van las FLOPs en MoE (spoiler: las FFNs de expertos dominan), la exacta reducción de 2× FFN de top‑2 → top‑1 y los caminos de comunicación que elimina, ajuste de factores de capacidad y pérdidas auxiliares para evitar el colapso del enrutamiento, poda de expertos basada en utilización, mecánicas de runtime en DeepSpeed‑MoE y MegaBlocks para un rendimiento estable, recuperación de calidad a través de una ligera LoRA y una lista de comprobación práctica con código para implementar y verificar los cambios.

Detalles de Arquitectura/Implementación

Dónde van realmente las FLOPs en MoE

En los bloques Transformer de MoE, la atención se mantiene densa mientras que el MLP/FFN se vuelve disperso mediante expertos. Por token, el enrutamiento selecciona k expertos; cada experto seleccionado aplica su FFN. Con k=2, se realizan dos pases FFN por token; con k=1, uno. Debido a que las FFNs superan en cantidad las FLOPs de atención en muchos ajustes de decodificadores, reducir k de 2 a 1 reduce el cómputo FFN de expertos en aproximadamente un 50% y también reduce el tráfico entre dispositivos cuando los expertos están fragmentados.

De top‑2 a top‑1: matemáticas de compuerta y caminos removidos

Sea x una representación de token, W_r la proyección del enrutador que produce logits g = W_r x. El enrutador selecciona los índices top‑k y aplica un softmax sobre aquellas k entradas para producir pesos de mezcla p. Con k ∈ {1,2}:

# Pseudocódigo para compuerta MoE
logits = router(x) # [num_experts]
indices = topk(logits, k) # k = 1 (top‑1) o 2 (top‑2)
weights = softmax(logits[indices]) # mezcla sobre expertos seleccionados
outputs = sum_j weights[j] * expert_ffn[indices[j]](x)

Operacionalmente, top‑2 introduce dos despachos de expertos por token, a menudo a través de dos colectivos todo-a-todo: los tokens se dividen por experto, se envían a los dispositivos que alojan esos expertos, se procesan y se devuelven. Top‑1 elimina un despacho y reduce a la mitad las invocaciones de FFN. También reduce el relleno: dentro de cada minibatch del experto, menos tokens significan menos espacios inactivos cuando los lotes están limitados por capacidad.

Caminos de comunicación simplificados por top‑1:

• Menos fases todo-a-todo y cargas útiles más pequeñas por paso.
• Menos colas de entrada/salida de expertos para drenar.
• Menos relleno por minibatch de experto cuando se respeta la capacidad.

Estos efectos se intensifican con planificadores conscientes de MoE que coalescen el trabajo y mitigan los rezagados.

Factor de capacidad y pérdida auxiliar: previniendo el colapso del enrutamiento

MoE estilo Switch usa un factor de capacidad C para limitar tokens por experto: capacidad = floor(C × tokens_por_lote / num_expertos). Los tokens que exceden la capacidad son descartados (en tiempo de entrenamiento) o re-rutados, y una pérdida auxiliar de balance de carga anima al enrutador a distribuir los tokens de manera uniforme mientras preserva asignaciones de alta probabilidad. Después de reducir k, dos cosas comúnmente fallan:

• Colapso: el enrutador se concentra demasiado en unos pocos expertos, saturando la capacidad y causando desbordamiento/relleno.
• Infrautilización: algunos expertos se enfrían, perjudicando la especialización y la calidad.

Contramedidas:

• Aumentar el coeficiente de pérdida de balance de carga y comenzar desde un checkpoint que ya usa esa pérdida auxiliar.
• Reajustar C. Con top‑1, aumentar ligeramente C (por ejemplo, 1.0 → 1.2) proporciona espacio para evitar desbordamientos en horizontes cortos; en contextos largos, se puede reducir C para disminuir el relleno.
• Aplicar una breve recuperación solo del enrutador o LoRA‑en‑enrutador para redistribuir el tráfico.

Poda de expertos mediante histogramas de utilización

Una vez que el enrutamiento es estable bajo top‑1, perfilar la utilización de expertos sobre cargas de trabajo representativas. Calcular la fracción de tokens (o tokens ponderados por puntuación de puerta) manejados por experto. Los expertos por debajo de un umbral (por ejemplo, <0.5-1.0% en largas ventanas) son candidatos para la poda. Después de la eliminación, reequilibrar el enrutador y ejecutar una ligera recuperación para recentrar la especialización.

Salvaguardias clave:

• Usar trazas largas y mensajes diversos para evitar podar expertos que sirven nichos raros pero esenciales.
• Preferir una poda escalonada con pequeños lotes de expertos eliminados a la vez, con validación entre etapas.

Problemas de rendimiento: rezagados, relleno, desequilibrio

Los cambios de enrutamiento a menudo destacan ineficiencias ocultas:

• Rezagados: un experto muy cargado alarga el paso mientras otros están inactivos.
• Relleno: los kernels de forma fija rellenan hasta la capacidad; una asignación de token a experto desigual magnifica los espacios desperdiciados.
• Sesgo de comunicación: un tráfico todo-a-todo desequilibrado produce latencia de cola.

MegaBlocks aborda estos problemas reagruparando minibatches en bloques equilibrados, reduciendo el relleno y suavizando la carga a nivel del dispositivo; permanece efectivo bajo distribuciones sesgadas y conjuntos de expertos podados. DeepSpeed‑MoE proporciona controles de tiempo de ejecución para capacidad, políticas de desbordamiento y comunicación/cálculo superpuestos que estabilizan p50/p99 bajo top‑1.

Runtimes conscientes de MoE: DeepSpeed‑MoE y MegaBlocks

• DeepSpeed‑MoE: un motor listo para producción con compuerta top‑k, controles de factor de capacidad, pérdidas de balance de carga y fragmentación paralela de expertos. Superpone todo-a-todo con el cálculo de experto y expone la configuración para enrutamiento, capacidad y partición de expertos.
• MegaBlocks: un enfoque de kernel y programación que divide la computación en bloques uniformes, limitando el relleno y mitigando los rezagados. Proporciona un rendimiento robusto cuando las skews de enrutamiento o la poda de expertos interrumpen la uniformidad.

Ambos enfoques son compatibles con las tuberías de entrenamiento/inferencia de MoE estilo Mixtral y DeepSeek‑MoE y funcionan bien con formulaciones de enrutador/pérdida auxiliar derivadas de Transformers Switch.

Notas de implementación en pilas de GPU sin tocar kernels densos

• CUDA/NVIDIA: Si tus rutas densas de atención/MLP ya usan kernels fusionados, puedes adoptar top‑1 y poda de expertos sin modificar kernels densos. Enfócate en los kernels de despacho de MoE y el solapamiento todo-a-todo en DeepSpeed‑MoE, o intercambia en MegaBlocks para el cálculo de expertos agrupados en bloques.
• Triton/ROCm: MegaBlocks se compila a través de Triton y se ejecuta en ROCm, brindando a los usuarios de AMD un camino portátil para lograr las ganancias de enrutamiento de MoE sin reescribir kernels densos.
• Entornos mixtos: Mantén los kernels densos intactos; los cambios de enrutamiento ocurren en el límite de despacho/agregación de MoE. Valida el rendimiento colectivo y la alineación de diseño de memoria, no la matemática de atención/FFN en sí.

Mantener la calidad después de cambios estructurales

El enrutamiento top‑1 y la poda de expertos alteran la especialización. Una recuperación ligera estabiliza las métricas:

• LoRA en el enrutador (y opcionalmente en las feed‑forwards de expertos) por unos pocos miles de pasos a menudo es suficiente para restaurar el equilibrio de utilización y recuperar 0.5-2 puntos en evaluaciones comunes, a una fracción del costo de ajuste fino completo.
• Una breve destilación de conocimiento desde el modelo previo a la poda puede preservar aún más el conocimiento de expertos en patrones raros tras la poda.

Tablas de Comparación

Opciones de enrutamiento y poda de un vistazo

Enfoque	Efecto en el cálculo	Efecto en la comunicación	Riesgos principales	Mejores mitigaciones	Runtime de soporte
Top‑2 (línea base)	2× FFN de experto por token	Dos despachos todo-a-todo	Más FLOPs, relleno	N/A (línea base)	Cualquier runtime de MoE
Top‑1	~50% menos FLOPs de FFN de expertos	Un despacho eliminado; cargas más pequeñas	Colapso de enrutamiento; expertos fríos	Ajustar pérdida auxiliar, ajustar C, corta recuperación de enrutador	DeepSpeed‑MoE, MegaBlocks
Poda de expertos	Reduce el conteo de parámetros y expertos activos	Menos destinos; puede aumentar la sesgo	Excesiva poda perjudica especialización	Histogramas de utilización, poda escalonada, recuperación LoRA	DeepSpeed‑MoE, MegaBlocks

Mecánicas de runtime: qué estabiliza el rendimiento

Runtime	Mecanismo clave	Fortalezas bajo top‑1/poda	Notas
DeepSpeed‑MoE	Todo-a-todo superpuesto + cálculo de experto; controles de capacidad/pérdida auxiliar	Configuraciones listas para producción; fácil ajuste del enrutador	Integración a través de tuberías de entrenamiento/inferencia
MegaBlocks	Programación y cálculo agrupados por bloques	Mitiga relleno/rezagados, robusto al sesgo	Kernels Triton, portátil a ROCm

Mejores Prácticas

• Comienza con top‑1, luego poda: Primero estabiliza el enrutamiento en k=1, luego poda expertos fríos basándote en la utilización a largo plazo.
• Ajusta el enrutador como un módulo de primera clase:
• Incrementa el peso de la pérdida de balance de carga después de cambiar k; monitorea la entropía y el tráfico por experto.
• Reajustar el factor de capacidad C por tamaño de lote y longitud de secuencia; favorece ligeramente un C más alto para secuencias cortas para evitar desbordamiento, un C más bajo para secuencias largas para reducir el relleno.
• Usa runtimes conscientes de MoE:
• Habilita el solapamiento todo-a-todo y la aplicación de capacidad de DeepSpeed‑MoE, o adopta MegaBlocks para reducir el relleno y los rezagados.
• Realiza un perfil antes de podar:
• Registra histogramas de asignación de expertos y mapas de calor de latencia por experto bajo tráfico similar al de producción; poda sólo expertos persistentemente fríos.
• Recupera ligeramente:
• Aplica LoRA al enrutador/FFNs de expertos para una recuperación corta; opcionalmente agrega una pequeña mezcla de destilación para recentrar la especialización.
• Verifica comportamiento en largo contexto: Algunos expertos capturan patrones largos raros; incluye evaluaciones de contexto largo al podar.

Ejemplos Prácticos

1) Configuración de DeepSpeed‑MoE para top‑1 y ajuste de capacidad

{
 "moe": {
 "enabled": true,
 "num_experts": 16,
 "top_k": 1,
 "capacity_factor": 1.25,
 "router_aux_loss_coef": 0.01,
 "router_topk": true,
 "moe_param_group": true,
 "drop_tokens": false,
 "alltoall": { "overlap_communication": true }
 }
}

• top_k: cambio de 2 → 1 para reducir a la mitad el cómputo FFN de expertos.
• capacity_factor: inicia ligeramente por encima de 1.0 para prevenir desbordamiento en contextos cortos; vuelve a ajustar por carga de trabajo.
• router_aux_loss_coef: aumentalo modestamente después del cambio de top‑k para prevenir colapso.
• overlap_communication: mantener las GPUs ocupadas durante el todo-a-todo.

2) Histograma de utilización para impulsar la poda de expertos

# Fragmento parecido a PyTorch capturando estadísticas de enrutamiento durante eval
util = torch.zeros(num_experts, device="cuda")
with torch.no_grad():
 for batch in loader:
 logits = router(batch.hiddens) # [B, num_experts]
 idx = torch.topk(logits, k=1, dim=-1).indices.squeeze(-1) # top‑1
 util.index_add_(0, idx.flatten(), torch.ones_like(idx.flatten(), dtype=util.dtype))

util = util / util.sum() # fracción de tokens por experto
cold = (util < 0.005).nonzero().flatten() # e.g., <0.5%
print("Candidatos para poda:", cold.tolist())

• Ejecuta sobre trazas largas y diversas. Confirma con mapas de calor de latencia por experto.
• Poda en etapas; después de cada etapa, ejecuta una corta recuperación LoRA sobre el enrutador y expertos afectados.

3) LoRA solo del enrutador para recuperación rápida

from peft import LoraConfig, get_peft_model

lora_cfg = LoraConfig(
 r=8, lora_alpha=16, target_modules=["router.proj"], lora_dropout=0.05
)
model.router = get_peft_model(model.router, lora_cfg)
# Afinar por unos pocos K pasos en tareas mixtas para recentrar la especialización

• Enfocarse primero en el enrutador; expandirse a FFNs de expertos selectivamente si la especialización cambió.

Conclusión

El enrutamiento top‑1 es la palanca más clara para reducir el cómputo FFN de MoE aproximadamente a la mitad, pero la ganancia solo se materializa de extremo a extremo cuando manejas el enrutador y el planificador como sistemas de primera clase. El factor de capacidad y las pérdidas de balance de carga mantienen la utilización uniforme; los histogramas de utilización identifican expertos que pueden ser podados de manera segura; y los runtimes conscientes de MoE como DeepSpeed‑MoE y MegaBlocks convierten la dispersión arquitectónica en tokens/seg estables al domar el relleno, los rezagados y el sesgo de comunicación. Aún más, un toque final de LoRA ayuda a recuperar la especialización después de cambios estructurales.

Puntos clave:

• Top‑2 → top‑1 elimina un pase de experto y una etapa todo-a-todo por token, reduciendo a la mitad las FLOPs de FFN de expertos y recortando el tráfico entre dispositivos.
• El ajuste de capacidad y la sintonización de pérdidas auxiliares previenen el colapso de enrutamiento; monitorea la utilización y la entropía como métricas de primera clase.
• Poda solo expertos persistentemente fríos, guiado por histogramas a largo plazo; recupera ligeramente con LoRA.
• DeepSpeed‑MoE y MegaBlocks estabilizan el rendimiento sobreponiendo la comunicación y reduciendo el relleno/rezagados, incluso bajo sesgo y poda.

Próximos pasos:

• Permutar a top‑1 en un entorno de puesta en escena, ajustar capacidad y pérdidas auxiliares, y validar utilización y p99.
• Recoger histogramas sobre tráfico real; estandarizar la poda de expertos con chequeos entre rondas.
• Adoptar características de runtimes conscientes de MoE (superposición de todo-a-todo o programación en bloques) y agregar una breve recuperación LoRA.

El resultado: con un enrutamiento y programación disciplinados, la dispersión teórica de MoE se convierte en rendimiento en el mundo real—sin reescribir tus kernels densos ⚙️.

Fuentes

• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity — https://arxiv.org/abs/2101.03961 — Establece la compuerta top‑k, factores de capacidad y pérdidas auxiliares de balance de carga que subyacen la estabilidad del enrutamiento top‑1.
• Mixtral of Experts — https://arxiv.org/abs/2401.04088 — Demuestra arquitecturas prácticas de MoE y discute comportamientos de enrutamiento y consideraciones de especialización relevantes para cambios de top‑k y poda.
• DeepSeek‑MoE (repositorio) — https://github.com/deepseek-ai/DeepSeek-MoE — Muestra controles de compuerta del mundo real y gestión de expertos usados al aplicar ajustes de enrutamiento y poda.
• MegaBlocks: Efficient Sparse Mixture‑of‑Experts Training and Inference — https://arxiv.org/abs/2211.15841 — Proporciona programación/kernels agrupados en bloques que mitigan el relleno, los rezagados y el desequilibrio bajo cargas de expertos sesgadas/podadas.
• DeepSpeed‑MoE Tutorial — https://www.deepspeed.ai/tutorials/moe/ — Documenta mecánicas de runtime, sintonización de capacidad/pérdida auxiliar y comunicación/computación sobrepuesta que estabilizan el rendimiento bajo top‑1.
• AMD ROCm Documentation — https://rocm.docs.amd.com/ — Confirma vías Triton/ROCm para desplegar kernels y MoE tipo MegaBlocks en AMD sin reescribir kernels densos.
• LoRA: Low‑Rank Adaptation of Large Language Models — https://arxiv.org/abs/2106.09685 — Soporta el ajuste fino de adaptadores ligeros para recuperar calidad después de enrutamiento/poda.
• AdaLoRA: Adaptive Budget Allocation for Parameter‑Efficient Fine‑Tuning — https://arxiv.org/abs/2303.10512 — Proporciona una alternativa eficiente para recuperación y recentrado de especialización después de poda.

Fuentes y Referencias