Las notas de la versión de Neuron 2.33 de AWS sitúan la beta pública de TorchNeuron el 19 de octubre. El anuncio de PyTorch, también fechado el 19 de octubre, indica que el lanzamiento de la beta pública y del código abierto se hará efectivo el 20 de octubre. El paquete torch-neuronx registra un dispositivo neuron en PyTorch para Trainium2 y Trainium3, lo que permite usar las API habituales de PyTorch para las operaciones de modelos y tensores compatibles. Para los ingenieros que consideran estos chips, el cambio útil es un camino más corto desde un programa existente de PyTorch hasta una prueba en el dispositivo. Que ese programa funcione correcta y eficientemente sigue dependiendo de sus operaciones, formas y entorno de software.

La ruta anterior de PyTorch para Trainium utilizaba una biblioteca torch-neuronx basada en XLA y código de dispositivo torch_xla. AWS afirma que el backend nativo utiliza la extensión PrivateUse1 de PyTorch y sustituye xm.xla_device() por la asignación al dispositivo neuron. Eso supone un cambio sustancial de interfaz para los proyectos que ya usan Neuron, aunque el nombre del paquete resulte familiar. La versión se aplica al entrenamiento y la inferencia en instancias Trn2 y Trn3; la matriz de compatibilidad no incluye Trn1.

Qué cambia en un programa de PyTorch

Las notas de la versión de AWS muestran un modelo trasladado mediante .to("neuron") y un tensor de entrada creado con device="neuron". Las operaciones compatibles de torch.nn, autograd y optimizadores se ejecutan entonces a través del backend. El código escrito para la API de PyTorch independiente del dispositivo torch.accelerator puede detectar Neuron sin una reescritura específica para ese dispositivo, según la descripción general de TorchNeuron.

Hay dos rutas de ejecución. El modo inmediato ejecuta las operaciones de PyTorch cuando Python las invoca, conservando el flujo de control, las impresiones y los puntos de interrupción habituales. Cada operación se compila la primera vez que se usa y las siguientes llamadas coincidentes pueden utilizar una caché. Los ingenieros también pueden compilar un modelo o una región seleccionada mediante torch.compile(model, backend="neuron", dynamic=False). El backend nativo compila los grafos compatibles de propagación hacia delante y hacia atrás; el código situado fuera de la región capturada continúa en modo inmediato. El anuncio de PyTorch también describe cómo establecer torch.compiler.set_default_backend("neuron") una vez antes de compilar, de modo que las llamadas de compilación posteriores puedan omitir el argumento del backend.

La versión incluye un backend neuron para torch.distributed, con operaciones colectivas y punto a punto documentadas. AWS enumera DistributedDataParallel, FSDP2, DTensor y DeviceMesh, así como el paralelismo de canalización, entre las API compatibles. Son rutas útiles que examinar para un proyecto de entrenamiento distribuido, pero las notas de la versión no demuestran que una tarea multinodo existente vaya a funcionar sin cambios ni a alcanzar un objetivo concreto de rendimiento.

La evaluación empieza con un entorno compatible

La guía rápida de instalación de AWS requiere una instancia Trn2 o Trn3 y una Neuron Base DLAMI o una AMI de Ubuntu compatible. La matriz de la versión especifica Linux x86_64 en Ubuntu 24.04 o 26.04. Ubuntu 22.04 llega al fin de su soporte en esta versión. Enumera Python 3.12, 3.13 y 3.14, con torch-neuronx 2.13.3.0 para PyTorch 2.13.0 o torch-neuronx 2.14.3.0 para PyTorch 2.14.0. La familia del paquete debe corresponder a la versión de PyTorch.

La guía rápida instala torch-neuronx==2.14.* junto con neuronx-cc==2.* y nki desde los índices de paquetes de Neuron y PyTorch. Después comprueba la disponibilidad del dispositivo y compara una pequeña operación matricial en Neuron con un resultado de CPU. Esto verifica la ruta del dispositivo instalado en el ejemplo de AWS. No valida el modelo del lector. Una siguiente evaluación razonable consiste en ejecutar las rutas hacia delante y hacia atrás del modelo real, inspeccionar las operaciones registradas que recurren a la CPU mediante torch_neuronx.get_fallback_ops(), comparar los resultados con el entorno actual y medir por separado la compilación y las ejecuciones en régimen estable. Estos últimos pasos son preguntas de evaluación propuestas por BIG CHANGE, no pruebas que hayamos realizado.

Dónde la beta aún puede cambiar el trabajo

En modo inmediato, una operación no compatible puede recurrir a la CPU y emitir una advertencia. Con torch.compile, una operación capturada en un grafo que el backend de Neuron no pueda transformar a su representación de ejecución provoca un fallo de compilación por defecto. AWS remite a los desarrolladores a su tabla de compatibilidad de operadores para consultar la cobertura de cada operación. Por tanto, una ejecución inmediata que termina correctamente no demuestra por sí sola una ejecución íntegra en el dispositivo ni compatibilidad con la compilación.

La versión también indica que torch.compile requiere dynamic=False; una nueva forma de entrada puede requerir otra compilación. AWS sugiere usar relleno o agrupación por tamaños para reducir las variaciones de forma. La primera ejecución incluye la compilación si no hay un artefacto coincidente en caché. Estos detalles importan al medir una carga de trabajo: la latencia de compilación y el patrón de las formas de entrada pueden cambiar el resultado de una evaluación.

Otras limitaciones explícitas acotan la afirmación de compatibilidad. La beta no admite torch.sparse, torch.func, torch.nested ni torch.distributions. Solo algunas transformadas de torch.fft se ejecutan en Neuron, y float64 se ejecuta como float32 con una advertencia. Esta versión no incluye la integración de TorchAO ni la compatibilidad con el tipo de datos compuesto MXFP8. Un equipo que use estas interfaces debe comprobar su propia ruta de código antes de considerar la asignación de dispositivos como un plan de migración.

El gran cambio

TorchNeuron sitúa Trainium detrás de las interfaces nativas de dispositivos y compilación de PyTorch en una beta pública. Esto hace más directa una primera prueba de los programas de PyTorch compatibles y ofrece a los usuarios actuales de Neuron una ruta distinta de su código basado en XLA. La evidencia decisiva para adoptarlo vendrá de cada carga de trabajo: si sus operadores permanecen en el dispositivo previsto, si sus resultados numéricos siguen siendo aceptables y si el rendimiento y el costo medidos justifican la instancia y el trabajo de migración. BIG CHANGE revisó la documentación y los anuncios para este análisis; no ejecutamos TorchNeuron ni verificamos de forma independiente las afirmaciones de rendimiento de AWS.

Fuentes y lecturas adicionales