CRAFT: Cost-aware Expert Replica Allocation with Fine-Grained Layerwise Estimations

O artigo apresenta o CRAFT, um framework eficiente de replicação de especialistas para modelos de linguagem do tipo MoE que maximiza o equilíbrio de carga sob restrições de memória através de estimativas granulares por camada, aumentando o throughput de serviço em até 1,2 vezes sem exigir alterações no modelo ou treinamento adicional.

O essencial

Imagine que você tem uma grande pizzaria (o modelo de Inteligência Artificial) que precisa atender milhares de pedidos ao mesmo tempo.

Para fazer isso rápido, a pizzaria não tem apenas um chef. Ela tem uma equipe gigante de especialistas (os "Mixture-of-Experts" ou MoE). Cada especialista é mestre em um tipo específico de pizza: um é o rei da pepperoni, outro é o mestre da pizza vegana, outro na de frutos do mar, e assim por diante.

O Problema: A Cozinha Caótica

Quando os clientes fazem pedidos, um "gerente de salão" (o roteador) decide qual especialista vai preparar cada pizza. O problema é que a maioria dos clientes pede as mesmas coisas (como pepperoni), enquanto poucos pedem as opções exóticas.

Isso cria um desequilíbrio:

• O Chef de Pepperoni está sobrecarregado, com uma fila enorme e demorando para atender.
• O Chef de Pizza Vegana está entediado, olhando para a parede, sem nada para fazer.
• O Chef de Frutos do Mar também está parado.

Na tecnologia atual, se você tentar distribuir os chefs entre várias cozinhas (GPUs) para equilibrar o trabalho, o Chef de Pepperoni ainda acaba sendo o gargalo. A pizzaria inteira fica parada esperando ele terminar, enquanto os outros chefs ficam ociosos. Isso é o desequilíbrio de carga.

A Solução Antiga: "Copiar e Colar" (Replicação Uniforme)

Para resolver isso, a solução antiga (chamada EPLB) era simples e burra: "Vamos ter uma cópia de TODOS os chefs em TODAS as cozinhas!".

• Se houver 100 tipos de pizza, você cria 100 cópias de cada chef em cada cozinha.
• O Resultado: O trabalho fica perfeitamente equilibrado. Ninguém espera.
• O Problema: A pizzaria precisa de muito mais espaço e equipamentos (memória da GPU). É como se você tivesse que alugar 100 cozinhas extras só para ter cópias de chefs que quase ninguém usa. Isso custa uma fortuna e faz a pizzaria ficar lenta porque o espaço está lotado de equipamentos inúteis.

A Inovação: CRAFT (O Gerente Inteligente)

Os autores deste artigo criaram o CRAFT. Pense no CRAFT como um gerente de salão superinteligente que usa um "orçamento de espaço" limitado.

Em vez de copiar todos os chefs, o CRAFT faz três coisas inteligentes:

1. Observa os Pedidos (Estimativa): Antes de abrir a pizzaria, ele analisa os pedidos das últimas semanas. Ele descobre exatamente quais chefs são os "queridinhos" (alta demanda) e quais são os "esquecidos" (baixa demanda).
2. Copia Apenas o Necessário (Alocação Fina):
   • Para o Chef de Pepperoni (que tem 100 pedidos), ele cria 8 cópias dele. Assim, o trabalho é dividido e a fila acaba rápido.
   • Para o Chef de Pizza Vegana (que tem apenas 2 pedidos), ele não cria nenhuma cópia. O único chef original dá conta do recado tranquilamente.
   • Para um chef intermediário, ele cria apenas 2 cópias.
3. Economiza Espaço: Como ele não copia os chefs que não precisam, sobra muito espaço na cozinha (memória da GPU).

Por que isso é incrível?

Esse espaço extra economizado é usado para algo muito importante: armazenar os ingredientes e as caixas de pizza prontas (o "KV Cache").

• Com a solução antiga (EPLB): A cozinha está tão cheia de cópias de chefs inúteis que não sobra espaço para as pizzas. O gerente tem que atender menos clientes ao mesmo tempo.
• Com o CRAFT: A cozinha está organizada. Sobrou espaço para preparar muitas pizzas ao mesmo tempo.

O Resultado na Vida Real

O artigo mostra que, ao usar o CRAFT em modelos gigantes (como o Kimi-K2 ou DeepSeek-R1):

• A pizzaria atende 14% a 20% mais clientes no mesmo tempo (aumento de throughput).
• O tempo de espera para o primeiro pedaço de pizza cai drasticamente.
• Tudo isso sem precisar comprar mais equipamentos ou gastar mais energia.

Resumo em uma frase

O CRAFT é como um gerente de restaurante que, em vez de contratar cópias de todos os funcionários (o que custa caro e ocupa espaço), contrata cópias extras apenas para os funcionários que estão sobrecarregados, deixando o resto da equipe trabalhar normalmente, resultando em um serviço mais rápido e barato.

Resumo técnico

1. O Problema

Os modelos de Linguagem de Grande Escala (LLMs) baseados em Mistura de Especialistas (MoE) tornaram-se a arquitetura padrão para escalar modelos mantendo custos computacionais por token constantes. No entanto, a implantação eficiente de MoE enfrenta um desafio crítico: o desequilíbrio de carga (load imbalance).

• Causa: O roteamento dinâmico dos tokens para especialistas (experts) segue uma distribuição de Zipf, onde poucos especialistas "quentes" (hot) recebem a maioria dos tokens, enquanto a maioria permanece "fria" (cold).
• Consequência: Em arquiteturas de Paralelismo de Especialistas (EP), isso gera gargalos em GPUs específicas, ociosidade em outras e congestionamento de rede durante as operações all-to-all.
• Solução Existente e Limitações: A técnica de replicação de especialistas (criar cópias de especialistas quentes em múltiplos dispositivos) é amplamente usada para equilibrar a carga. Contudo, métodos atuais (como o EPLB) utilizam uma replicação uniforme (ex: uma cópia por camada por GPU), o que consome memória de GPU excessiva. Isso leva a um trade-off negativo: o ganho de equilíbrio de carga diminui rapidamente após certo ponto, enquanto o custo de memória continua a subir, reduzindo o tamanho do KV Cache e degradando o throughput geral.

2. Metodologia: O Framework CRAFT

O CRAFT é um framework de alocação de réplicas de especialistas que visa maximizar o equilíbrio de carga sob um orçamento de memória estrito, utilizando estimativas finas e por camada.

Principais Observações que Guiaram o Design:

1. Retornos Decrescentes: O ganho de equilíbrio de carga diminui rapidamente à medida que o número de réplicas aumenta.
2. Heterogeneidade por Camada: Nem todas as camadas do modelo MoE se beneficiam igualmente da replicação. Camadas com alta assimetria (poucos especialistas dominam a carga) beneficiam-se muito, enquanto camadas já balanceadas beneficiam-se pouco.

Fluxo de Trabalho do CRAFT:

O sistema opera em três etapas principais:

1. Estimativa de Benefício por Camada (Offline):

   • Analisa a distribuição de carga de especialistas offline (reproduzindo lotes de inferência).
   • Calcula o "ganho de equilíbrio" (benefício de replicação) para diferentes contagens de réplicas em cada camada individualmente.
   • Constrói uma matriz de benefícios que mapeia quantas réplicas valem a pena para cada camada específica.

2. Alocação Otimizada de Réplicas (Otimização):

   • Formula o problema de alocação como um Problema da Mochila com Múltiplas Escolhas (MCKP).
   • O objetivo é selecionar o número ótimo de réplicas para cada camada de modo a maximizar o benefício total de equilíbrio de carga, respeitando o orçamento total de memória (número total de réplicas permitido).
   • Utiliza Programação Dinâmica para resolver isso eficientemente, permitindo que camadas críticas recebam mais réplicas e camadas balanceadas recebam nenhuma.

3. Atribuição e Posicionamento Consciente de Capacidade:

   • Para lidar com o fato de que diferentes camadas terão números diferentes de réplicas (desafio de balanceamento de memória entre GPUs), o CRAFT utiliza um algoritmo ganancioso intercalado.
   • Atribui especialistas e réplicas de forma a garantir que a capacidade de memória (número de especialistas) em cada GPU seja quase idêntica, evitando fragmentação de memória e garantindo tamanhos de KV Cache uniformes.

3. Contribuições Chave

• Caracterização Detalhada: Primeira análise detalhada de como o throughput e o equilíbrio de carga escalam com a replicação, demonstrando que a replicação uniforme é ineficiente devido aos retornos decrescentes.
• Framework CRAFT: Introdução de um sistema de alocação de réplicas end-to-end que é:
  • Consciente de Custos: Otimiza o uso de memória.
  • Granularidade Fina: Aloca réplicas por camada, não uniformemente.
  • Integrável: Pode ser inserido em frameworks de serviço existentes (como SGLang) sem necessidade de re-treinamento ou alterações no modelo.
• Algoritmo de Otimização: Desenvolvimento de uma solução baseada em Programação Dinâmica para resolver o problema de alocação de réplicas sob restrições de memória.

4. Resultados e Avaliação

O CRAFT foi avaliado em um cluster de 8 nós (AWS p4de.24xlarge) com modelos de grande escala (DeepSeek-R1-671B e Kimi-K2-1000B) e diversos conjuntos de dados.

• Aumento de Throughput: O CRAFT aumentou o throughput de serviço end-to-end em média 1,14x (chegando a 1,2x) em comparação com a técnica de replicação uniforme existente (EPLB).
• Eficiência de Memória: O CRAFT alcança ganhos de equilíbrio de carga comparáveis ao EPLB, mas utilizando 7,25x a 7,5x menos réplicas (dependendo do modelo). Isso libera memória para um KV Cache maior, permitindo maior concorrência de requisições.
• Latência (TTFT): Redução de 29% no Time-to-First-Token (TTFT) em comparação com a base sem replicação, mantendo-se competitivo com o EPLB.
• Robustez: O sistema demonstrou ganhos consistentes em diferentes tamanhos de cluster (6, 8 e 12 nós) e em conjuntos de dados com diferentes níveis de assimetria de carga.
• Custo de Overhead: O tempo de inicialização para estimar os benefícios é de aproximadamente 10 segundos (negligenciável) e não há overhead adicional durante a inferência.

5. Significado

O trabalho do CRAFT é significativo porque resolve o dilema fundamental da implantação de MoE em larga escala: como equilibrar a carga sem esgotar a memória da GPU.

Ao demonstrar que a replicação uniforme é ineficiente e propor uma alocação baseada em benefícios estimados por camada, o CRAFT permite que sistemas de inferência de LLMs operem com maior eficiência de recursos. Isso resulta em:

1. Redução de Custos Operacionais: Maior throughput por GPU disponível.
2. Escalabilidade: Possibilidade de implantar modelos com trilhões de parâmetros em clusters menores ou com maior eficiência.
3. Adoção Prática: Como não requer mudanças no modelo ou treinamento, pode ser adotado imediatamente em frameworks de produção modernos.

Em resumo, o CRAFT transforma a replicação de especialistas de uma estratégia "bruta" e custosa em uma técnica de otimização precisa e econômica, essencial para a próxima geração de modelos de IA eficientes.