Perfilamento do TorchNeuron (Parte 1): Perfilamento em nível de sistema no AWS Trainium
huggingface ·
Este artigo apresenta o perfilamento em nível de sistema para o TorchNeuron na AWS Trainium, com o objetivo de ajudar profissionais de machine learning a analisar gargalos de desempenho na inferência. Ele detalha a configuração do NeuronConfig e do NeuronProfiler para capturar rastros do dispositivo e do tempo de execução usando a extensão PrivateUse1 do PyTorch. O guia demonstra como utilizar o Neuron Explorer para visualizar o ciclo de vida das operações, desde o despacho do framework até a execução no hardware. Três implementações da RMSNorm são comparadas: operações ingênuas do PyTorch, a função integrada torch.rms_norm e uma função ingênua compilada com torch.compile. Os resultados mostram que tanto a fusão por meio do torch.rms_norm quanto a compilação reduzem significativamente o tempo do dispositivo Neuron ao eliminar a materialização de tensores intermediários e lacunas entre as operações. As abordagens fundidas consomem aproximadamente um sexto do tempo de dispositivo exigido pela implementação ingênua. O autor conclui que identificar pontos críticos nas tabelas key_averages e verificá-los nos rastros é um fluxo de trabalho eficaz para otimização.
huggingface