O TorchNeuron da AWS traz um backend PyTorch nativo para o Trainium

As notas de versão do Neuron 2.33 da AWS datam a beta pública do TorchNeuron de 19 de outubro. O anúncio do PyTorch, também datado de 19 de outubro, diz que a beta pública e o lançamento em código aberto ficam disponíveis a 20 de outubro. O pacote torch-neuronx regista no PyTorch um dispositivo neuron para Trainium2 e Trainium3, permitindo que as operações de modelos e tensores suportadas usem as API habituais do PyTorch. Para os engenheiros que ponderam estes chips, a mudança útil é um caminho mais curto entre um programa PyTorch existente e uma experiência no dispositivo. Se esse programa funciona correta e eficientemente continua a depender das suas operações, formas e ambiente de software.

O percurso anterior do PyTorch para Trainium usava uma biblioteca torch-neuronx baseada em XLA e código de dispositivo torch_xla. A AWS diz que o backend nativo usa a extensão PrivateUse1 do PyTorch e substitui xm.xla_device() pela colocação no dispositivo neuron. É uma alteração significativa da interface para projetos que já usam Neuron, embora o nome do pacote seja familiar. A versão aplica-se ao treino e à inferência em instâncias Trn2 e Trn3; a matriz de suporte não inclui Trn1.

O que muda num programa PyTorch

As notas de versão da AWS mostram um modelo movido com .to("neuron") e um tensor de entrada criado em device="neuron". As operações suportadas de torch.nn, autograd e otimizadores passam então pelo backend. O código escrito para a API do PyTorch independente do dispositivo torch.accelerator pode detetar Neuron sem uma reescrita específica para esse dispositivo, segundo a visão geral do TorchNeuron.

Há dois percursos de execução. O modo imediato executa as operações PyTorch quando o Python as invoca, preservando o fluxo de controlo, as saídas de impressão e os pontos de interrupção habituais. Cada operação é compilada na primeira utilização e as chamadas seguintes que correspondam podem usar uma cache. Os engenheiros também podem compilar um modelo ou uma região selecionada com torch.compile(model, backend="neuron", dynamic=False). O backend nativo compila os grafos suportados de propagação direta e inversa; o código fora da região capturada continua em modo imediato. O anúncio do PyTorch também descreve a definição de torch.compiler.set_default_backend("neuron") uma vez antes de compilar, para que as chamadas de compilação posteriores possam omitir o argumento do backend.

A versão inclui um backend neuron para torch.distributed, com operações coletivas e ponto a ponto documentadas. A AWS enumera DistributedDataParallel, FSDP2, DTensor e DeviceMesh, bem como o paralelismo em pipeline, entre as API suportadas. São percursos úteis a analisar num projeto de treino distribuído, mas as notas de versão não demonstram que uma tarefa existente com vários nós funcione sem alterações ou atinja uma meta específica de débito.

A avaliação começa com um ambiente compatível

O guia de instalação rápida da AWS exige uma instância Trn2 ou Trn3 e uma Neuron Base DLAMI ou uma Ubuntu AMI suportada. A matriz da versão especifica Linux x86_64 em Ubuntu 24.04 ou 26.04. O suporte para Ubuntu 22.04 termina nesta versão. São indicados Python 3.12, 3.13 e 3.14, com torch-neuronx 2.13.3.0 para PyTorch 2.13.0 ou torch-neuronx 2.14.3.0 para PyTorch 2.14.0. A família do pacote tem de corresponder à versão do PyTorch.

O guia rápido instala torch-neuronx==2.14.* juntamente com neuronx-cc==2.* e nki a partir dos índices de pacotes Neuron e PyTorch. Depois verifica a disponibilidade do dispositivo e compara uma pequena operação matricial em Neuron com um resultado da CPU. Isto verifica o percurso do dispositivo instalado no exemplo da AWS. Não valida o modelo do leitor. Uma avaliação seguinte sensata consiste em executar os percursos direto e inverso do modelo real, inspecionar as operações de recurso à CPU registadas com torch_neuronx.get_fallback_ops(), comparar os resultados com o ambiente atual e medir separadamente a compilação e as execuções em regime estável. Estes últimos passos são questões de avaliação propostas pela BIG CHANGE, não testes que tenhamos realizado.

Onde a beta ainda pode alterar o trabalho

No modo imediato, uma operação não suportada pode recorrer à CPU e emitir um aviso. Com torch.compile, uma operação capturada num grafo que o backend Neuron não consiga converter para uma representação de nível inferior faz a compilação falhar por predefinição. A AWS remete os programadores para a sua tabela de compatibilidade de operadores para conhecerem a cobertura por operação. Assim, uma execução imediata concluída não prova, por si só, a execução integral no dispositivo nem a compatibilidade com a compilação.

A versão também indica que torch.compile exige dynamic=False; uma nova forma de entrada pode exigir outra compilação. A AWS sugere padding ou agrupamento em buckets para reduzir a variação das formas. A primeira execução inclui a compilação se não houver um artefacto correspondente em cache. Estes pormenores são importantes ao medir o tempo de uma carga de trabalho: a latência de compilação e o padrão das formas de entrada podem alterar o resultado de uma avaliação.

Outros limites explícitos restringem a alegação de compatibilidade. A beta não suporta torch.sparse, torch.func, torch.nested ou torch.distributions. Apenas algumas transformadas de torch.fft são executadas em Neuron, e float64 é executado como float32 com um aviso. A integração TorchAO e o suporte para o tipo de dados composto MXFP8 estão ausentes desta versão. Uma equipa que use estas interfaces precisa de verificar o seu próprio percurso de código antes de considerar a colocação no dispositivo como um plano de migração.

A grande mudança

O TorchNeuron coloca o Trainium por trás das interfaces nativas de dispositivos e compilação do PyTorch numa beta pública. Isto torna a primeira experiência mais direta para programas PyTorch suportados e dá aos utilizadores atuais do Neuron um percurso diferente do seu código baseado em XLA. A prova decisiva para a adoção virá de cada carga de trabalho: se os operadores permanecem no dispositivo pretendido, se os resultados numéricos continuam a ser aceitáveis e se o desempenho e o custo medidos justificam a instância e o trabalho de migração. A BIG CHANGE analisou a documentação e os anúncios para esta análise; não executámos TorchNeuron nem verificámos de forma independente as alegações de desempenho da AWS.

Fontes e leituras adicionais

• Notas de versão do AWS Neuron 2.33 TorchNeuron: data de lançamento, estado da beta, matriz de versões e sistemas operativos, API e limitações conhecidas.
• Visão geral do TorchNeuron da AWS: integração de dispositivos, execução imediata e compilada, comportamento de recurso à CPU e percurso do software através do Neuron.
• Guia rápido de instalação e verificação da AWS: requisitos da instância e da imagem, instalação de pacotes e verificação documentada do dispositivo.
• Blogue PyTorch: apresentação do TorchNeuron: explicação, do lado do PyTorch, sobre a mudança de dispositivo, compilação e API distribuídas. O artigo é um anúncio de colaboradores da AWS no blogue PyTorch, não um teste de desempenho independente.
• Anúncio da Amazon: descrição da AWS dos seus objetivos e primeiros utilizadores. As alegações de migração, desempenho e economia exigem verificação específica para cada carga de trabalho.