TorchNeuron: siga a operação até o hardware
Um exemplo de RMSNorm conecta a tabela do profiler à execução no Trainium e mostra onde investigar antes de mudar código.
torch.neuron.synchronize() tem uma função concreta no exemplo de profiling da Amazon: esperar o dispositivo terminar antes de avançar a coleta. Sem essa sincronização, explica Josh Longenecker, a exportação pode ocorrer cedo demais e deixar vazia a coluna de tempo do Neuron.
Publicado em 21 de outubro de 2026, o tutorial acompanha RMSNorm no Trainium por caminhos distintos: operações PyTorch separadas, torch.rms_norm e a função original compilada. No exemplo, tanto a operação pronta quanto a compilação reúnem o cálculo em um kernel.
A tabela key_averages() aponta onde o dispositivo gastou tempo. O Neuron Explorer permite seguir o mesmo trabalho pelo framework, agendamento e hardware. É essa conexão que torna a comparação útil: o rastro mostra os intervalos entre execuções que a fusão remove.
Eu começaria por esse percurso antes de escrever um kernel próprio. Escolha uma operação cara na tabela e confirme seu comportamento no rastro. O limite está no recorte: melhorar RMSNorm nesse exemplo não demonstra aceleração equivalente do modelo inteiro.
Fonte: Josh Longenecker, da Amazon, em Profiling TorchNeuron (Part 1): System-level profiling on AWS Trainium..