A Outra Metade da Memory Wall: Servindo MoEs de 35B a Partir de SSD com Previsão de Roteamento Treinada

Guia de leitura

Problema e motivação

Modelos Mixture-of-Experts (MoE) escalam o total de parâmetros mantendo os parâmetros ativos baixos por token (por exemplo, ativando ~3B de parâmetros por token em um modelo de 35B). No entanto, a inferência em hardware de consumo continua limitada pela memória estática de pesos: um MoE de 35B requer cerca de 19,5 GB mesmo quando quantizado para int4. Em uma máquina de 24 GB, isso sobrecarrega completamente o sistema operacional e a KV cache, levando a um severo paging de memória e degradação da taxa de transferência de decodificação (por exemplo, 3,9 tok/s para um modelo residente em um Mac mini M4 Pro).

O offloading de pesos de especialistas para o SSD reduz a pegada de memória, mas o carregamento sob demanda ingênuo introduz um gargalo rígido de serialização: a seleção de especialistas para a camada $N+1$ depende da saída da camada $N$. As leituras do SSD não conseguem começar cedo o suficiente para se ocultarem sob a computação, paralisando a execução na latência do disco uma vez por camada a cada token.

Ideias centrais

• SSD como Nível de Pesos: Os pesos dos especialistas de MoE permanecem no SSD no formato int4 e são mapeados em memória (mmap) em um pool de streaming compacto. O pico de memória é determinado pelo conjunto de trabalho ativo, e não pelo tamanho total dos parâmetros do modelo.
• Previsão como Roteamento: Em vez de esperar que a camada $N$ termine, um prerouter por camada prevê o roteamento da camada $N+1$ com um token inteiro de antecedência. Fundamentalmente, o modelo trata a previsão como o próprio roteamento no momento da decodificação. Isso elimina tokens descartados ou carregamentos de fallback por construção, porque o conjunto preparado é identicamente igual ao conjunto roteado.
• LoRA de Recuperação Não Fundido: A quantização de 4 bits e o roteamento aproximado causam perda de qualidade. O Edge0 treina um adaptador de baixo posto (low-rank adapter) no caminho de roteamento do aluno via destilação. Como a fusão do LoRA na base int4 e a consequente requantização apagam a maior parte do delta (apenas 18% de retenção no nível de logits sobrevive), o Edge0 serve o adaptador como um delta paralelo não fundido com sobrecarga insignificante.

Método

Arquitetura da Cabeça do Prerouter

Cada cabeça de prerouter pertencente à camada $N$ é executada no token $t$ sobre a saída da normalização de camada pós-atenção e prevê o roteamento para a camada $N+1$ no token $t+1$. As características de entrada concatenam o estado oculto com dois vetores one-hot de roteamento top-$k$ (token atual e anterior):

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

Para a classe 35B, isso resulta em $2048+2\times 256=2560$. A cabeça computa:

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

onde $\ell (x)$ é um caminho residual linear inicializado a partir dos pesos do roteador da próxima camada.

Execução do Adaptador Não Fundido

A camada MoE servida preserva os bytes base quantizados e executa o delta LoRA paralelo:

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

Treinamento em Três Fases

1. Destilação da Cabeça: As cabeças do prerouter são treinadas para imitar as decisões reais do roteador da próxima camada do professor, mantendo o restante da rede congelado.
2. SFT no Caminho do Aluno: Adaptadores de recovery LoRA (anexados às projeções de atenção, atenção linear e especialistas compartilhados, mas não aos especialistas em streaming) são treinados usando texto gerado pelo professor enquanto o roteamento é forçado através do prerouter.
3. Destilação On-Policy: O aluno gera texto sob execução pré-roteada, o professor base fp16 pontua as gerações e os parâmetros do LoRA são otimizados via divergência reversa de Kullback–Leibler com um objetivo de probabilidade top-$k$ sensível à cauda.

Resultados

• Taxa de Transferência e Memória (Mac mini M4 Pro, 24 GB):
  • edge0-35b ($K=4$): Atinge 20,4 tok/s de decodificação usando apenas 2,9 GiB de pico de memória ativa a partir de um checkpoint int4 de 19,5 GB (em comparação com 3,9 tok/s e 18,2 GiB quando totalmente residente).
  • edge0-8b ($K=8$): Atinge 28,0 tok/s de decodificação dentro de 1,5 GiB de pico de memória.
• Ganho de Velocidade com Prefetching (MacBook M2, 16 GB):
  • Em comparação com o streaming sob demanda puramente reativo, o prefetch do prerouter de 1 token aumenta a velocidade de decodificação em +80% a +84% em diferentes larguras de roteamento ($K=2$: 4,8 para 8,6 tok/s; $K=4$: 3,5 para 6,4 tok/s; $K=8$: 1,8 para 3,3 tok/s).
  • O tempo de carregamento bloqueado na thread principal cai drasticamente (por exemplo, de 244,0 ms para 101,9 ms por passo com $K=4$).
• Qualidade do Modelo (OpenCompass):
  • Em cinco benchmarks (AIME 2026, HumanEval, GPQA-Diamond, MMLU-Pro, IFBench), o edge0-35b pontua uma média de 79,2 contra 83,2 da base Qwen3.6 fp16 (uma diferença de 3,9 pontos).
  • O edge0-8b atinge 69,9 contra 72,7 da base Ling 3.0 tiny fp16 (uma diferença de 2,8 pontos).

Limitações e questões em aberto

• Concorrência de Serviço: O Edge0 opera serializado em FIFO para requisições únicas; o processamento em lote (batching) altera o conjunto de trabalho dos especialistas e a dinâmica de cache de maneiras não contempladas pelos perfis de fluxo único.
• Sobrecarga de Construção de Grafos na CPU: A velocidade de decodificação é limitada pela sobrecarga do framework no lado da CPU (~44 ms/passo na construção de grafos MLX em 40 camadas), o que as otimizações de SSD não conseguem aliviar.
• Queda de Qualidade em Raciocínio: A degradação de desempenho concentra-se em benchmarks de raciocínio complexo (por exemplo, uma queda de 6,1 pontos no AIME para o modelo 35B e de 10,0 pontos para o de 8B).
• Dependência de Hardware: Os ganhos relativos de taxa de transferência do prerouter dependem da latência do disco e diminuem quando os caches de pesos estão quentes ou ao rodar em camadas de armazenamento ultrarrápidas.