Amazon ha anunciado TorchNeuron, un backend PyTorch nativo de código abierto para su familia Trainium de chips de IA, en una medida que elimina una de las mayores barreras prácticas para la adopción de silicio personalizado para cargas de trabajo de aprendizaje automático. Según una publicación oficial en el blog de PyTorch publicada antes de la Conferencia PyTorch de Norteamérica el 20 de octubre, los desarrolladores ahora pueden ejecutar el código del modelo PyTorch existente en Trainium con un cambio de una sola línea, en lugar de reescribir sus flujos de trabajo en torno a herramientas patentadas.
El anuncio llamó la atención de inmediato porque apunta a lo que los analistas han descrito durante mucho tiempo como la ventaja más duradera de Nvidia: no el rendimiento bruto del chip, sino el ecosistema de software CUDA que vincula a los desarrolladores con sus GPU. Reuters enmarcó el lanzamiento como si Amazon dijera que sus chips de IA funcionarán con el marco de desarrollo más popular en IA, mientras que TechCrunch calificó a TorchNeuron como la mayor oportunidad de la compañía en Nvidia hasta el momento. Para obtener una visión más amplia de cómo está cambiando el mercado de aceleradores, nuestra cobertura de noticias sobre hardware de IA continua sigue estos desarrollos a medida que se desarrollan.
Lo que realmente cambia TorchNeuron
La promesa principal de TorchNeuron es la simplicidad. En el ejemplo publicado en el blog de PyTorch, se ejecuta un bucle de entrenamiento existente en Trainium después de simplemente mover el modelo al nuevo dispositivo con `.to('neuron')`. El código escrito en `torch.accelerator`, la API independiente del dispositivo recomendada por PyTorch, se ejecuta sin modificaciones, según la publicación. Para cargas de trabajo compiladas, los desarrolladores configuran el backend predeterminado una vez con `torch.compiler.set_default_backend("neuron")` y luego invocan `model.compile()` sin más cambios.
Debajo del capó, TorchNeuron registra Trainium a través de PrivateUse1, el punto de extensión oficial de PyTorch para backends de hardware personalizados, lo que significa que el chip se comporta como un dispositivo PyTorch estándar en lugar de un apéndice atornillado al marco. AWS describe la integración como que abarca varios modos de ejecución:
- Ejecución ansiosa para investigación y depuración, con un "modo ansioso adaptativo" experimental que reduce la sobrecarga de lanzamiento
- torch.compile para rendimiento de producción, capturando modelos en gráficos que el compilador Neuron optimiza, con soporte nativo para FlexAttention
- La Neuron Kernel Interface (NKI) para investigadores que necesitan kernels personalizados de bajo nivel
El backend funciona tanto en los chips Trainium2 como en los más nuevos Trainium3, según el anuncio.
Un ataque directo al foso de software de Nvidia
El control de Nvidia sobre la informática de IA nunca se ha basado únicamente en un hardware más rápido. CUDA, la plataforma de software de la compañía con dos décadas de antigüedad, es donde se escribe, ajusta e implementa la mayor parte del código de investigación de IA, y cada chip desafiante, desde la línea Instinct de AMD hasta los TPU de Google y Trainium de Amazon, ha tenido que lidiar con el costo de pedir a los desarrolladores que lo abandonen. La cobertura del lanzamiento de TorchNeuron en GeekWire, Benzinga y GuruFocus encuadró uniformemente el anuncio como un asalto a ese foso.
Amazon no está intentando sacar a los desarrolladores de PyTorch para ganárselos. En cambio, está incorporando Trainium a PyTorch, para que la enorme base de usuarios del marco pueda adoptar el chip sin cambiar sus hábitos. La publicación del blog de PyTorch incluyó un respaldo de Meta, que mantiene el marco: "El soporte nativo para AWS Trainium amplía esa elección, permitiendo a los investigadores e ingenieros experimentar e iterar utilizando la ejecución entusiasta y las API distribuidas que ya conocen", dijo Jana van Greunen, directora de ingeniería de PyTorch en Meta, en la publicación.
Ese marco es importante para el negocio de la nube de AWS. Las instancias de Trainium se han comercializado como una alternativa de menor costo a las GPU de Nvidia para entrenamiento e inferencia, y los chips ya han atraído a importantes clientes de IA. Hacer que los chips parezcan nativos de PyTorch aborda directamente la objeción que ha frenado una adopción más amplia: el esfuerzo de portabilidad.
Alcance técnico, desde cuadernos hasta capacitación a escala de clúster
Las capacidades enumeradas en el blog de PyTorch sugieren que Amazon apunta mucho más allá de los experimentos de aficionados. TorchNeuron cubre los conjuntos de operadores vistos en modelos de texto, multimodales y de lenguaje de visión, así como generación de imágenes y videos, modelos de espacio de estados y modelos mundiales, con precisión BF16, FP32 y FP16; El soporte de TorchAO y FP8 se encuentra en versión beta privada antes de la disponibilidad general en la próxima versión.
La capacitación distribuida, la disciplina que convierte el código de una sola GPU en ejecuciones a escala de clúster, se admite a través de las API estándar de PyTorch: DDP, FSDP2, SimpleFSDP, paralelismo de tensor y canalización y fragmentación estilo ZeRO, junto con DTensor y DeviceMesh para grupos de procesos multidimensionales. Las transmisiones y los eventos asignan la semántica de PyTorch a las colas de hardware de Trainium, según la publicación.
A modo de inferencia, las pilas Neuron de Amazon se basan en la misma base: vLLM Neuron ofrece modelos de lenguaje grandes, mientras que vLLM Omni Neuron ofrece modelos de difusión para generación de imágenes y vídeos. La interfaz Neuron Kernel permite a los investigadores escribir operadores personalizados para ejecución activa y compilada, y AWS es NKI-lib de código abierto, un conjunto de núcleos optimizados para Trainium que cubren cálculos en bloque FlashAttention y MoE para escenarios multimodales y LLM convencionales.
También se incluyen herramientas de rendimiento. Neuron Explorer extiende el estándar `torch.profiler` para rastrear la ejecución desde el marco hasta el hardware, vinculando cada evento a su línea de origen y puede exportar métricas con formato Parquet que los agentes pueden consumir en las notas de publicación para el análisis automatizado de cuellos de botella.
Disponibilidad y lo que viene después
TorchNeuron está programado para su lanzamiento de código abierto en la Conferencia PyTorch de Norteamérica el 20 de octubre, según la publicación del blog, y la compatibilidad se extiende a Hugging Face Transformers v5. AWS describe Trainium como una cartera de chips diseñados desde cero para el entrenamiento e inferencia de IA a escala de la nube.
El contexto estratégico es un mercado en el que todos los principales proveedores de nube son simultáneamente el mayor comprador de aceleradores Nvidia y un posible competidor. Microsoft, Google y Meta desarrollan silicio personalizado, y empresas chinas, incluida Huawei, están construyendo sus propias pilas de software para escapar de la gravedad de CUDA. Al reunirse con los desarrolladores dentro de PyTorch en lugar de pedirles que migren, Amazon ha elegido el camino menos disruptivo en ese concurso y el más probable, si la experiencia del desarrollador se mantiene, para trasladar cargas de trabajo de capacitación reales del hardware de Nvidia.
Que TorchNeuron realmente erosione la ventaja de CUDA dependerá del rendimiento a escala y de la larga cola de modelos del mundo real, preguntas que los materiales de lanzamiento no responden. Pero la barrera que impedía a la mayoría de los desarrolladores de PyTorch probar Trainium (una reescritura) ahora se ha reducido a una sola línea de código.
---
Manténgase a la vanguardia de la IAObtenga las últimas noticias, análisis y avances en IA, todo en un solo lugar.
Leer más noticias sobre IA →
