Build on Trainium

Build on Trainium, is a $110M investment program focused on AI research and university education to support the next generation of innovation and development on AWS Trainium. AWS Trainium is an AI systolic array chip uniquely designed for advancing state-of-the-art AI ideas and applications. Build on Trainium funds novel AI research on Trainium, investing in leading academic teams to build innovations in critical areas including new model architectures, ML libraries, optimizations, large-scale distributed systems, and more. This multi-year initiative lays the foundation for the future of AI by inspiring the academic community to leverage, invest in, and contribute to the open-source community around Trainium. Combining these benefits with Neuron software development kit (SDK) and recent launch of the Neuron Kernel Interface (NKI), Trainium customers can now innovate at scale in the cloud.

We have created a dedicated Trainium research cluster with up to 40,000 Trainium chips that will be available through Amazon EC2 Trn1 instances connected on a single non-blocking peta-bit scale network using Amazon EC2 UltraClusters. Research teams and students can access these chips through self-managed capacity block reservations using Amazon EC2 Capacity Blocks for ML.

We are conducting multiple rounds of Amazon Research Awards (ARA) call for proposals (CFP) to the broad research community, with selected proposals receiving AWS Trainium credits and access to the Trainium research cluster. Build on Trainium welcomes research proposals that will leverage popular open-source ML libraries and Frameworks, and contribute back to open-source to enhance resources for the ML developer community.

Neuron Kernel Interface (NKI) is a new programming interface for AWS AI chips, Trainium and Inferentia. NKI provides direct access to hardware primitives and instructions available on AWS Trainium and Inferentia, enabling researchers to build and tune compute kernels for optimal performance. It is a Python-based programming environment which adopts commonly used Triton-like syntax and tile-level semantics. Researchers can use NKI to enhance deep learning models with new functionalities, optimizations, and science innovations. Please visit the NKI documentation page to learn more.

" Trainium is beyond programmable—not only can you run a program, you get low-level access to tune features of the hardware itself. The knobs of flexibility built into the architecture at every step make it a dream platform from a research perspective. AWS is really enabling unexpected innovation. I walk across the lab and every project needs compute cluster resources for something different. The Build on Trainium resources will be immensely useful—from day-to-day work, to the deep research we do in the lab. "

Christopher Fletcher, Associate Professor of Computer Science, University of California, Berkeley
 

" AWS’s new Build on Trainium initiative enables our faculty and students large scale access to modern accelerators, like AWS Trainium, with an open programming model, and allow us to greatly expand our research on tensor program compilation, ML parallelization, and language model serving and tuning. "

Todd C. Mowry, Professor of Computer Science, Carnegie Mellon University 

" Trainium accelerates our generative AI research across multiple
scientific domains. It's a game-changer for academics looking to scale. "

Adam Kilvans, Professor of Computer Science, University of Texas

" The AWS Build on Trainium program greatly strengthens my lab’s research at Oxford on neuro-symbolic AI to revolutionize hardware design. Access to Trainium accelerators allows our researchers and students to innovate at an ambitious and impactful scale. "

Tom Melham, Professor of Computer Science, University of Oxford