MoE desde cero + híbridos de frontera (cierre del curso)

5.1La idea: capacidad sin coste proporcional

Un MLP denso usa todos sus parámetros para cada token. Un Mixture of Experts (MoE) tiene muchos "expertos" (MLPs) pero un router que envía cada token solo a unos pocos (top-k). Resultado: un modelo con, digamos, 235B parámetros totales pero solo ~22B activos por token (Qwen3-235B-A22B). Ganas capacidad (más parámetros = más conocimiento) sin pagar el cómputo de activarlos todos.

Por eso casi todos los modelos de frontera 2026 (DeepSeek V4, Kimi K2, GLM, Qwen3, gpt-oss) son MoE: es la forma de escalar parámetros sin que el coste de inferencia explote.

5.2Las piezas de un MoE

• Expertos: N MLPs independientes (típicamente los FFN de cada bloque).
• Router (gating): una capa lineal que, para cada token, produce logits sobre los N expertos; se eligen los top-k (k=1 o 2 normalmente) y se ponderan por el softmax de sus logits.
• Balance de carga: sin incentivo, el router colapsa a usar siempre los mismos pocos expertos (los demás no aprenden). Una auxiliary loss penaliza el desequilibrio para que la carga se reparta.
• Expertos compartidos (DeepSeek): algunos expertos se activan siempre (capturan conocimiento común), el resto son especializados.

5.3Laboratorio L5.1 — Un micro-MoE desde cero

1# lab_n5l5_moe.py — capa MoE con routing top-2 y auxiliary loss de balance, desde cero
2import torch, torch.nn as nn, torch.nn.functional as F
3
4class Expert(nn.Module):
5    def __init__(self, d):
6        super().__init__()
7        self.net = nn.Sequential(nn.Linear(d, 4*d), nn.GELU(), nn.Linear(4*d, d))
8    def forward(self, x): return self.net(x)
9
10class MoELayer(nn.Module):
11    def __init__(self, d, n_experts=8, k=2):
12        super().__init__()
13        self.k = k; self.n_experts = n_experts
14        self.experts = nn.ModuleList([Expert(d) for _ in range(n_experts)])
15        self.router = nn.Linear(d, n_experts, bias=False)   # gating
16
17    def forward(self, x):                                   # x: [B, T, d]
18        B, T, d = x.shape
19        x_flat = x.reshape(-1, d)                           # [B*T, d] cada token enruta solo
20        logits = self.router(x_flat)                        # [N_tokens, n_experts]
21        probs = F.softmax(logits, dim=-1)
22        topk_val, topk_idx = probs.topk(self.k, dim=-1)     # top-k expertos por token
23        topk_val = topk_val / topk_val.sum(-1, keepdim=True)  # renormaliza los pesos elegidos
24
25        out = torch.zeros_like(x_flat)
26        for i in range(self.k):
27            idx = topk_idx[:, i]                            # qué experto le toca a cada token (slot i)
28            w = topk_val[:, i].unsqueeze(-1)
29            for e in range(self.n_experts):                # (didáctico; en prod se hace con scatter)
30                mask = (idx == e)
31                if mask.any():
32                    out[mask] += w[mask] * self.experts[e](x_flat[mask])
33
34        # --- auxiliary load-balancing loss ---
35        # fracción de tokens que va a cada experto vs probabilidad media asignada
36        importance = probs.mean(0)                          # prob media por experto
37        load = torch.zeros(self.n_experts, device=x.device)
38        load.scatter_add_(0, topk_idx.reshape(-1),
39                          torch.ones(topk_idx.numel(), device=x.device))
40        load = load / load.sum()
41        aux_loss = self.n_experts * (importance * load).sum()  # mínima cuando la carga es uniforme
42        return out.reshape(B, T, d), aux_loss

Líneas no triviales explicadas:

• router + topk: el corazón del MoE. Cada token recibe logits sobre expertos; solo los k mejores se activan. Esa esparsidad es lo que ahorra cómputo.
• Renormalizar topk_val: los pesos de los k expertos elegidos se renormalizan para sumar 1, de modo que la salida sea una combinación convexa.
• aux_loss de balance: penaliza cuando importance (prob media) y load (uso real) se concentran en pocos expertos. Sin esto, el router colapsa (usa 2 expertos, los otros 6 mueren). Se suma a la loss principal con un peso pequeño.
• El doble bucle es didáctico; en producción se usa dispatch con scatter/gather o kernels especializados (Megablocks, grouped GEMM) para eficiencia.

5.4Laboratorio L5.2 — Observar y arreglar el colapso de routing

1# Entrena tu MoE en el dataset de N5·L2 con y sin aux_loss, y registra el uso de expertos.
2# 1) SIN aux_loss (peso 0): verás que pocos expertos acaparan casi todos los tokens (colapso).
3# 2) CON aux_loss (peso ~0.01): la carga se reparte; todos los expertos aprenden.
4# Loguea un histograma del uso de expertos a lo largo del entrenamiento -> la prueba visual.
5loss_total = loss_principal + 0.01 * aux_loss

Ver el colapso y cómo la aux loss lo evita es la lección central del MoE. Es un fallo que sorprende a quien no lo entiende.

5.5Híbridos de frontera (dónde converge todo el curso)

Los modelos punteros de 2026 combinan las tres ideas de este nivel:

• MoE para escalar parámetros baratos (capacidad).
• Atención para recall preciso.
• Mamba/SSM para eficiencia en contexto largo.

Ejemplo concreto que cabe en tu 5090 para fine-tuning: Nemotron Nano (híbrido Mamba-Transformer + MoE latente) o Qwen3-30B-A3B (MoE, 30B totales / 3B activos, QLoRA en ~17.5 GB con Unsloth, N2). Tomar uno de estos y hacerle QLoRA cierra el círculo: aplicas todo lo aprendido (arquitectura N5, fine-tuning N2, presupuesto VRAM N0, serving N1) sobre un modelo de arquitectura de frontera.

5.6Ejercicios

E1. Entrena tu MoE con y sin aux_loss y grafica el uso de expertos. ¿En cuántos pasos colapsa sin balance?

E2. Varía n_experts (4, 8, 16) y k (1, 2). ¿Cómo cambian calidad, cómputo activo y balance?

E3. Compara tu MoE (8 expertos, top-2) con un MLP denso del mismo cómputo activo. ¿Gana el MoE en perplexity a igualdad de FLOPs por token?

5.7Trampas comunes

• Olvidar la aux loss → colapso de routing, expertos muertos.
• Confundir parámetros totales con activos (el coste de inferencia depende de los activos; la VRAM, de los totales — N0·L3).
• Implementar el dispatch con bucles en producción → lentísimo; usa scatter/gather o kernels MoE.

5.8Referencias

• Switch Transformer (Fedus et al. 2021), GShard, Mixtral (Mistral 2024), DeepSeek-V3 (expertos compartidos). CS336 L4 "Mixture of experts". Repos: Megablocks. Unsloth (Qwen3-30B-A3B QLoRA, Nemotron Nano).

CIERRE DEL CURSO

Has recorrido los seis niveles hoja por hoja:

• N0 — Fundamentos: el roofline, el stack Blackwell, qué cabe en 32 GB, tu benchmark base.
• N1 — Serving (spine): del motor de inferencia al multi-LoRA, dentro del roofline.
• N2 — Post-training + RL (spine principal): de SFT a GRPO/RLVR y agentes verificables.
• N3 — Breadth: kernels, interpretabilidad, difusión, audio/long-context.
• N4 — World models (frontera): V-JEPA 2, DIAMOND, mundos interactivos.
• N5 — Pretraining + arquitecturas: desde backprop hasta GPT-2 reproducido, Mamba-2 y MoE.

Eso es el perfil T completo: profundidad world-class en dos spines, breadth de reproductor en el resto, frontera tocada, cimientos sólidos. Medido contra benchmarks objetivos, no contra el aplauso.

El último paso es el capstone (programa maestro, sección 7): un proyecto que encadena varios niveles —pre-entrenar/tomar un modelo pequeño, fine-tunearlo con RLVR, escribir un kernel que lo acelere, cuantizarlo y servirlo dentro del roofline, evaluarlo con rigor e interpretar sus fallos—. Cuando lo completes y cada paso pase su checkpoint, estarás en el percentil más alto y lo sabrás, con números, sin necesitar que nadie lo valide.

Ese era el objetivo. El curso está completo; el trabajo es tuyo.