TorchNeuron e TorchTPU: como PyTorch roda hoje em chips da Amazon e do Google

O que mudou: PyTorch ganha ponte oficial para Trainium e TPU

Até agora, quem escrevia um modelo em PyTorch (o framework de código aberto mais usado para treinar e rodar redes neurais) contava com otimizações pensadas para GPU da Nvidia. Em 19 e 20 de outubro de 2026, Amazon e Google lançaram, com um dia de diferença, duas ferramentas que prometem rodar esse mesmo código nos próprios chips de IA: TorchNeuron, da AWS, para os chips Trainium, e TorchTPU, do Google, para os TPU. A informação é da ActuIA.

As duas peças substituem uma ponte mais antiga, o PyTorch/XLA, e passam por uma interface que o PyTorch reserva a hardware de terceiros, chamada PrivateUse1. O TorchNeuron está em beta pública, sob licença Apache 2.0, desde 20 de outubro, e cita oficialmente apenas as instâncias Trn2 e Trn3, com Python 3.12 a 3.14. O TorchTPU chegou na versão 0.2.0, mesma licença, e sua documentação valida quatro gerações de chip: TPU v5e, v5p, v6e e v7 (Ironwood), com v8i e v8t ainda em mira.

O que roda sem reescrever (e o que ainda exige kernel novo)

Segundo a AWS, mover um modelo para o chip Neuron e compilar com torch.compile usando o backend do mesmo nome costuma bastar. Código escrito com a API torch.accelerator roda sem alteração, de acordo com o blog técnico publicado no site do PyTorch. A Hugging Face informou que 89% das 100 arquiteturas de geração de texto mais usadas já funcionam de imediato com a biblioteca Transformers, ante 46% na biblioteca anterior, Optimum Neuron 0.4.5, que a própria empresa recomenda abandonar.

No lado do Google, basta o argumento device_map="tpu" com a versão 5 da Transformers, segundo o billet da empresa. A compilação ainda passa por OpenXLA e StableHLO, não pelo Inductor que o PyTorch usa nativamente.

O problema aparece quando o modelo depende de kernel personalizado, aquele trecho de código de baixo nível escrito direto para acelerar uma operação específica. Nenhum dos dois comunicados prevê rodar kernels escritos em CUDA ou Triton (as linguagens usadas para GPU Nvidia): eles precisam ser reescritos nas ferramentas próprias de cada fabricante, NKI na Amazon e Pallas ou Tokamax no Google. O Helion, linguagem de kernel portátil que o próprio PyTorch desenvolve, ainda é classificado como experimental na documentação do TorchTPU.

As limitações documentadas pela Amazon incluem:

• torch.compile só aceita formas fixas de entrada: cada formato novo dispara uma recompilação
• cálculos em float64 rodam em float32
• quantização FP8 pela TorchAO segue em beta privada
• cerca de 22 mil dos testes padrão do PyTorch falham por operadores sem suporte nativo, incluindo células LSTM e GRU fundidas
• tensores esparsos, NestedTensor, máscaras de atenção arbitrárias, quantização por canal e formas dinâmicas ficaram fora da primeira versão por falta de demanda dos clientes

O Google não lista limites no próprio anúncio, mas a documentação do TorchTPU marca como experimentais inteiros simbólicos, tensores irregulares, LSTM, GRU, modelos Mamba e FSDPv2, com a ressalva de baixa performance. Operações dinâmicas ainda disparam recompilação lá também.

O caso Deepgram: meses de trabalho para um ganho que ninguém mediu de fora

O único cliente que detalhou publicamente uma migração completa foi a Deepgram, com seu modelo de transcrição Nova-3. A equipe primeiro rodou o modelo em modo eager (sem compilação prévia) no Trainium2 para identificar trechos que caíam de volta no processador, depois compilou com torch.compile até eliminar esses desvios por completo.

O trabalho, feito com a equipe de engenharia aplicada da AWS, levou vários meses, não dias. Parte das camadas foi trocada por kernels prontos da biblioteca NKILib; outra parte precisou ser reescrita do zero, incluindo uma convolução 2D que a NKILib simplesmente não tinha. Essa convolução nova chegou a rodar até 50 vezes mais rápido que a versão de referência em lotes pequenos. A Deepgram também construiu, em Rust, um motor de execução sem Python e sem PyTorch, que ganhou 7,5% de desempenho sobre a versão em Python.

A configuração final roda oito cópias do modelo em uma única instância trn2.3xlarge. E é aqui que a história fica interessante para quem gosta de números batidos: três fontes diferentes deram três valores diferentes para o mesmo ganho de custo.

Nenhum dos três números vem de uma medição independente. A taxa de erro de transcrição, segundo o material, ficou dentro das referências internas da Deepgram, mas nenhum valor foi publicado.

Google mostra números de treinamento, mas não fecha a conta

Para o treinamento do modelo Qwen 3 8B com a ferramenta TorchTitan, o Google reivindica até 52% de aproveitamento da capacidade teórica de cálculo dos chips, com escala quase linear até 8.192 TPU Ironwood. A configuração completa não foi detalhada.

Um segundo caso, o cliente Zenteiq, foi de 1,47 bilhão para 5,97 bilhões de tokens gerados por dia a cada mil TPU v6e, mais de quatro vezes mais. O billet atribui esse salto ao roteamento feito por GKE Inference Gateway e llm-d, duas ferramentas de serviço de inferência, sem dizer quanto desse ganho veio do TorchTPU em si.

O que ainda trava e o que fica em aberto

O TorchNeuron está fixado na versão 2.14.0 do PyTorch, enquanto a versão 2.15 tem lançamento previsto para 28 de outubro de 2026, uma semana depois da publicação do beta. O TorchTPU não informa, nem no anúncio nem na documentação, qual versão do PyTorch exige.

No mesmo dia 20 de outubro, AMD e Intel anunciaram que pretendem reunir suas otimizações de PyTorch para processadores x86 numa única biblioteca, a oneDNN, por recomendação do x86 Ecosystem Advisory Group. A AMD vai contribuir com sua biblioteca ZenDNN e entra como membro diretor da fundação que mantém o oneDNN. O comunicado não traz versões, processadores específicos nem prazo.

O que muda para quem constrói software no Brasil

Se o seu time roda modelos PyTorch hoje em GPU Nvidia só porque não há alternativa simples, essa dependência técnica diminuiu, mas não some. Para cargas de inferência padrão, sem kernel personalizado, trocar para Trainium ou TPU pode ser, de fato, questão de mudar o dispositivo de destino e compilar de novo.

O caso Deepgram mostra o outro lado da moeda: ganho de custo real existiu, mas exigiu meses de engenharia dedicada, acesso direto às equipes da AWS e reescrita de pelo menos duas peças críticas do pipeline. Não é algo que uma equipe pequena replica num sprint. Antes de prometer economia para o cliente ou para a diretoria, vale rodar o próprio benchmark com a carga real do seu modelo, porque nenhum dos números publicados até agora veio de uma fonte independente.

Para quem decide onde hospedar a próxima carga de treinamento ou inferência, o ponto prático é este: liste as operações específicas que seu modelo usa (atenção customizada, convolução, LSTM, quantização) e confira se elas já têm suporte nativo no TorchNeuron ou no TorchTPU antes de orçar a migração. Se usam NestedTensor, formas dinâmicas ou tensores esparsos, o caminho ainda passa por reescrever kernel, não por trocar uma linha de configuração.

Fonte: ActuIA