En una publicación en la plataforma Hugging Face, Amazon inicia una serie dedicada al análisis del rendimiento de la inferencia, es decir, de la ejecución de un modelo ya entrenado. Esta primera parte explica cómo identificar dónde se gasta el tiempo, utilizando herramientas de perfilado aplicadas al chip Trainium de AWS, diseñado para el entrenamiento y la inferencia de modelos de aprendizaje profundo.
- Amazon inicia una serie sobre el perfilado del rendimiento de inferencia en AWS Trainium.
- TorchNeuron y el perfilador localizan dónde se gasta el tiempo de ejecución.
- La versión fusionada y la compilada reducen el tiempo de cálculo de RMSNorm a una sexta parte.
- La fusión elimina interrupciones entre ejecuciones y valores intermedios; la próxima parte analizará el archivo NEFF.
La guía está dirigida a lectores sin experiencia previa en perfilado. Se basa en TorchNeuron, el backend de código abierto nativo de PyTorch para Trainium, que admite flujos de trabajo estándar de PyTorch, incluida la ejecución en modo eager y torch.compile, con pocas modificaciones del código. TorchNeuron también se integra con la Neuron Kernel Interface (NKI), que permite crear operadores personalizados.
Configurar el perfilador
El perfilado se basa en dos objetos. NeuronConfig define los flujos de datos a capturar, los NeuronCores a rastrear y la carpeta de salida, y luego se pasa a torch.profiler.profile. NeuronProfiler proporciona una función de exportación que escribe el rastro del framework, en formato trace.json, junto a los rastros del sistema de Neuron. El parámetro correlate_device_time=True permite mostrar el tiempo pasado en el dispositivo en la tabla key_averages(), donde puede compararse con el tiempo de CPU.
Los rastros pueden luego abrirse en Neuron Explorer, una herramienta de visualización que funciona en el navegador. Permite seguir una operación desde la distribución en el framework hasta su ejecución en el hardware, pasando por la planificación en el runtime. Para un NeuronCore, es decir, un núcleo de cálculo del chip, el archivo NEFF contiene el ejecutable compilado, mientras que el archivo NTFF correspondiente contiene el rastro de las instrucciones de hardware.
Tres versiones de una misma operación
La publicación estudia la función RMSNorm, una operación simple compuesta por una serie de cálculos aritméticos y reducciones. Se comparan tres implementaciones: una versión escrita operación por operación en PyTorch, la llamada a la función integrada torch.rms_norm(), y la versión escrita operación por operación pero compilada con torch.compile.
En la versión ingenua, la tabla muestra que tres llamadas a aten::mul representan la mayor parte del tiempo pasado en el dispositivo. El tiempo total de Neuron alcanza 2 355 µs para seis llamadas de operadores. Según el autor, la función integrada torch.rms_norm(), que se transforma en un solo núcleo de cálculo fusionado, ocupa solo una operación, aten::_fused_rms_norm, con un tiempo de Neuron de 403,3 µs. La versión compilada con torch.compile aparece como una sola región compilada, de una duración de 375,5 µs.
La publicación precisa que, en estas pruebas, ambos enfoques toman aproximadamente una sexta parte del tiempo de cálculo de la versión ingenua. La fusión elimina las interrupciones entre las ejecuciones de hardware y evita materializar valores intermedios, como x_sq y mean_sq, entre operaciones enviadas por separado.
El autor indica que el método puede aplicarse a modelos más grandes: identificar un punto crítico en la tabla, verificarlo en el rastro y luego determinar si una operación fusionada o torch.compile puede reducir su costo. La segunda parte de la serie tratará sobre el perfil a nivel del dispositivo, con el análisis del archivo NEFF.
