La otra mitad del muro de la memoria: Servir MoEs de 35B desde SSD con predicción de enrutamiento entrenada

Guía de lectura

Problema y motivación

Los modelos Mixture-of-Experts (MoE) escalan los parámetros totales manteniendo un número bajo de parámetros activos por token (por ejemplo, activando ~3B de parámetros por token en un modelo de 35B). Sin embargo, la inferencia en hardware de consumo sigue limitada por la memoria estática de los pesos: un MoE de 35B requiere aproximadamente 19.5 GB incluso cuantizado a int4. En una máquina de 24 GB, esto satura por completo el espacio para el sistema operativo y la caché KV, provocando una severa paginación de memoria y degradando el rendimiento de decodificación (por ejemplo, 3.9 tok/s para un modelo residente en una Mac mini M4 Pro).

Descargar los pesos de los expertos al SSD reduce la huella de memoria, pero la carga bajo demanda ingenua introduce un cuello de botella de serialización estricto: la selección de expertos para la capa $N+1$ depende de la salida de la capa $N$. Las lecturas desde el SSD no pueden comenzar con suficiente anticipación para ocultarse tras el cómputo, pausando la ejecución por la latencia del disco una vez por capa por cada token.

Ideas clave

• SSD como nivel de almacenamiento de pesos: Los pesos de los expertos del MoE permanecen en el SSD en formato int4 y se mapean en memoria (mmap) dentro de un pool de transmisión compacto. El pico de memoria lo determina el conjunto de trabajo activo y no el tamaño total de los parámetros del modelo.
• Predicción como enrutamiento: En lugar de esperar a que termine la capa $N$, un prerouter por capa predice el enrutamiento de la capa $N+1$ con un token completo de antelación. De forma crucial, el modelo trata la predicción como el enrutamiento real en el momento de la decodificación. Esto elimina por diseño los tokens descartados o las cargas de respaldo (fallback), ya que el conjunto preparado en memoria es idéntico al conjunto enrutado.
• Recovery LoRA no fusionado: La cuantización a 4 bits y el enrutamiento aproximado generan una pérdida de calidad. Edge0 entrena un adaptador de bajo rango sobre la ruta de enrutamiento del estudiante mediante destilación. Como fusionar el LoRA en la base int4 y volver a cuantizar elimina la mayor parte del delta (solo sobrevive un 18% de retención a nivel de logits), Edge0 sirve el adaptador como un delta paralelo no fusionado con un costo marginal insignificante.

Método

Arquitectura del cabezal Prerouter

Cada cabezal prerouter perteneciente a la capa $N$ se ejecuta en el token $t$ sobre la salida de layer norm posterior a la atención y predice el enrutamiento para la capa $N+1$ en el token $t+1$. Las características de entrada concatenan el estado oculto con dos vectores one-hot top-$k$ de enrutamiento (token actual y anterior):

$$\text{dim}=d_{\text{model}}+2E$$

Para el modelo de 35B, esto equivale a $2048+2\times 256=2560$. El cabezal calcula:

$$\text{head}(x)={\mathrm{fc}}_2(\mathrm{erf}\text{-}\mathrm{gelu}({\mathrm{fc}}_1(x)))+\ell (x)$$

donde $\ell (x)$ es una ruta residual lineal inicializada con los pesos del enrutador de la siguiente capa.

Ejecución del adaptador no fusionado

La capa MoE servida conserva los bytes base cuantizados y ejecuta el delta LoRA paralelo:

$$y=W_{int4}(x)+\frac{\alpha }{r}BAx$$

Entrenamiento en tres fases

1. Destilación del cabezal: Los cabezales prerouter se entrenan para imitar las decisiones del enrutador real de la siguiente capa del profesor, manteniendo congelado el resto de la red.
2. SFT en la ruta del estudiante: Los adaptadores de recovery LoRA (incorporados en las proyecciones de atención, atención lineal y expertos compartidos, pero no en los expertos transmitidos) se entrenan utilizando texto generado por el profesor mientras se fuerza el enrutamiento a través del prerouter.
3. Destilación On-Policy: El estudiante genera texto bajo ejecución prerouteada, el profesor base en fp16 evalúa las generaciones y los parámetros de LoRA se optimizan mediante divergencia inversa de Kullback–Leibler con un objetivo de probabilidad top-$k$ sensible a las colas (tail-aware).

Resultados

• Rendimiento y memoria (Mac mini M4 Pro, 24 GB):
  • edge0-35b ($K=4$): Logra una decodificación de 20.4 tok/s usando solo 2.9 GiB de memoria activa máxima a partir de un checkpoint int4 de 19.5 GB (frente a 3.9 tok/s y 18.2 GiB cuando es completamente residente).
  • edge0-8b ($K=8$): Logra 28.0 tok/s de decodificación dentro de un pico de memoria de 1.5 GiB.
• Aceleración por precarga (MacBook M2, 16 GB):
  • En comparación con la transmisión puramente reactiva bajo demanda, la precarga con prerouter de 1 token incrementa la velocidad de decodificación entre un +80% y un +84% en distintos anchos de enrutamiento ($K=2$: de 4.8 a 8.6 tok/s; $K=4$: de 3.5 a 6.4 tok/s; $K=8$: de 1.8 a 3.3 tok/s).
  • El tiempo de carga que bloquea el hilo principal se reduce drásticamente (por ejemplo, de 244.0 ms a 101.9 ms por paso con $K=4$).
• Calidad del modelo (OpenCompass):
  • A lo largo de cinco benchmarks (AIME 2026, HumanEval, GPQA-Diamond, MMLU-Pro, IFBench), edge0-35b obtiene un promedio de 79.2 frente a 83.2 de la base fp16 Qwen3.6 (una diferencia de 3.9 puntos).
  • edge0-8b alcanza 69.9 frente a 72.7 de la base fp16 Ling 3.0 tiny (una diferencia de 2.8 puntos).

Limitaciones y preguntas abiertas

• Concurrencia en el servicio: Edge0 opera serializado por FIFO para solicitudes individuales; el procesamiento por lotes (batching) altera el conjunto de trabajo de los expertos y la dinámica de caché de formas no analizadas en los perfiles de flujo único.
• Sobrecarga de construcción de grafos en CPU: La velocidad de decodificación está limitada por la sobrecarga del framework en la CPU (~44 ms/paso en la construcción del grafo MLX a través de 40 capas), lo cual no puede ser mitigado por optimizaciones de SSD.
• Caída de calidad en razonamiento: La degradación del rendimiento se concentra en evaluaciones de razonamiento complejo (por ejemplo, una caída de 6.1 puntos en AIME para el modelo de 35B y de 10.0 puntos para el de 8B).
• Dependencia del hardware: Las ganancias relativas de rendimiento del prerouter dependen de la latencia del disco y disminuyen cuando las cachés de pesos están calientes o al ejecutarse en niveles de almacenamiento ultrarrápidos.