Mixture of Experts optimiza Transformers y LLMs | Keryc
En los últimos años, escalar modelos densos fue la receta principal para mejorar LLMs: más datos + más parámetros = mejor rendimiento. Pero eso choca con límites prácticos: entrenamiento caro, latencia de inferencia creciente y necesidades de memoria que no cualquiera puede cubrir.
¿La solución? Mixture of Experts (MoEs). Mantienen la columna vertebral Transformer, pero reemplazan capas feed-forward densas por un conjunto de expertos (sub-redes entrenables). Un router decide, token por token, qué pocos expertos procesan cada entrada. Resultado: capacidad total alta por parámetros totales, pero velocidad de inferencia depende solo de los parámetros activos por token. ¿No suena atractivo?
¿Por qué MoEs importan para ingenieros y equipos que despliegan LLMs?
MoEs ofrecen tres ventajas claras:
Eficiencia de cómputo: con el mismo presupuesto de FLOPs de entrenamiento, MoEs tienden a superar a modelos densos, permitiendo iteraciones más rápidas.
Eje natural de paralelización: como diferentes tokens activan distintos expertos, puedes paralelizar por experto (expert parallelism), útil en clusters.
Adopción industrial: modelos recientes como Qwen 3.5, MiniMax M2 o GLM-5 usan MoEs; DeepSeek R1 marcó un punto de inflexión.
Un ejemplo práctico: gpt-oss-20b tiene 21B parámetros totales, pero sólo 4 expertos activos por token de un conjunto de 32. Eso se traduce en ~3.6B parámetros activos por token. En un M3 Ultra (ancho de banda ~800 GB/s) y usando bfloat16 (2 bytes por parámetro), la estimación rápido da ~111 tokens/s; la medición real fue ~115 tok/s. Es decir: latencia comparable a un modelo de 3.6B, calidad cercana a 21B.
Ingeniería: el principal dolor de MoEs y cómo Transformers lo resolvió
El ecosistema estaba pensado para modelos densos: carga de checkpoints, colocación en dispositivos, cuantización y ejecución. MoEs rompen varias suposiciones:
En checkpoints, cada experto suele estar serializado como tensores separados: model.layers.3.mlp.experts.0.gate_proj.weight ... experts.255.gate_proj.weight.
En tiempo de ejecución, los kernels modernos esperan pesos empaquetados en un único tensor contiguo para poder usar grouped GEMMs o kernels fusionados.
Esa discrepancia entre checkpoint y runtime exige transformar pesos al cargar el modelo. La solución fue refactorizar la carga con una abstracción llamada WeightConverter.
¿Qué hace WeightConverter?
Piensa en la carga no como un simple copiado clave a clave, sino como una pipeline de conversión desde la estructura serializada del checkpoint hasta el layout que espera el runtime.
WeightConverter define patrones fuente → destino + operaciones. Operaciones primitivas (chunk, concatenate, etc.) son componibles. Dos operaciones claves para MoEs:
MergeModulelist: fusiona una lista de tensores en uno solo. Útil para apilar expertos.
SplitModulelist: divide un tensor en múltiples tensores. Útil para volver al formato por experto.
La carga además se planifica mejor: escanea claves una vez, agrupa por convertidor, materializa tensores con un thread pool y corre conversiones sólo cuando dependencias están listas. Eso reduce picos de memoria y evita múltiples escaneos.
Resultados prácticos: carga más rápida y cuantización integrada
Benchmarks comparando ramas v4 vs v5 de transformers muestran mejoras importantes en tiempos de carga para grandes modelos MoE (ejemplo con Qwen/Qwen1.5-110B-Chat en 1× A100 80GB):
v4 (device_map="auto", threadpool): 66.24 s
v5 (device_map="auto", async): 20.71 s
v5 (TP, async): 10.1 s
No es solo usar más hilos: la combinación de enrutamiento de una sola pasada, materialización asíncrona y planificación consciente de conversiones permite empaquetar expertos y fusionar proyecciones al cargar.
Además, la refactorización permite integrar la cuantización dentro de la misma pipeline de conversión. Cuantizar "por experto" sólo tiene sentido si los expertos están en un layout empaquetado y predecible.
Ejecución eficiente: backends de expertos
Empacar pesos es solo mitad del trabajo. En tiempo de inferencia, cada token se enruta a un subconjunto de expertos y hay que:
Despachar tokens a pesos de experto seleccionados.
Ejecutar proyecciones eficientemente.
Aplicar pesos de routing y reagrupar resultados en orden original.
La solución en Transformers fue un sistema de Experts Backend pluggable que desacopla la computación de los expertos de la implementación del modelo. Se selecciona el backend en tiempo de ejecución mediante un decorador @use_experts_implementation.
Backends actuales:
eager: bucle por experto; útil para debug y referencia de corrección.
batched_mm: duplica pesos seleccionados por token y usa torch.bmm; bueno para batch pequeño en GPUs con memoria disponible.
grouped_mm: ordena tokens por ID de experto y usa torch._grouped_mm; destaca en batches grandes o setups con memoria limitada.
Expert Parallelism: cómo escalar más allá de una GPU
MoEs pueden alcanzar cientos de miles de millones de parámetros. El truco es que cada token activa pocos expertos, así que no necesitas cargar todos los pesos en cada GPU. enable_expert_parallel=True activa el plan de expero-paralelismo (EP) y cambia el sharding.
Componentes clave:
GroupedGemmParallel: divide pesos de expertos por la dimensión de expertos (dim=0). Cada dispositivo carga num_experts / num_devices.
RouterParallel: remapea índices de expertos globales a locales, enmascara expertos no asignados al rank y usa un all-reduce para combinar salidas parciales.
model = AutoModelForCausalLM.from_pretrained("openai/gpt-oss-120b", dtype="auto", distributed_config=distributed_config)
Lanzamiento recomendado: torchrun --nproc-per-node N script.py, donde N divide uniformemente el número total de expertos.
Entrenamiento: todavía desafiante, pero con avances
Entrenar MoEs es más complejo que inferirlos: comunicación distribuida entre expertos, inestabilidades en routing, gran conteo de parámetros. La colaboración con Unsloth muestra optimizaciones importantes:
~12× más rápido en entrenamiento MoE
35% reducción de VRAM
~6× contexto más largo
12–30× speedup general frente a v4
Estas mejoras combinan el backend de expertos, estandarización en torch._grouped_mm y kernels Triton custom para grouped-GEMM y LoRA.
¿Qué significa esto para ti como desarrollador o responsable de producto?
Si construyes modelos o pipelines de inferencia a escala, MoEs te permiten mantener calidad con menor latencia e inversión en hardware.
Si trabajas en infra, presta atención a la carga de checkpoints y layout de pesos: la nueva pipeline de WeightConverter y la carga asíncrona cambian mucho la experiencia de traer modelos grandes a memoria.
Si investigas, hay espacio para nuevas estrategias de routing, medición de estabilidad y kernels optimizados que reduzcan comunicación y mejoren la eficiencia.
La adopción de MoEs ya no es solo una curiosidad de laboratorio: las librerías y herramientas están evolucionando para hacerlos prácticos en producción.
