Skip to content

English · Español

Break 00 — MoE con un router roto (todos los tokens a un expert); el caso degenerado

🇪🇸 Si quitas la auxiliary loss del MoE, el router rápidamente colapsa a un experto único. La pérdida de entrenamiento se ve "bien" — todavía baja — pero el modelo se comporta como una FFN normal con N-1 experts sin entrenar. Este /break lo demuestra desactivando la aux loss.

Lo que vas a hacer

Desactivar la pérdida auxiliar de balanceo de carga en el MoE (Mixture of Experts) de juguete construido en docs/phase-36-frontier-architectures/lab/00-moe-on-grammar-tutor.md. Observar cómo el router colapsa a un único expert en menos de 200 pasos; observar que la pérdida principal no señala el fallo.

Paso 1 — Localizar el bloque MoE

src/minimoe/moe_block.py          # the toy MoE layer (Phase 36 lab 00)
src/minimoe/loss.py               # the combined train loss = main + aux

Paso 2 — Introducir el bug

En src/minimoe/loss.py, la pérdida combinada actualmente pondera el término auxiliar de balanceo de carga por alpha=0.01. Ponlo a cero:

# OLD
total_loss = main_loss + 0.01 * aux_load_balancing_loss(gates, mask)

# NEW (the broken version)
total_loss = main_loss + 0.0 * aux_load_balancing_loss(gates, mask)

Una constante numérica. La aux loss se sigue computando y logueando (así tenemos una métrica para detectar el colapso), simplemente no propaga gradiente.

Paso 3 — Registrar el break

learners/borja/phase-36/notes/breaks.md:

- bug-id: 36-01
  concept: MoE load-balancing aux loss
  symptom: main train loss falls normally; but the per-expert token count
           collapses — after ~200 steps, expert 0 sees 95%+ of tokens and the
           other 3 experts see ~0. Held-out val loss is worse than the dense
           baseline despite 4× the params.
  hidden_cause: the aux loss coefficient is zero; nothing pushes the router
                to spread tokens across experts.
  hint_1: "Log per-expert token counts. Plot them across steps. What's the trend?"
  hint_2: "What does the aux loss penalize, and what does its coefficient control?"
  hint_3: "diff src/minimoe/loss.py. Has anything multiplying the aux loss changed?"
  fix_diff: restore alpha=0.01 in the loss combination.

Paso 4 — Verificar que es observable

Ejecuta just moe-train --steps 500. Salida esperada con el bug:

step=  10  main=2.401  aux=0.124  expert_counts=[ 240  238  244  254 ]
step= 100  main=2.018  aux=0.310  expert_counts=[ 540  191  119  126 ]
step= 200  main=1.872  aux=0.892  expert_counts=[ 893  62   18   3  ]
step= 300  main=1.748  aux=0.988  expert_counts=[ 945  20   8    3  ]
step= 400  main=1.665  aux=0.996  expert_counts=[ 962  10   3    1  ]
step= 500  main=1.601  aux=0.998  expert_counts=[ 970  4    2    0  ]   <-- collapse
val_loss=2.143  (dense_baseline=1.890)   <-- 4× params, *worse* val loss

La métrica aux_load_balancing_loss sube hacia 1.0 (su máximo bajo la formulación Switch: \(f_e P_e\) sumado con toda la masa en un expert tiende a 1). La pérdida principal pinta bien. La val loss held-out es la prueba del crimen.

El test tests/phase36/test_moe_balance.py::test_no_expert_above_50_percent se pone en rojo.

Paso 5 — El momento docente

Dos lecciones, un bug:

  1. La pérdida principal por sí sola no detecta la patología del MoE. Va a seguir bajando porque el único expert activo sigue aprendiendo una FFN decente. La arquitectura está degenerando silenciosamente a un modelo de 1 expert con 3 experts muertos.
  2. La aux loss es estructural, no cosmética. Sin ella, el sistema no es un MoE en ningún sentido significativo. El coeficiente 0.01 parece pequeño; borrarlo borra la arquitectura.

El fix es una constante; la lección es que la función de pérdida es parte de la definición arquitectónica, no un detalle ajustable del entrenamiento.

Hard rules respetadas

  • Un único bug (una constante numérica).
  • Reversible en 1 línea.
  • Observable (los conteos de tokens por expert divergen; la val loss regresiona).
  • Sin impacto de seguridad.
  • Tests no modificados.

Siguiente: cuando esté en verde, relee ../theory/05-moe-routing-math-and-mamba-intuition.md para la derivación formal de \(f_e P_e\).