Build on Trainium

Build on Trainium es un programa de inversión de 110 000 000 USD centrado en la investigación de la IA y la educación universitaria para apoyar la próxima generación de innovación y desarrollo en AWS Trainium. AWS Trainium es un chip de matriz sistólica de IA diseñado exclusivamente para promover ideas y aplicaciones de IA de última generación. Build on Trainium financia investigaciones novedosas sobre IA en Trainium, mediante la inversión en equipos académicos líderes para crear innovaciones en áreas críticas, como arquitecturas de nuevos modelos, bibliotecas de ML, optimizaciones, sistemas distribuidos a gran escala y más. Esta iniciativa de varios años sienta las bases para el futuro de la IA al inspirar a la comunidad académica a aprovechar, invertir y contribuir a la comunidad de código abierto en torno a Trainium. Al combinar estos beneficios con el kit de desarrollo de software (SDK) de Neuron y el reciente lanzamiento de la Neuron Kernel Interface (NKI), los clientes de Trainium ahora pueden innovar a escala en la nube.

Hemos creado un clúster de investigación de Trainium dedicado con hasta 40 000 chips de Trainium que estarán disponibles a través de instancias Trn1 de Amazon EC2 conectadas a una única red a escala de petabits sin bloqueo mediante Amazon EC2 UltraClusters. Los equipos de investigación y los estudiantes pueden acceder a estos chips mediante reservas de bloques de capacidad autoadministradas usando bloques de capacidad de Amazon EC2 para ML.

Estamos realizando varias rondas de convocatorias de propuestas (CFP) de los Amazon Research Awards (ARA) dirigidas a la comunidad investigadora en general y las propuestas seleccionadas recibirán créditos de AWS Trainium y acceso al clúster de investigación de Trainium. Build on Trainium da la bienvenida a las propuestas de investigación que aprovecharán las populares bibliotecas y marcos de machine learning (ML) de código abierto y contribuirán al código abierto para mejorar los recursos para la comunidad de desarrolladores de ML.

Neuron Kernel Interface (NKI) es una nueva interfaz de programación para los chips de IA de AWS, Trainium e Inferentia. NKI proporciona acceso directo a las primitivas de hardware y a las instrucciones disponibles en AWS Trainium e Inferentia, lo que permite a los investigadores crear y ajustar kernels de computación para lograr un rendimiento óptimo. Es un entorno de programación basado en Python que adopta una sintaxis similar a la de Triton y una semántica a nivel de mosaico de uso común. Los investigadores pueden usar NKI para mejorar los modelos de aprendizaje profundo con nuevas funcionalidades, optimizaciones e innovaciones científicas. Visite la página de la documentación de NKI para obtener más información.

«En el laboratorio de desarrollo de dispositivos del MIT, utilizamos AWS Trainium para traspasar los límites de la investigación médica con inteligencia artificial. Nuestros modelos de segmentación por ultrasonido en 3D y estimación de la velocidad del sonido se entrenan más rápido y de manera más eficiente que nunca, lo que reduce el tiempo de experimentación a más de la mitad y logra una precisión de vanguardia. AWS Trainium nos ha permitido escalar nuestra investigación de maneras que no eran factibles con los sistemas de GPU tradicionales. Al entrenar nuestras redes neuronales 3D totalmente convolucionales en AWS Trainium (trn.32xlarge), logramos un rendimiento de vanguardia un 50% más alto y un costo menor en comparación con las instancias NVIDIA A100. Con un clúster Trainium de 32 nodos, realizamos de manera eficiente más de 180 experimentos de ablación, lo que redujo el tiempo total de entrenamiento de meses a semanas y aceleró la innovación médica con inteligencia artificial en el MIT. En el futuro, planeamos usar Trainium para entrenar modelos de agentes de IA que puedan operar y automatizar el flujo de trabajo de la ecografía digital, para ahorrar mucho tiempo a los médicos y brindar una mejor atención a los pacientes».

«El grupo de investigación CMU Catalyst trabaja en la optimización de los sistemas de aprendizaje automático. Nuestro proyecto tiene como objetivo facilitar la optimización en diferentes sistemas de aprendizaje automático. Trainium es único porque nos proporciona al mismo tiempo un control de bajo nivel y una interfaz de programación accesible a través de Neuron Kernel Interface (NKI).

Con el apoyo de AWS a través del programa Build on Trainium, nuestro investigador pudo explorar optimizaciones avanzadas en un núcleo fundamental: FlashAttention. Lo que más nos sorprendió fue la rapidez con la que pudimos iterar: logramos mejoras significativas con respecto al estado actual de la técnica en solo una semana, utilizando documentación sobre NKI, perfiladores de neuronas y arquitectura disponibles al público. La combinación de herramientas potentes y conocimientos claros sobre el hardware hizo que nuestro equipo pudiera acceder a una optimización sofisticada y de bajo nivel.

AWS Trainium y Neuron Kernel Interface (NKI) permiten a los investigadores como nosotros innovar más rápido y eliminar las barreras que suelen ralentizar el trabajo de optimización específico del hardware».

 

«Gracias al programa Build on Trainium, su equipo ha obtenido acceso total a la nueva pila de compiladores de código abierto NKI de AWS Neuron, que incluye visibilidad directa de la ISA y las API de Trainium para una programación y una asignación de memoria precisas. Este nivel de visibilidad y control permite a sus estudiantes analizar con mayor facilidad las oportunidades de optimización y descubrir de manera más efectiva las implementaciones que funcionan».

Christopher Fletcher, profesor asociado de Ciencias de la Computación, Universidad de California, Berkeley
 

«El acceso a AWS Trainium e Inferentia ha sido fundamental para avanzar en nuestra investigación y educación sobre sistemas de IA eficientes y a gran escala. Usamos estas plataformas para la capacitación de expertos y las optimizaciones de inferencia, creando prototipos de nuevas técnicas de ejecución y programación que mejoran la escalabilidad, la eficiencia y la portabilidad en las arquitecturas de aceleradores emergentes. Al aprovechar el paquete de desarrolladores de Neuron, los investigadores de la UIUC están desarrollando nuevas técnicas de ejecución y programación que mejoran la eficiencia y la portabilidad de las cargas de trabajo de IA. El equipo está particularmente impresionado por la apertura del paquete de desarrolladores de Neuron, que hace que estas plataformas sean valiosas para la investigación en tiempo de ejecución, además de permitir innovaciones en cuanto a la escasez, las jerarquías de memoria y la eficiencia de la comunicación que van más allá de las arquitecturas de GPU tradicionales».

«Al aprovechar AWS Trainium y el programa Build on Trainium, mis alumnos y yo pudimos acelerar significativamente nuestras simulaciones de circuitos cuánticos. El proyecto reunió a un grupo sólido de estudiantes que, de manera colaborativa, crearon un simulador de alto rendimiento, lo que permitió una experimentación más profunda y un aprendizaje práctico a una escala que antes simplemente no era posible».

«Nuestro equipo de investigación de la UTS está estudiando la integración de algoritmos de marcas de agua de anillos de árboles mediante el desarrollo de núcleos Neuron NKI personalizados. Tener acceso a la pila de Neuron de código abierto a través del programa Build on Trainium ha sido transformador. Nos brindó una visibilidad sin precedentes de la arquitectura Trainium y nos permitió trabajar directamente a nivel de hardware. El acceso a Trainium ha permitido a nuestro equipo acelerar significativamente nuestras cargas de trabajo con marcas de agua, lo que reduce los ciclos de iteración y nos permite explorar modelos y técnicas más complejos. Esta profundidad de acceso permite a nuestros investigadores crear prototipos de nuevas ideas, experimentar con optimizaciones de bajo nivel y ampliar los límites de lo que los sistemas de marcas de agua pueden lograr con los aceleradores de IA modernos».