A portabilidade do modelo Gemma-4 31B do Google para um AWS Inferentia2 inf2.24xlarge produziu uma correspondência perfeita token por token em relação à referência de CPU — no entanto, cada frase gerada era um amontoado de palavras sem sentido. A lacuna entre "corresponder" e "funcionar" serve agora como um alerta para quem tenta espremer LLMs massivos nos chips de inferência customizados da Amazon.

Por que uma correspondência token por token não é suficiente

O desenvolvedor comparou cada token de saída do dispositivo Inferentia com o token produzido por uma execução do modelo em CPU. Os fluxos eram idênticos, portanto, o hardware parecia ter reproduzido a implementação de referência exatamente. Na realidade, ambos os fluxos alimentaram um prompt malformado em um modelo desprovido de seu chat template e fornecido com os marcadores de turno incorretos. A ausência do template enviou o modelo para um loop infinito, gerando algo sem sentido. O hardware fez o seu trabalho — ele reproduziu um bug que já existia no código de referência.

A lição é simples: SEQ_MATCH (igualdade sequencial de tokens) não é sinônimo de correção. Se a implementação de referência estiver quebrada, uma réplica fiel de hardware herdará a mesma falha. A validação deve ir além da paridade no nível de token; ela precisa de verificações funcionais de ponta a ponta com entradas formatadas corretamente.

Buffers mascarados como parâmetros

Durante a fase de carregamento, o carregador do modelo ignorou um componente chamado layer_scalar. O código registrou esse objeto como um buffer em vez de um parameter na definição do modelo PyTorch. Buffers são tensores estáticos que o treinamento não atualiza, e muitos carregadores os ignoram ao converter para formatos compatíveis com Neuron. Ao ignorá-lo, os fatores de escala de várias camadas permaneceram em seus valores padrão, distorcendo os cálculos em toda a rede. Nenhum erro foi gerado; o modelo compilou e o pipeline de inferência foi executado, mas os resultados numéricos estavam incorretos.

Para quem estiver movendo modelos grandes para o Inferentia, audite cada tensor que não seja um parâmetro. Mesmo que um tensor não deva ser aprendido, ele ainda pode ser essencial para o cálculo correto do forward-pass. Verificar manualmente a inclusão de buffers pode evitar erros de escala silenciosos que, de outra forma, seriam difíceis de diagnosticar.

Volatilidade de instâncias spot e a compilação de 39 minutos

Executar um modelo de 31 bilhões de parâmetros em uma instância spot parece barato, mas a economia vem acompanhada de eventos de recuperação imprevisíveis. O tempo de compilação do desenvolvedor — cerca de 39 minutos para traduzir o modelo em código compatível com Neuron — desapareceu quando a AWS recuperou a instância. Para sobreviver às interrupções, ele construiu uma rede de segurança de três frentes:

  • ModelBuilder manteve o uso de memória dentro do limite de 384 GB do host, evitando travamentos que forçariam um reinício.
  • O espelhamento imediato no S3 tanto dos arquivos de pesos brutos quanto dos "neffs" compilados (arquivos executáveis Neuron) permitiu que uma nova instância retomasse exatamente de onde a anterior parou.
  • Um poller multi-região varreu as regiões da AWS em busca de capacidade spot disponível e lançou uma nova instância assim que uma apareceu.

Essas etapas transformaram uma compilação frágil de ponto único em um pipeline resiliente que sobrevive à volatilidade dos mercados spot.

Armadilhas de sharding com layouts de atenção mistos

O Gemma-4 31B utiliza duas configurações de atenção. Algumas camadas empregam quatro cabeças de chave-valor (KV), outras utilizam uma contagem diferente. Dividir o modelo uniformemente entre oito ranks paralelos falha quando a contagem de cabeças KV de uma camada não é divisível de forma exata. Tentar fazer o sharding de uma camada de 4 cabeças em oito ranks forçaria cada rank a lidar com meia cabeça — uma impossibilidade matemática que desencadeia incompatibilidades de formato (shape mismatches) e erros de tempo de execução.

A solução foi replicar as camadas com sharding global (aquelas com contagens de cabeças compatíveis) em todos os ranks e fazer o sharding apenas das camadas "sliding", cujas contagens de cabeças permitiam uma divisão uniforme. Essa estratégia híbrida manteve a eficiência do tensor-parallel, evitando a divisão ilegal de cabeças KV e eliminando os erros de paralelização de tensores que assombraram as tentativas anteriores.

Conclusão

Portar um LLM gigante para o Inferentia é mais do que um exercício de "compilar e executar". Exige testes funcionais rigorosos além da igualdade de tokens, verificação meticulosa de que cada tensor — parâmetro ou buffer — é tratado corretamente e uma estratégia de implantação que antecipe a recuperação de instâncias spot. Por fim, o sharding deve respeitar a geometria de atenção interna do modelo; caso contrário, o paralelismo que promete velocidade torna-se uma fonte de falha silenciosa.