TorchNeuron leva AWS Trainium nativo ao PyTorch

Tecnologia

Novo plugin da AWS integra Trainium ao PyTorch, ampliando suporte e facilitando treinamentos em larga escala.

Leia também: Startup de IA Tolan capta US$ 30 mi para assistente virtual

O AWS Trainium ganha, a partir de amanhã, um atalho direto para o dia a dia de pesquisadores e desenvolvedores: o TorchNeuron, plugin que integra o chip de treinamento da AWS como dispositivo nativo no PyTorch, com execução eager, suporte a torch.compile e APIs distribuídas padrão da biblioteca, dobrando a compatibilidade com modelos de texto da Hugging Face.

O que muda com o TorchNeuron no PyTorch

Até agora, quem queria treinar ou fazer inferência em AWS Trainium ou AWS Inferentia no ecossistema Hugging Face precisava passar pelo Optimum Neuron, camada extra baseada em PyTorch XLA. Isso exigia pensar como compilador, lidar com execução lazy, formas estáticas e recompilações a cada quebra de grafo, além de manter código específico para esses chips.

O TorchNeuron rompe com essa lógica. “TorchNeuron integra diretamente no PyTorch e expõe neuron como dispositivo nativo, com execução eager e APIs distribuídas nativas”, descreve o blog do PyTorch em anúncio técnico. Na prática, o Trainium passa a se comportar como mais um device do PyTorch, ao lado de CPU e GPU, usando as próprias primitivas distribuídas da biblioteca, como DTensor, FSDP2 e paralelismo de tensor.

Para começar a usar, o fluxo é o mesmo de qualquer backend suportado oficialmente: basta instalar o pacote via pip e selecionar o dispositivo neuron no código. Em muitos casos, especialmente com Transformers da Hugging Face, um simples device_map=neuron já leva o modelo para o Trainium.

Entenda o contexto

Trainium e Inferentia são chips de IA da AWS desenhados para reduzir custo e ampliar escala no treinamento e na inferência de modelos. O Optimum Neuron, usado até aqui como caminho principal no Hugging Face Hub, fazia a ponte entre essas arquiteturas e o mundo PyTorch, mas à custa de complexidade e de compatibilidade limitada com bibliotecas e modelos mais novos.

Com o TorchNeuron, o Neuron deixa de ser um ecossistema paralelo e passa a viver dentro do PyTorch “oficial”. Isso significa que scripts padrão, tutoriais e exemplos que já rodam em GPU podem, em muitos casos, ser reaproveitados quase sem mudanças em instâncias com Trainium, o que reduz a barreira de entrada principalmente para times menores, laboratórios acadêmicos e startups.

Inferentia vs Trainium: qual é o foco do TorchNeuron?

O TorchNeuron é desenhado como backend nativo para o AWS Trainium no PyTorch, com foco em treinamento e fine-tuning em larga escala, mas também abre caminho para fluxos de inferência de alta performance em modelos grandes.

Na prática, a AWS segue posicionando o Trainium como chip principal para treinamento de modelos de linguagem, visão e multimodais, enquanto o Inferentia permanece especializado em inferência em produção. O avanço de hoje está concentrado no lado do Trainium: é ele que ganha integração direta ao PyTorch, execução eager e acesso às APIs distribuídas nativas. Para quem já usa Inferentia via Optimum Neuron, a mensagem é que o uso direto das bibliotecas Hugging Face tende a se tornar o padrão à medida que o suporte nativo ao Trainium amadurecer e influenciar o restante do ecossistema Neuron.

Hugging Face muda de rota e recomenda abandonar Optimum Neuron

No anúncio conjunto, a Hugging Face indica uma mudança clara de estratégia. “Recomendamos que usuários Neuron usem as bibliotecas Hugging Face diretamente em vez de Optimum Neuron”, afirma a empresa, dando sinal verde para que Transformers, Diffusers, Accelerate, PEFT e TRL passem a ser o caminho oficial em instâncias com Trainium.

Com o TorchNeuron, o Trainium passa a ser apenas mais um device nesses pacotes: as mesmas classes e scripts usados em GPU funcionam também no chip da AWS, sem forks de código ou trainers dedicados. Isso simplifica manutenção de projetos, reduz risco de bugs específicos e facilita a colaboração em equipes que misturam ambientes on-premise, GPU em outras nuvens e hardware dedicado da AWS.

Compatibilidade de modelos praticamente dobra

Um dos efeitos mais imediatos da nova integração aparece na compatibilidade com modelos de geração de texto. Segundo a Hugging Face, “89% dos 100 modelos de geração de texto mais usados são suportados nativamente com TorchNeuron, contra 46% com Optimum Neuron v0.4.5”.

Na prática, isso significa que a maioria dos modelos de linguagem populares disponíveis hoje no Hugging Face Hub pode ser carregada no Trainium sem adaptações profundas. Do ponto de vista de empresas e laboratórios, a mudança reduz o tempo entre experimentação em GPU e migração para um ambiente de treinamento mais barato, aliviando gargalos de custo e disponibilidade de hardware.

Treinamento distribuído com FSDP e paralelismo de tensor

Outro ponto central do TorchNeuron é o suporte nativo a treinamento distribuído usando os próprios recursos do PyTorch e das bibliotecas da Hugging Face. O plugin permite combinar FSDP (Fully Sharded Data Parallel) e paralelismo de tensor em arranjos flexíveis para escalar modelos grandes em múltiplos NeuronCores.

Em um exemplo divulgado, um treinamento de modelo de linguagem roda em 8 chips Trainium2, totalizando 32 NeuronCores, com um esquema de paralelismo 4 x 8, ou seja, 32 processos. Cada camada é dividida entre 4 núcleos em grupos de paralelismo de tensor, enquanto o FSDP fragmenta parâmetros, gradientes e estados de otimizador em 8 grupos, aproveitando toda a máquina. Ao final, o script gera checkpoints padrão da biblioteca Transformers, compatíveis com qualquer hardware e prontos para serem enviados ao Hugging Face Hub.

Para o desenvolvedor, o detalhe importante é que nada no código de treino fica amarrado ao Trainium: quem já fine-tunou um modelo com a TRL usando o SFTTrainer em GPU encontra praticamente o mesmo script rodando sobre o TorchNeuron.

Diffusers, imagens e kernels customizados

O impacto não se limita ao texto. Pipelines importantes da biblioteca Diffusers, como FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 e H3, já funcionam com TorchNeuron em modos eager e compilado, incluindo suporte a paralelismo de tensor. Com isso, tarefas de geração e edição de imagens também podem aproveitar o Trainium de forma mais direta, sem camadas específicas de integração.

O ecossistema é ampliado ainda pela Neuron Kernel Interface (NKI), que permite escrever kernels customizados para operações críticas e publicá-los no Hub. A equipe Neuron promete disponibilizar cada vez mais kernels otimizados, incluindo opções que aceleram transformers de visão com ganhos de performance declarados de até 2,4 vezes no Trainium, dependendo do cenário.

Impacto para pesquisa, startups e mercado de chips de IA

A chegada do TorchNeuron em beta público abre um período de testes intensos com a comunidade, que deve apontar ajustes de desempenho, bugs e lacunas de compatibilidade. A expectativa é que, ao longo dos próximos meses, atualizações consolidem o plugin como padrão de fato para executar modelos no AWS Trainium dentro do ecossistema PyTorch.

Para pesquisadores acadêmicos, equipes de ciência de dados e startups de IA, o ganho é duplo: acesso mais simples a hardware especializado e menor esforço de engenharia para portar projetos existentes. Setores que dependem de grandes modelos de linguagem, visão computacional e aplicações multimodais — de assistentes virtuais a ferramentas de geração de conteúdo — tendem a se beneficiar diretamente.

No horizonte mais amplo, a aposta da AWS em integração nativa com PyTorch e Hugging Face reforça a disputa no mercado de chips de IA, hoje dominado por GPUs. Ao tornar o Trainium mais plug-and-play, a empresa tenta reduzir uma das principais barreiras à adoção de aceleradores alternativos: a fricção no desenvolvimento.

Fonte: matéria original