Perfilado de TorchNeuron (Parte 1): Perfilado a nivel del sistema en AWS Trainium
huggingface ·
Este artículo introduce el perfilado a nivel del sistema para TorchNeuron en AWS Trainium, con el objetivo de ayudar a los profesionales de aprendizaje automático a analizar cuellos de botella en el rendimiento de inferencia. Detalla la configuración de NeuronConfig y NeuronProfiler para capturar trazas del dispositivo y del tiempo de ejecución mediante la extensión PrivateUse1 de PyTorch. La guía demuestra cómo utilizar Neuron Explorer para visualizar el ciclo de vida de las operaciones desde el despacho del marco de trabajo hasta la ejecución en hardware. Se comparan tres implementaciones de RMSNorm: operaciones ingenuas de PyTorch, torch.rms_norm integrado y una función ingenua compilada con torch.compile. Los resultados muestran que tanto la fusión mediante torch.rms_norm como la compilación reducen significativamente el tiempo del dispositivo Neuron al eliminar la materialización de tensores intermedios y los huecos entre operaciones. Los enfoques fusionados logran aproximadamente una sexta parte del tiempo de dispositivo requerido por la implementación ingenua. El autor concluye que identificar puntos calientes en las tablas key_averages y verificarlos en las trazas constituye un flujo de trabajo eficaz para la optimización.
huggingface