A Amazon anunciou o TorchNeuron, um back-end PyTorch nativo de código aberto para sua família Trainium de chips de IA, em um movimento que remove uma das maiores barreiras práticas à adoção de silício personalizado para cargas de trabalho de aprendizado de máquina. De acordo com uma postagem oficial no blog PyTorch publicada antes da PyTorch Conference North America em 20 de outubro, os desenvolvedores agora podem executar o código do modelo PyTorch existente no Trainium com uma alteração de linha única, em vez de reescrever seus fluxos de trabalho em torno de ferramentas proprietárias.
O anúncio atraiu atenção imediata porque visa o que os analistas há muito descrevem como a vantagem mais durável da Nvidia: não o desempenho bruto do chip, mas o ecossistema de software CUDA que liga os desenvolvedores às suas GPUs. A Reuters enquadrou o lançamento como a Amazon dizendo que seus chips de IA funcionarão com a estrutura de desenvolvimento mais popular em IA, enquanto o TechCrunch chamou o TorchNeuron de a maior chance da empresa na Nvidia até agora. Para uma visão mais ampla de como o mercado de aceleradores está mudando, nossa cobertura de notícias sobre hardware de IA contínua acompanha esses desenvolvimentos à medida que eles se desenrolam.
O que o TorchNeuron realmente muda
A principal promessa do TorchNeuron é a simplicidade. No exemplo publicado no blog PyTorch, um loop de treinamento existente é executado no Trainium depois de nada mais do que mover o modelo para o novo dispositivo com `.to('neuron')`. O código escrito em `torch.accelerator`, a API independente de dispositivo recomendada pelo PyTorch, é executado sem modificações, de acordo com a postagem. Para cargas de trabalho compiladas, os desenvolvedores definem o backend padrão uma vez com `torch.compiler.set_default_backend("neuron")` e então invocam `model.compile()` sem mais alterações.
Nos bastidores, o TorchNeuron registra o Trainium por meio do PrivateUse1, o ponto de extensão oficial do PyTorch para back-ends de hardware personalizados, o que significa que o chip se comporta como um dispositivo PyTorch padrão, em vez de um apêndice aparafusado na estrutura. A AWS descreve a integração abrangendo vários modos de execução:
- Execução rápida para pesquisa e depuração, com um "modo ansioso adaptativo" experimental que reduz a sobrecarga de inicialização
- torch.compile para desempenho de produção, capturando modelos em gráficos que o compilador Neuron otimiza, com suporte nativo para FlexAttention
- A Neuron Kernel Interface (NKI) para pesquisadores que precisam de kernels personalizados de baixo nível
O backend funciona tanto no Trainium2 quanto nos chips Trainium3 mais recentes, de acordo com o anúncio.
Um ataque direto ao software Moat da Nvidia
O domínio da Nvidia na computação de IA nunca se baseou apenas em hardware mais rápido. CUDA, a plataforma de software da empresa com duas décadas de existência, é onde a maior parte do código de pesquisa de IA é escrita, ajustada e implantada, e todos os chips desafiadores – da linha Instinct da AMD aos TPUs do Google e Trainium da Amazon – tiveram que enfrentar o custo de pedir aos desenvolvedores que o abandonassem. A cobertura do lançamento do TorchNeuron no GeekWire, Benzinga e GuruFocus enquadrou uniformemente o anúncio como um ataque àquele fosso.
A Amazon não está tentando tirar os desenvolvedores do PyTorch para conquistá-los. Em vez disso, está trazendo o Trainium para o PyTorch, para que a enorme base de usuários da estrutura possa adotar o chip sem mudar de hábitos. A postagem do blog PyTorch incluiu um endosso da Meta, que mantém a estrutura: “O suporte nativo para AWS Trainium expande essa escolha, permitindo que pesquisadores e engenheiros experimentem e iterem usando a execução ávida e APIs distribuídas que eles já conhecem”, disse Jana van Greunen, Diretora de Engenharia PyTorch da Meta, na postagem.
Esse enquadramento é importante para os negócios em nuvem da AWS. As instâncias Trainium foram comercializadas como uma alternativa de baixo custo às GPUs Nvidia para treinamento e inferência, e os chips já atraíram grandes clientes de IA. Fazer com que os chips pareçam nativos do PyTorch aborda diretamente a objeção que retardou a adoção mais ampla: o esforço de portabilidade.
Escopo técnico, de notebooks a treinamento em escala de cluster
Os recursos listados no blog PyTorch sugerem que a Amazon está almejando muito além dos experimentos amadores. TorchNeuron cobre os conjuntos de operadores vistos em modelos de texto, multimodais e de linguagem de visão, bem como geração de imagem e vídeo, modelos de espaço de estado e modelos mundiais, com precisão BF16, FP32 e FP16; O suporte TorchAO e FP8 estão em beta privado antes da disponibilidade geral na próxima versão.
O treinamento distribuído, a disciplina que transforma código de GPU único em execuções em escala de cluster, é suportado por meio de APIs PyTorch padrão – DDP, FSDP2, SimpleFSDP, paralelismo de tensor e pipeline e sharding estilo ZeRO, junto com DTensor e DeviceMesh para grupos de processos multidimensionais. Streams e eventos mapeiam a semântica do PyTorch nas filas de hardware do Trainium, dizia o post.
Para inferência, as pilhas Neuron da Amazon baseiam-se na mesma base: o vLLM Neuron atende a grandes modelos de linguagem, enquanto o vLLM Omni Neuron atende a modelos de difusão para geração de imagens e vídeos. A Neuron Kernel Interface permite que os pesquisadores escrevam operadores personalizados para execução antecipada e compilada, e a AWS é o NKI-lib de código aberto, um conjunto de kernels otimizados para Trainium que cobrem cálculos em blocos FlashAttention e MoE para LLM convencionais e cenários multimodais.
Ferramentas de desempenho também estão incluídas. O Neuron Explorer estende o padrão `torch.profiler` para rastrear a execução da estrutura ao hardware, vinculando cada evento à sua linha de origem, e pode exportar métricas no formato Parquet que as notas postais podem ser consumidas pelos agentes para análise automatizada de gargalos.
Disponibilidade e o que vem a seguir
TorchNeuron está pronto para seu lançamento de código aberto na PyTorch Conference North America em 20 de outubro, de acordo com a postagem do blog, com compatibilidade estendendo-se ao Hugging Face Transformers v5. A AWS descreve o Trainium como um portfólio de chips projetados desde o início para treinamento e inferência de IA em escala de nuvem.
O contexto estratégico é um mercado em que cada grande fornecedor de nuvem é simultaneamente o maior comprador de aceleradores Nvidia e um potencial concorrente. Microsoft, Google e Meta desenvolvem silício personalizado, e empresas chinesas, incluindo a Huawei, estão construindo suas próprias pilhas de software para escapar da gravidade do CUDA. Ao reunir os desenvolvedores dentro do PyTorch em vez de pedir-lhes que migrassem, a Amazon escolheu o caminho menos perturbador para esse concurso – e o mais provável, se a experiência do desenvolvedor se mantiver, para transferir cargas de trabalho de treinamento reais do hardware da Nvidia.
Se o TorchNeuron realmente destruirá a vantagem do CUDA dependerá do desempenho em escala e da longa cauda dos modelos do mundo real, questões que os materiais de lançamento não respondem. Mas a barreira que impedia a maioria dos desenvolvedores do PyTorch de tentar o Trainium – uma reescrita – agora foi reduzida a uma única linha de código.
---
Fique à frente da IAReceba as últimas notícias, análises e avanços sobre IA — tudo em um só lugar.
Leia mais notícias sobre IA →
