English · Español

Lab 00 — Una variante MoE de 2 experts de MiniGPT-grammar (y por qué no ayuda)

Objetivo: cambiar la FFN de un bloque transformer por un MoE (Mixture of Experts) de 2 experts. Entrenar localmente en CPU. Comparar honestamente con la baseline densa. Confirmar el resultado negativo: MoE no ayuda al tutor de gramática.

Tiempo estimado: 3–4 horas.

Prerrequisito: MiniGPT-grammar de la Fase 17 entrenado; bucle de entrenamiento de la Fase 18 en src/minitrain/; hooks de inspección de la Fase 19 para monitorizar la pérdida. Borja ha leído theory/01-moe.md.

---

Lo que produces

Un directorio experiments/36-moe-on-grammar-tutor/ que contenga:

• moe_block.py — módulo local (alcance del experimento) con una clase MoEBlock que implementa routing + 2 FFNs expert + aux loss de balanceo de carga. Según el plan de la fase, esto vive en el directorio del experimento (o como src/minimodel/_moe_block.py si Borja lo prefiere — un único archivo opcional bajo un módulo existente, no un módulo nuevo de primer nivel).
• train_moe.py — script de entrenamiento: construye una variante de MiniGPT-grammar con el bloque MoE reemplazando una o dos de las FFNs densas.
• compare.py — corre la baseline (MiniGPT denso de la Fase 18) y la variante MoE, lado a lado: mismos datos, misma semilla, mismo número de pasos.
• manifest.json — métricas finales, conteos de parámetros, perplejidad, tiempo de entrenamiento.
• findings.md — el informe honesto sobre la comparación.

TODOs

Bloque A — implementar el bloque MoE

# experiments/36-moe-on-grammar-tutor/moe_block.py
# Skeleton — Borja escribe el cuerpo. Target ~150 LOC.

import torch
from torch import nn
import torch.nn.functional as F

class MoEBlock(nn.Module):
    """Top-2 routing of E experts. Returns (output, aux_loss).

    For the grammar tutor at d_model=64, d_ff=256, set E=2, k=2 to start
    (every token goes to every expert — degenerate, mostly for testing).
    Then E=4, k=2 for the actual experiment.
    """
    def __init__(self, d_model: int, d_ff: int, num_experts: int, top_k: int):
        super().__init__()
        self.gate = nn.Linear(d_model, num_experts, bias=False)
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_ff),
                nn.GELU(),
                nn.Linear(d_ff, d_model),
            )
            for _ in range(num_experts)
        ])
        self.top_k = top_k
        self.num_experts = num_experts

    def forward(self, x):
        # x: [B, T, d_model]
        # 1. compute gate logits, softmax, top-k indices and weights
        # 2. for each expert, gather the tokens assigned to it
        # 3. compute the expert's FFN on its assigned tokens
        # 4. scatter back into the output, weighted by gate
        # 5. compute auxiliary load-balance loss (theory/01-moe.md §2)
        # 6. return (output [B, T, d_model], aux_loss [scalar])
        ...

Notas de implementación:

• Sin capacity dropping. A escala de tutor de gramática, cada token cabe. Capacity factor = ∞.
• Top-k diferenciable. El torch.topk de PyTorch no es diferenciable a través de los índices, pero los pesos del gate (los valores de la softmax) propagan gradientes. Con eso vale.
• Escala de la aux loss. Empezar con \(\alpha = 0.01\). Sintonizar si el routing colapsa.

Tests en tests/minimodel/test_moe_block.py (si colocas el bloque en src/minimodel/):

• test_routing_is_top_k — cada token recibe exactamente \(k\) pesos de gate no nulos.
• test_aux_loss_is_one_for_uniform — para routing uniforme, la aux loss vale 1.
• test_aux_loss_grows_under_imbalance — alimentar todos los tokens al expert 0; la aux loss debería ser \(E\).
• test_total_params_grow_as_E — el conteo de parámetros del MoE con \(E\) experts ≈ \(E\) × FFN densa.

Bloque B — modificar el modelo

Toma el MiniGPT-grammar de la Fase 17 y reemplaza la FFN de exactamente un bloque por el bloque MoE. Reemplazar una sola FFN es suficiente — reemplazarlas todas añade coste proporcional sin perspectiva proporcional a esta escala.

La factoría del modelo debe aceptar --moe-layer=1 (qué índice de bloque recibe el MoE) y --num-experts=4 --top-k=2.

Bloque C — el protocolo de comparación

Para una comparación justa:

• Misma semilla (p. ej., 36000).
• Mismos shards de datos (splits del corpus de la Fase 12).
• Mismo schedule de LR (cosine + warmup de la Fase 18).
• Mismo número de pasos (p. ej., 5000 pasos).
• Mismo batch size.
• Mismo set de probes de eval (harness de eval de la Fase 20).

Corre ambos modelos. Registra:

• Pérdida final de entrenamiento.
• Perplejidad final de validación.
• Parámetros totales.
• Wall time de entrenamiento.
• Distribución de carga por expert (en el MoE) — media y stddev de asignación de tokens entre los 4 experts.

Bloque D — el análisis de routing

Para el modelo MoE, instrumenta el forward para loggear a qué expert(s) va cada token. Tras entrenar, produce dos artefactos de análisis:

• Histograma de asignación de experts por token — para un batch de validación representativo, dibuja la distribución del routing. Comprueba: ¿es uniforme? ¿sesgada? ¿hay experts muertos?
• Asignación de expert por clase de token — agrupa tokens por su etiqueta del corpus (p. ej., "verb-infinitive", "verb-past", "subject-pronoun"). ¿Se especializan los experts en categorías lingüísticas, o en algo arbitrario como la frecuencia de token?

Entrega ambos como gráficos en experiments/36-moe-on-grammar-tutor/routing-analysis.png y expert-by-class.png.

Bloque E — los findings honestos

findings.md — 300–600 palabras. Aborda cada uno de:

1. ¿La variante MoE consiguió menor perplejidad que la baseline densa? (Esperado: no, o marginalmente — dentro del ruido.)
2. ¿Cuántos más parámetros tiene el MoE? (Esperado: ~2× lo que vale la FFN.)
3. ¿Funcionó la pérdida de balanceo de carga — recibieron los experts conteos de tokens aproximadamente iguales?
4. ¿Se especializaron los experts lingüísticamente, o arbitrariamente?
5. ¿El entrenamiento tardó más? ¿Por qué?
6. Resumen: ¿debería el tutor de gramática adoptar MoE? Respuesta: no, con razonamiento.

Sé honesto. Si por casualidad el MoE sí bate a la baseline (puede pasar por ruido, efectos de regularización o un init aleatorio afortunado), documéntalo, pero también reporta la varianza: reejecuta ambos con 3 semillas distintas y reporta media ± std. No cherry-pickees.

Bloque F — manifest

experiments/36-moe-on-grammar-tutor/manifest.json:

{
  "seed": 36000,
  "lab": "00-moe-on-grammar-tutor",
  "num_experts": 4,
  "top_k": 2,
  "moe_layer_index": 1,
  "alpha_aux_loss": 0.01,
  "steps": 5000,
  "dense_baseline": {
    "params": "<filled in>",
    "final_val_ppl": "<filled in>",
    "wall_time_s": "<filled in>"
  },
  "moe_variant": {
    "params": "<filled in>",
    "final_val_ppl": "<filled in>",
    "wall_time_s": "<filled in>",
    "load_balance_stddev": "<filled in>",
    "experts_dead_count": "<filled in>"
  },
  "verdict_for_grammar_tutor": "<adopt / defer / never>",
  "lesson_notes": "<the bottom-line paragraph from findings.md>"
}

Restricciones

• Solo CPU. Sin gasto en la nube.
• Máximo 5000 pasos. La convergencia es rápida sobre este corpus; más pasos no cambian la respuesta.
• Sin optimizadores exóticos. AdamW de la Fase 18. Mismo schedule de LR que la baseline.
• Sin trucos especiales de MoE. Sin expert dropout, sin inyección de ruido en experts, sin routing fancy. Top-2 vainilla + aux loss. El objetivo es ver el MoE, no sobreajustarlo para ganar.
• Bench con 3 semillas, no 1. Si la varianza eclipsa al delta denso-vs-MoE, repórtalo honestamente. No reclames una victoria que está dentro del ruido.

Condiciones de parada

Has terminado cuando:

1. experiments/36-moe-on-grammar-tutor/{moe_block.py, train_moe.py, compare.py, findings.md, manifest.json, routing-analysis.png, expert-by-class.png} existen todos.
2. findings.md responde a las seis preguntas del Bloque E.
3. manifest.json tiene los números de comparación.
4. El veredicto para el tutor de gramática es defer o never — y si es adopt, tienes evidencia de 3 semillas mostrando que la ganancia supera la varianza.

Pista de último recurso

Si el routing colapsa (un expert se lleva >80% de los tokens):

• Primera comprobación: ¿es alpha_aux_loss demasiado bajo? Prueba 0.05.
• Segunda comprobación: ¿pesos del gate inicializados a cero? Inicializa \(W_g\) con valores aleatorios pequeños (std=0.02).
• Tercera comprobación: top-\(k\) vs \(E\) — para \(E=2, k=2\), el routing es degenerado (cada token va a cada expert), y la aux loss no tiene palanca. Usa \(E=4, k=2\).

Si el entrenamiento del MoE es inestable (spikes en la pérdida): baja la LR 2×. El MoE es más sensible a la LR que el denso.

Si el MoE bate a la baseline por un margen cómodo (>10% de reducción de perplejidad con baja varianza): tómatelo en serio, pero también comprueba que no hayas hecho accidentalmente la baseline densa más débil que la de la Fase 18. La comparación debe ser de manzanas con manzanas.

Cuándo consultar solutions/

Tras commitear findings.md. La solución vive en solutions/00-moe-ref.md — escrita al abrir la fase, una vez que las Fases 17 + 18 de Borja están en su sitio. La solución de referencia incluye los rangos esperados de los números de comparación (dentro de bandas de varianza) para que Borja pueda hacer sanity-check de su ejecución.

---

Siguiente lab: lab/01-mla-math-exercise.md.