TorchNeuron: de 6 operaciones a 1 kernel en Trainium

Qué enseña el primer capítulo de la serie de TorchNeuron

TorchNeuron es el backend nativo de PyTorch para los chips Trainium de AWS: el mismo código de PyTorch corre en el chip cambiando una línea de dispositivo, sin reescribir el modelo. Sobre esa base, el equipo de Amazon publicó en Hugging Face la Parte 1 de una serie sobre perfilado, y abre con una premisa incómoda: muchos practicantes de machine learning tienen intuición sobre cómo hacer que un modelo corra rápido, pero pocos saben analizar por qué corre lento.

El contexto importa. TorchNeuron fue presentado el 19 de octubre de 2026, según CryptoBriefing, y su beta pública abrió el 20 de octubre en la PyTorch Conference de San José bajo licencia Apache 2.0. KuCoin News reportó que alrededor del 85% de las aplicaciones de IA están escritas en PyTorch, lo que explica por qué Amazon apunta a ese framework y no a uno propio.

Bajar la fricción de migración es solo la mitad del problema. La otra mitad es que, una vez que el código corre en Trainium, el equipo sepa por qué corre a la velocidad que corre. De eso trata el post.

¿Cómo se perfila un workload en Trainium paso a paso?

El artículo explica que el backend de Neuron se apoya en el punto de extensión PrivateUse1 de PyTorch, el mecanismo que permite integrar un acelerador sin agregar una clave de dispatch dedicada como cuda. Sobre esa base, el flujo de perfilado usa dos objetos:

• NeuronConfig: controla qué flujos de datos y qué NeuronCores se capturan, y dónde se escribe la salida. Se pasa a torch.profiler.profile como experimental_config.
• NeuronProfiler: recibe un NeuronConfig y expone un callback export_trace que escribe el trace del framework (trace.json) junto a los traces del sistema Neuron.

Dos detalles del código son los que suelen arruinar un perfil la primera vez. Primero, hay que incluir ProfilerActivity.PrivateUse1 en las actividades del profiler; sin eso no hay actividad de dispositivo que capturar. Segundo, hay que llamar a torch.neuron.synchronize() dentro del loop del profiler: si no se sincroniza, el profiler puede exportar antes de que el dispositivo termine y la columna de tiempo de Neuron queda vacía.

El schedule del ejemplo usa wait=2, warmup=1 y active=1 a lo largo de cuatro llamadas a prof.step(), de modo que se registra una iteración ya caliente sin arrastrar el overhead de arranque.

El caso RMSNorm: 2.355 µs contra 376 µs

Para que el ejercicio sea comparable, el post perfila tres implementaciones de la misma operación, RMSNorm, con los mismos inputs, la misma configuración de profiler y el mismo schedule:

• RMSNorm naive con operaciones PyTorch sueltas: multiplicación, media, suma, raíz cuadrada recíproca y multiplicación.
• torch.rms_norm(), el built-in de PyTorch que baja a un único kernel fusionado de Neuron.
• RMSNorm naive con torch.compile, usando las mismas operaciones que la primera versión.

Los resultados, medidos en tiempo de dispositivo:

Variante	Tiempo NEURON	Self CPU	Llamadas a operadores
Ops naive	2.355 µs	~125 µs	6
torch.rms_norm	403 µs	~48 µs	1
Naive + torch.compile	376 µs	~98 µs	1

La lectura es contundente: fusionar lleva el costo de dispositivo a alrededor de un sexto del naive. Y no es solo el número total. En la tabla de key_averages() del naive, aten::mul aparece con 3 llamadas y 1.164,5 µs de tiempo de dispositivo (49,5% del total), porque la función multiplica tres veces: x * x, x * rms y el escalado final por weight. Le siguen aten::add con 434,6 µs (18,5%), aten::rsqrt con 415,0 µs (17,6%) y aten::mean con 340,8 µs (14,5%), para un total de 2,355 ms.

El detalle que más enseña es el de aten::add: suma el epsilon y, como el escalar arranca en la CPU, dispara una copia host-to-device. Es exactamente el tipo de costo invisible que un perfil de CPU nunca te muestra.

¿Cómo leer un trace en Neuron Explorer?

El post usa Neuron Explorer, un visor en el navegador —accesible por túnel SSH o desde la extensión de VS Code— para confirmar visualmente lo que la tabla reporta en números. En el trace, el tiempo corre de izquierda a derecha y el ciclo de vida del evento va de arriba hacia abajo. Cada operación ATen sigue el mismo patrón: creación en el lane del framework, scheduling en el lane de runtime como nrta_execute_schedule, y ejecución en el lane de hardware como nc_exec_running.

Los lanes se organizan por subsistema: dispatch del framework, dispositivo del framework, runtime de Neuron y hardware de Neuron. En modo eager, cada operación recorre ese camino por separado; con el kernel fusionado hay un solo dispatch y un solo evento de hardware, sin huecos entre operaciones.

Un detalle práctico que suele frenar a quien perfila por primera vez: la carpeta de salida contiene el trace del sistema y, cuando se capturan datos de dispositivo, archivos .neff y .ntff emparejados por un hash numérico. El NEFF es el ejecutable compilado; el NTFF es su trace de instrucciones de hardware. Para abrir un perfil de sistema se sube la carpeta completa; para uno solo de dispositivo, se suben el .neff y el .ntff correspondientes.

¿Qué significa esto para tu startup?

El post es técnico, pero la decisión de negocio detrás es muy concreta. AWS sostiene que las instancias Trn2 basadas en Trainium2 ofrecen entre 30% y 40% mejor precio-rendimiento en entrenamiento frente a instancias EC2 con GPU, y que las Inf2 recortan los costos de inferencia hasta 50% en LLMs de alto throughput, según AndroGuider. KuCoin News recogió además el caso de Deepgram, que migró a Trainium con soporte nativo de PyTorch y logró 70% más de audio procesado por dólar.

Ahí está el punto: el ahorro no se cobra solo por mover el workload, se cobra por entender dónde se va el tiempo una vez que está ahí. Un equipo que migra a Trainium y no perfila puede quedarse con un kernel naive pagando varias veces el costo de dispositivo que necesita.

Dos acciones concretas para esta semana:

• Perfila antes de migrar, no después. Toma tu operación más caliente y corre las tres variantes del post —naive, built-in fusionado y torch.compile(backend="neuron")— con el mismo schedule. Si tu cuello de botella es una secuencia de ops sueltas, el built-in o el compilador probablemente ya te resuelven el problema sin escribir un kernel a mano.
• Ordena por self_privateuse1_time_total, no por CPU. El propio post lo indica: ordenar por tiempo de dispositivo pone arriba lo que realmente cuesta. Y revisa los aten::add y las conversiones de escalares: cada copia host-to-device es latencia que no aparece en un perfil de CPU.

Conclusión

La Parte 1 de la serie deja una metodología más que un truco: encontrar el hotspot en la tabla, verificarlo en el trace y decidir si una operación fusionada o torch.compile puede bajarlo. En el ejemplo de RMSNorm, esa disciplina convierte 2.355 µs en 376 µs. La Parte 2 promete el nivel de dispositivo: abrir el NEFF, mirar el desglose de TensorEngine, VectorEngine y DMA, y ubicar el kernel en la roofline para saber si está limitado por cómputo o por memoria.

Para un founder con presupuesto de cómputo ajustado, esa es la diferencia entre "Trainium es más barato" como folleto y "Trainium es más barato para mi modelo" como número propio.

Fuentes

• Profiling TorchNeuron (Part 1): System-level profiling on AWS Trainium (fuente original)
• Amazon unveils TorchNeuron to make its Trainium AI chips easier to use - CryptoBriefing
• Amazon Takes Aim at Nvidia with TorchNeuron - AndroGuider
• AWS Launches TorchNeuron to Challenge NVIDIA in the AI Ecosystem - KuCoin News