Num artigo publicado na plataforma Hugging Face, a Amazon inicia uma série dedicada à análise de desempenho de inferência, ou seja, da execução de um modelo já treinado. Esta primeira parte explica como identificar onde o tempo é gasto, utilizando ferramentas de análise de desempenho aplicadas ao chip Trainium da AWS, concebido para o treino e a inferência de modelos de deep learning.
O guia dirige-se a leitores sem experiência prévia em análise de desempenho. Baseia-se no TorchNeuron, o backend open source nativo do PyTorch para o Trainium, que suporta fluxos de trabalho padrão do PyTorch, incluindo a execução em modo eager e torch.compile, com poucas modificações de código. O TorchNeuron também se integra na Neuron Kernel Interface (NKI), que permite criar operadores personalizados.
Configurar o analisador de desempenho
A análise de desempenho baseia-se em dois objetos. NeuronConfig define os fluxos de dados a capturar, os NeuronCores a rastrear e a pasta de saída, sendo depois transmitido a torch.profiler.profile. NeuronProfiler fornece uma função de exportação que escreve o rastreio do framework, no formato trace.json, juntamente com os rastreios de sistema da Neuron. O parâmetro correlate_device_time=True permite mostrar o tempo gasto no dispositivo na tabela key_averages(), onde pode ser comparado com o tempo de CPU.
Os rastreios podem depois ser abertos no Neuron Explorer, uma ferramenta de visualização que funciona no navegador. Permite seguir uma operação desde a distribuição no framework até à sua execução no hardware, passando pelo agendamento no runtime. Para um NeuronCore, ou seja, um núcleo de computação do chip, o ficheiro NEFF contém o executável compilado, enquanto o ficheiro NTFF correspondente contém o rastreio das instruções de hardware.
Três versões de uma mesma operação
O artigo estuda a função RMSNorm, uma operação simples composta por uma sequência de cálculos aritméticos e reduções. São comparadas três implementações: uma versão escrita operação a operação em PyTorch, a chamada à função integrada torch.rms_norm(), e a versão escrita operação a operação mas compilada com torch.compile.
Na versão ingénua, a tabela mostra que três chamadas a aten::mul representam a maior parte do tempo gasto no dispositivo. O tempo total da Neuron atinge 2 355 µs para seis chamadas de operadores. Segundo o autor, a função integrada torch.rms_norm(), que é transformada num único núcleo de computação fundido, ocupa apenas uma operação, aten::_fused_rms_norm, com um tempo Neuron de 403,3 µs. A versão compilada com torch.compile aparece como uma única região compilada, com uma duração de 375,5 µs.
O artigo esclarece que, nestas experiências, as duas abordagens demoram cerca de um sexto do tempo de computação da versão ingénua. A fusão elimina as interrupções entre as execuções de hardware e evita materializar valores intermédios, como x_sq e mean_sq, entre operações enviadas separadamente.
O autor indica que o método pode ser aplicado a modelos maiores: identificar um ponto crítico na tabela, verificá-lo no rastreio e depois determinar se uma operação fundida ou o torch.compile pode reduzir o seu custo. A segunda parte da série abordará o perfil ao nível do dispositivo, com a análise do ficheiro NEFF.
