As notas de versão do Neuron 2.33 da AWS datam o beta público do TorchNeuron em 19 de outubro. O anúncio do PyTorch, também datado de 19 de outubro, informa que o lançamento do beta público e do código aberto ocorre em 20 de outubro. O pacote torch-neuronx registra um dispositivo neuron no PyTorch para Trainium2 e Trainium3, permitindo que operações de modelos e tensores compatíveis usem as APIs conhecidas do PyTorch. Para os engenheiros que consideram esses chips, a mudança útil é um caminho mais curto entre um programa PyTorch existente e um teste no dispositivo. Se esse programa funciona corretamente e com eficiência ainda depende de suas operações, dos formatos dos tensores e do ambiente de software.

O caminho anterior do PyTorch para o Trainium usava uma biblioteca torch-neuronx baseada em XLA e código de dispositivo torch_xla. A AWS afirma que o backend nativo usa a extensão PrivateUse1 do PyTorch e substitui xm.xla_device() pelo posicionamento no dispositivo neuron. É uma mudança significativa de interface para projetos que já usam o Neuron, embora o nome do pacote seja conhecido. A versão se aplica a treinamento e inferência em instâncias Trn2 e Trn3; a matriz de suporte não lista o Trn1.

O que muda em um 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 compatíveis de torch.nn, autograd e otimizadores passam então pelo backend. O código escrito para a API independente de dispositivo torch.accelerator do PyTorch pode descobrir o Neuron sem uma reescrita específica para o dispositivo, segundo a visão geral do TorchNeuron.

Há dois caminhos de execução. O modo eager executa as operações do PyTorch conforme o Python as chama, preservando o fluxo de controle normal, as impressões e os pontos de interrupção. Cada operação é compilada no primeiro uso, e chamadas posteriores correspondentes podem usar um 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 compatíveis de propagação direta e retropropagação; o código fora da região capturada continua em modo eager. O anúncio do PyTorch também descreve como configurar torch.compiler.set_default_backend("neuron") uma vez antes de compilar, para que 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 lista DistributedDataParallel, FSDP2, DTensor, DeviceMesh e paralelismo de pipeline entre as APIs compatíveis. São caminhos úteis a examinar em um projeto de treinamento distribuído, mas as notas de versão não demonstram que uma tarefa existente em vários nós funcionará sem alterações nem que atingirá uma meta específica de vazão.

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 AMI Ubuntu compatível. A matriz da versão especifica Linux x86_64 no Ubuntu 24.04 ou 26.04. O suporte ao Ubuntu 22.04 terminou para essa versão. Ela lista 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 precisa corresponder à versão do PyTorch.

O guia rápido instala torch-neuronx==2.14.* junto com neuronx-cc==2.* e nki a partir dos índices de pacotes do Neuron e do PyTorch. Em seguida, verifica a disponibilidade do dispositivo e compara uma pequena operação matricial no Neuron com um resultado na CPU. Isso verifica o caminho do dispositivo instalado no exemplo da AWS. Não valida o modelo do leitor. Uma próxima avaliação sensata é executar os caminhos de propagação direta e retropropagação do modelo real, examinar os fallbacks para CPU registrados com torch_neuronx.get_fallback_ops(), comparar as saídas com as do ambiente atual e medir separadamente a compilação e as execuções em regime estável. Essas últimas etapas são questões de avaliação propostas pelo BIG CHANGE, e não testes que realizamos.

Onde o beta ainda pode mudar o trabalho

No modo eager, uma operação não compatível pode recorrer à CPU e emitir um aviso. Com torch.compile, uma operação capturada em um grafo que o backend Neuron não consegue converter faz a compilação falhar por padrão. A AWS direciona os desenvolvedores à sua tabela de compatibilidade de operadores para consultar a cobertura de cada operação. Portanto, uma execução eager concluída não comprova, por si só, que toda a execução ocorre no dispositivo nem que há compatibilidade com a compilação.

A versão também informa que torch.compile exige dynamic=False; um novo formato de entrada pode exigir outra compilação. A AWS sugere preenchimento ou agrupamento por formatos, conhecido como bucketing, para reduzir essa variação. A primeira execução inclui a compilação quando não há um artefato correspondente no cache. Esses detalhes importam ao cronometrar uma carga de trabalho: a latência de compilação e o padrão dos formatos de entrada podem alterar o resultado de uma avaliação.

Outros limites explícitos restringem a compatibilidade anunciada. O beta não oferece suporte a torch.sparse, torch.func, torch.nested ou torch.distributions. Apenas algumas transformações de torch.fft são executadas no Neuron, e float64 é executado como float32 com um aviso. A integração com TorchAO e o suporte ao tipo de dados composto MXFP8 estão ausentes nessa versão. Uma equipe que usa essas funções precisa verificar seu próprio caminho de código antes de tratar o posicionamento do dispositivo como um plano de migração.

A grande mudança

O TorchNeuron coloca o Trainium por trás das interfaces nativas de dispositivo e compilação do PyTorch em um beta público. Isso torna um primeiro teste mais direto para programas PyTorch compatíveis e oferece aos usuários atuais do Neuron um caminho diferente do código baseado em XLA. A evidência decisiva para a adoção virá de cada carga de trabalho: se seus operadores permanecem no dispositivo pretendido, se seus resultados numéricos continuam aceitáveis e se o desempenho e o custo medidos justificam a instância e o trabalho de migração. O BIG CHANGE analisou a documentação e os anúncios para esta análise; não executamos o TorchNeuron nem verificamos de forma independente as alegações de desempenho da AWS.

Fontes e leituras complementares