Mistura de especialistas divide em dois o preço de um parâmetro

O modelo Mixtral, da Mistral, lançado em janeiro de 2024, tem oito blocos feed-forward em cada camada e usa dois deles para cada token. Sua contagem de parâmetros vem em duas cifras: 47 bilhões no total, 13 bilhões ativos durante a inferência. Nos benchmarks do artigo, igualou ou superou o Llama 2 70B e o GPT-3.5, com ganhos especialmente grandes em matemática, geração de código e tarefas multilíngues. Onze meses depois, o DeepSeek-V3 levou a razão muito mais longe. Tem 671 bilhões de parâmetros, ativa 37 bilhões a cada token, foi treinado com 14,8 trilhões de tokens e consumiu 2,788 milhões de horas de GPU H800 para isso.

Os dois modelos trazem dois números nas especificações, e a distância entre eles é o assunto deste ensaio. Num modelo denso, um número cobre tudo: a contagem de parâmetros determina quanta aritmética cada token custa e quanta memória o modelo ocupa. Um modelo de mistura de especialistas separa as duas coisas. A aritmética segue os parâmetros ativos. A memória segue o total. Quase tudo o que surpreende na economia desses modelos, tanto os treinos baratos quanto os custos de serviço por vezes decepcionantes, vem dessa separação.

Uma ideia antiga que esperou o hardware

A ideia remonta a Noam Shazeer e colegas em 2017, entre eles Geoffrey Hinton e Jeff Dean. Eles inseriram, entre as camadas de uma LSTM, uma camada com até milhares de pequenas redes feed-forward, os especialistas, e treinaram uma rede de porta para escolher, a cada entrada, uma combinação esparsa deles. O maior modelo tinha 137 bilhões de parâmetros, e eles relataram um aumento de capacidade de mais de mil vezes com perdas apenas pequenas de eficiência computacional. Em 2021, os Switch Transformers, de William Fedus, Barret Zoph e Shazeer, simplificaram o roteamento para que cada token vá a um único especialista, treinaram modelos de até um trilhão de parâmetros e relataram uma aceleração de quatro vezes em relação ao T5-XXL.

Durante vários anos a abordagem ficou quase só dentro dos grandes laboratórios. Era difícil de treinar com estabilidade e difícil de servir, por razões que ficarão claras. Em 2024, a engenharia havia amadurecido o suficiente para que projetos MoE se tornassem comuns entre os maiores lançamentos de pesos abertos. O padrão é o mesmo em toda parte: em cada bloco Transformer, a camada de atenção permanece como está, e a rede feed-forward única é substituída por um banco de especialistas e um pequeno roteador que escolhe alguns deles por token.

A aritmética segue os parâmetros ativos

Um transformer denso gasta cerca de duas operações de ponto flutuante por parâmetro em cada token que processa. Um modelo denso de 70 bilhões de parâmetros custa, portanto, uns 140 gigaflops por token. O Mixtral gasta uns 26 gigaflops, já que 13 bilhões de parâmetros estão ativos. O DeepSeek-V3 gasta uns 74 gigaflops por token, onde um modelo denso de 671 bilhões gastaria 1,34 teraflops. É um fator de 18 em computação para a mesma capacidade total.

O treino segue o mesmo padrão. Com a aproximação usual de seis operações por parâmetro por token, a execução do DeepSeek-V3 com 14,8 trilhões de tokens e 37 bilhões de parâmetros ativos soma cerca de 3,3 × 10²⁴ operações. Um modelo denso de 671 bilhões nos mesmos dados precisaria de cerca de 6,0 × 10²⁵, dezoito vezes mais. O relatório faz a sua própria contabilidade de custo: 2,788 milhões de horas de GPU a um preço de aluguel presumido de US$ 2 cada, ou US$ 5,576 milhões, valor que cobre a execução oficial de treino e, como diz o relatório, exclui os custos de pesquisa prévia e de experimentos de ablação. Essas ressalvas importam, já que os experimentos que levaram a um projeto são parte real do seu custo, mas a aritmética acima mostra por que uma execução desse porte se tornou viável.

Uma verificação cruzada com um modelo denso confirma a aritmética. A ficha do modelo Llama 3.1 405B, da Meta, um modelo denso treinado com cerca de 15 trilhões de tokens, lista 30,84 milhões de horas de GPU H100. As 2,788 milhões de horas de GPU H800 do DeepSeek-V3 são cerca de onze vezes menos. A regra das seis operações prevê a mesma razão só a partir dos tamanhos dos modelos: 6 × 405 bilhões × 15 trilhões dá cerca de 3,6 × 10²⁵ operações, contra 3,3 × 10²⁴ do DeepSeek-V3, um fator de 11,1. Os dois chips entregam vazão aritmética parecida, de modo que a economia em horas de GPU é a economia em parâmetros ativos, quase até a casa decimal.

Além do custo, o projeto traz um argumento de qualidade. O resultado do Mixtral sugere que espalhar conhecimento por muitos blocos especializados e consultar apenas alguns deles pode igualar um modelo denso com cinco vezes os parâmetros ativos. Apesar do nome, os especialistas não são especialistas no sentido humano. Os autores do Mixtral analisaram o roteador e não encontraram uma atribuição óbvia de especialistas por tema, com algumas regularidades que seguem a sintaxe. Chamá-los de especialistas é uma metáfora, e o benefício se descreve melhor como capacidade que não custa nada até ser usada.

Mais especialistas, menores

O Mixtral usa 8 especialistas por camada e escolhe 2, o que permite 28 combinações distintas para cada token em cada camada. O DeepSeek-V3 divide a mesma capacidade feed-forward em 256 especialistas roteados e escolhe 8, o que permite cerca de 4,1 × 10¹⁴ combinações, além de um especialista compartilhado pelo qual todo token passa. O raciocínio por trás da granularidade mais fina é que um roteador com mais opções, e menores, pode compor uma mistura mais específica para cada token, enquanto o especialista compartilhado guarda o conhecimento de que quase todo token precisa, o qual de outro modo seria copiado em muitos dos roteados. Se o maior número de combinações é bem aproveitado é uma questão empírica, e os resultados do modelo sugerem que sim. A granularidade mais fina tem um preço no serviço, porque espalha o trabalho de um token por mais destinos, como as duas próximas seções mostram.

Os pesos ainda precisam morar em algum lugar

Todo especialista precisa estar disponível em todo passo, porque o roteador decide de quais um token precisa só depois que o token chega. Todos os 671 bilhões de parâmetros do DeepSeek-V3 precisam, portanto, estar residentes na memória do acelerador. Com precisão de 8 bits são 671 GB, que não cabem nos 640 GB de um nó com oito GPUs de 80 GB. Cabem num nó com oito placas de 141 GB (1.128 GB ao todo), sobrando uns 450 GB para caches e ativações. Um modelo denso de 37 bilhões de parâmetros caberia numa única GPU com 8 bits, com folga.

É o primeiro lugar em que um comprador encontra a divisão. Um modelo vendido como "37B ativos" é precificado e avaliado como um modelo de porte médio e hospedado como um modelo muito grande. Para quem roda modelos de pesos abertos em hardware próprio, seja por soberania de dados ou para controlar custo em alto volume, o hardware precisa ser dimensionado pela contagem total de parâmetros. O orçamento de computação sozinho leva a um nó que não comporta o modelo.

A decodificação encolhe a economia

O segundo lugar é mais sutil. Lembre que, durante a geração, o hardware passa a maior parte do tempo esperando pela memória: cada passo precisa ler os pesos que participam daquele passo. Para um modelo denso, são todos os pesos, uma vez por passo, não importa quantas conversas dividam o passo. Para uma mistura de especialistas, um passo lê os especialistas para os quais pelo menos um token do lote foi roteado.

Um cálculo pequeno mostra com que rapidez esse conjunto cresce. Suponha que cada token escolha seus especialistas de modo uniforme ao acaso, o que exagera a dispersão, já que roteadores reais têm favoritos. O Mixtral escolhe 2 de 8 especialistas por camada. A chance de um dado especialista não ser tocado por um lote de B tokens é 0,75^B, de modo que a fração de especialistas lidos é 1 − 0,75^B. Com B = 4 são 68%, com B = 8 são 90% e com B = 16 são 99%. Quando dezesseis conversas dividem um passo, o servidor lê praticamente todos os 47 bilhões de parâmetros por passo em vez de 13 bilhões, e o custo de largura de banda da decodificação passa a parecer o de um modelo denso de 47 bilhões. Um modelo denso de 13 bilhões atendendo as mesmas dezesseis conversas leria 3,6 vezes menos bytes.

O DeepSeek-V3 escolhe 8 de 256 especialistas roteados por camada, mais um especialista compartilhado que todo token usa. Um especialista roteado não é tocado com probabilidade (1 − 8/256)^B. Com B = 32, cerca de 64% dos especialistas são lidos; com B = 128, cerca de 98%; com B = 256, praticamente todos. A economia que parecia um fator de 18 no papel existe no regime limitado pela memória apenas quando o lote é pequeno o bastante para que a maioria dos especialistas fique ociosa, e lotes pequenos são exatamente onde o hardware é menos eficiente.

As duas fases do serviço se comportam de modo diferente aqui. O prefill processa um prompt inteiro numa passada, com milhares de tokens fluindo por cada camada, então todo especialista recebe bastante trabalho e o passo é limitado pela aritmética e não pela memória. Nesse caso a esparsidade rende por inteiro, e um prompt longo custa o que os parâmetros ativos dizem que deve custar. É na decodificação que a conta de memória reaparece. Uma carga de trabalho dominada por prompts longos e respostas curtas, como classificar ou resumir documentos, combina por isso melhor com esses modelos do que uma dominada por gerações longas para um punhado de usuários.

A saída é fazer cada GPU guardar apenas alguns especialistas e enviar cada token para onde moram os seus especialistas. Isso se chama paralelismo de especialistas e transforma o problema em um de lotes muito grandes espalhados por muitas GPUs, de modo que cada GPU leia só os seus poucos especialistas e ainda receba muitos tokens. O relatório da DeepSeek descreve a escala que isso exige. O estágio de prefill usa uma unidade mínima de implantação de 4 nós e 32 GPUs, e o estágio de decodificação usa 40 nós e 320 GPUs, com cada GPU hospedando um especialista e algumas hospedando cópias redundantes dos mais populares para equilibrar a carga. É uma configuração de serviço para um provedor com tráfego suficiente para encher centenas de GPUs. Não é uma configuração que uma empresa com vinte usuários simultâneos consiga reproduzir, e para ela a promessa de custo baixo por token não se transfere.

O que isso significa para escolher e precificar

Algumas regras práticas decorrem daí.

Primeiro, leia os dois números separadamente e calcule seus próprios custos a partir deles. Os parâmetros ativos governam a computação por token e, portanto, a conta de treino e a velocidade de um lote grande. Os parâmetros totais governam a memória e, portanto, o hardware mínimo.

Segundo, faça benchmark na sua concorrência. Os números de vazão publicados para modelos MoE costumam ser medidos com lotes grandes em clusters grandes. Um teste sensato é rodar a sua própria carga, com o seu número esperado de requisições simultâneas, e medir latência e tokens por segundo. Um modelo que parece quatro vezes mais barato com lote 256 pode parecer não mais barato que uma alternativa densa com lote 8.

Terceiro, entenda por que modelos MoE hospedados podem ter preço baixo. Um provedor com tráfego grande e estável consegue encher os lotes que tornam o paralelismo de especialistas eficiente. O preço que você paga reflete a escala dele, além da arquitetura. Isso significa também que um modelo MoE hospedado e um modelo auto-hospedado não são o mesmo objeto econômico, mesmo com pesos idênticos.

Quarto, observe os riscos do lado do treino, que aparecem como diferenças de comportamento sob carga. Roteadores tendem a colapsar em alguns especialistas favoritos a menos que algo os force a distribuir a carga. A correção usual é uma perda auxiliar que penaliza o desequilíbrio, o que também pode degradar a qualidade, e o relatório do DeepSeek-V3 descreve uma alternativa que chama de balanceamento de carga sem perda auxiliar, apresentada como inédita. Que um dado modelo se equilibre bem afeta como ele se comporta sob tráfego real.

Quando um modelo denso ainda é a melhor resposta

O desenho esparso não é uma melhoria universal. Se uma carga de trabalho roda com baixa concorrência, digamos um assistente interno usado por algumas dezenas de pessoas, um modelo denso com o mesmo tamanho ativo pode ser mais barato de hospedar e mais simples de operar, porque precisa de menos memória e não tem variação de roteamento. Se a memória é a restrição, como em dispositivos de borda ou servidores locais modestos, os parâmetros ativos são a régua errada e os totais são a certa.

O inverso também vale. Se o tráfego de um provedor é grande e estável, a mistura de especialistas lhe permite oferecer uma qualidade que de outro modo custaria muito mais, e é o desenho por trás de vários dos maiores modelos de pesos abertos. A resposta razoável de um comprador é parar de tratar a contagem de parâmetros de manchete como medida de custo e pedir os dois números, a concorrência em que a vazão informada foi medida e o hardware que essa vazão supôs.

O primeiro passo antes de qualquer compra é uma tabela com uma linha por modelo candidato e quatro colunas: parâmetros totais, parâmetros ativos, memória dos pesos na precisão que você vai usar e tokens por segundo medidos na sua própria concorrência. A tabela leva um dia para ser preenchida e vai ordenar os candidatos de modo diferente dos números de marketing.

Fontes

1. Jiang, A. Q., Sablayrolles, A., Roux, A., et al. (2024). Mixtral of Experts. Mistral AI. arXiv:2401.04088.
2. DeepSeek-AI (2024). DeepSeek-V3 Technical Report. arXiv:2412.19437.
3. Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., Dean, J. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. ICLR 2017. arXiv:1701.06538.
4. Fedus, W., Zoph, B., Shazeer, N. (2022). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. Journal of Machine Learning Research 23. arXiv:2101.03961.
5. NVIDIA. Folhas de dados das GPUs H100 e H200 Tensor Core (capacidade de memória: 80 GB e 141 GB).
6. Meta (2024). Ficha do modelo Llama 3.1, tempo de treino em horas de GPU. github.com/meta-llama/llama-models.