Cómo averiguar por qué tu modelo de IA va lento: guía para hacer perfiles en AWS Trainium
Amazon y Hugging Face publican un tutorial paso a paso para analizar el rendimiento de modelos de inteligencia artificial y detectar dónde se pierde tiempo al ejecutarlos en chips especializados.
Amazon Web Services y Hugging Face publicaron un tutorial que enseña a hacer perfiles de rendimiento en AWS Trainium (un chip para entrenar y ejecutar modelos de IA) usando TorchNeuron (el backend que adapta PyTorch a estos chips). La guía muestra cómo medir dónde se pierde tiempo al ejecutar modelos y cómo usar Neuron Explorer para visualizar cuellos de botella en el hardware.
Muchos desarrolladores de machine learning tienen intuición sobre cómo hacer que un modelo funcione rápido, pero pocos saben analizar por qué un modelo va lento. Amazon Web Services y Hugging Face han publicado la primera parte de una guía técnica que enseña a hacer perfiles de rendimiento (profiling, la técnica que mide dónde se gasta el tiempo de computación) en AWS Trainium, un chip diseñado específicamente para entrenar y ejecutar modelos de IA. El tutorial usa TorchNeuron, el backend de código abierto que permite ejecutar PyTorch (el framework de programación de redes neuronales más popular) en estos chips sin apenas cambiar el código.
La guía arranca desde cero y no asume conocimientos previos de optimización. Explica cómo configurar torch.profiler (la herramienta integrada en PyTorch para medir el rendimiento) con el backend Neuron, cómo leer la tabla key_averages() que muestra qué operaciones consumen más tiempo en el chip, y cómo seguir una operación completa desde que el framework la despacha hasta que se ejecuta en el hardware. Para ilustrar el proceso, el tutorial perfila tres implementaciones distintas de RMSNorm, una operación matemática habitual en modelos de lenguaje: una versión ingenua con operaciones básicas de PyTorch, la función integrada torch.rms_norm(), y la versión ingenua compilada con torch.compile.
El texto introduce Neuron Explorer, una herramienta que se abre en el navegador y permite ver la traza del sistema, la línea de tiempo del chip y el desglose por motor de ejecución. Tras ejecutar el perfilado, la herramienta genera archivos NEFF (Neuron Executable File Format, el formato del ejecutable compilado para el chip) y NTFF (que contiene la traza de instrucciones del hardware), emparejados con un código numérico que los identifica. Explorer ayuda a confirmar visualmente lo que la tabla key_averages() reporta en números y a inspeccionar qué pasa entre operaciones.
El tutorial muestra que en la versión ingenua de RMSNorm, tres llamadas a aten::mul (la operación de multiplicación a nivel interno de PyTorch) acumulan el 49,5% del tiempo total en el chip Neuron. La tabla key_averages() enumera cada operación ATen que se ejecutó, cuántas veces corrió, cuánto tiempo pasó en la CPU y cuánto en el chip. En Neuron Explorer, cada operación sigue el mismo patrón visual: creación en el carril del framework, planificación en el carril del runtime como nrta_execute_schedule, y ejecución en el carril del hardware como nc_exec_running. Este flujo de arriba abajo y de izquierda a derecha permite identificar cuellos de botella con más facilidad que mirando solo números.
Según la publicación en el blog de Hugging Face, el objetivo de la serie es que cualquier practicante de machine learning pueda analizar de forma sistemática el rendimiento de inferencia (la fase en que un modelo ya entrenado genera respuestas) e identificar dónde se gasta el tiempo. Todos los scripts para reproducir los ejemplos están disponibles en el directorio de ejemplos de torch-neuronx en GitHub. La segunda parte de la serie promete profundizar en técnicas más avanzadas de optimización.
