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.

At MIT’s Device Realization Lab, we’re using AWS Trainium to push the limits of medical AI research. Our 3D ultrasound segmentation and speed-of-sound estimation models train faster and more efficiently than ever, cutting experimental time by more than half while achieving state-of-the-art accuracy. AWS Trainium has enabled us to scale our research in ways that were not feasible with traditional GPU systems. By training our 3D fully convolutional neural networks on AWS Trainium (trn.32xlarge), we achieved state-of-the-art performance with 50% higher throughput and lower cost compared to NVIDIA A100 instances. Using a 32-node Trainium cluster, we efficiently conducted over 180 ablation experiments, reducing total training time from months to weeks and accelerating medical AI innovation at MIT. In the future, we plan to use Trainium to train AI agent models that can operate and automate the digital ultrasound workflow, saving significant clinician time and providing better care to patients."

" CMU Catalyst research group works on optimizing ML systems. Our project aims to make it easier to optimize across different ML systems. Trainium is unique in providing both low-level control and an accessible programming interface through Neuron Kernel Interface (NKI).

with the support of AWS through the Build on Trainium program, our researcher was able to explore advanced optimizations on a critical kernel—FlashAttention. What amazed us most was the speed at which we could iterate: we achieved meaningful improvements on top of the prior state of the art in just a week using publicly available NKI, Neuron profiler, and architecture documentation. The combination of powerful tools and clear hardware insights made sophisticated, low-level optimization accessible to our team.

AWS Trainium and Neuron Kernel Interface (NKI) empowers researchers like us to innovate faster, removing barriers that typically slow down hardware-specific optimization work."

 

"Through the Build on Trainium program, his team has gained full access to AWS Neuron’s new NKI open-source compiler stack — including direct visibility into the Trainium ISA and APIs for precise scheduling and memory allocation. This level of visibility and control allows his students to more easily analyze the opportunities for optimization and more effectively discover performant implementations."

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

"Access to AWS Trainium and Inferentia has been instrumental in advancing our research and education on large-scale, efficient AI systems. We use these platforms for Mixture-of-Experts training and inference optimizations, prototyping new runtime and scheduling techniques that improve scalability, efficiency, and portability on emerging accelerator architectures. By leveraging the Neuron Developer stack, UIUC researchers are developing new runtime and scheduling techniques that advance the efficiency and portability of AI workloads. The team is particularly impressed by the openness of Neuron Developer stack, which make these platforms valuable for runtime research and enable innovations in sparsity, memory hierarchies, and communication efficiency that go beyond traditional GPU architectures."

"By leveraging AWS Trainium and the Build on Trainium program, my students and I were able to accelerate our quantum circuit simulations significantly. The project brought together a strong group of students who collaboratively built a high-performance simulator, enabling deeper experimentation and hands-on learning at a scale that simply wasn't possible before."

"Our research team at UTS is exploring the integration of tree-ring watermarking algorithms by developing custom Neuron NKI kernels. Having access to the open-source Neuron stack through the Build on Trainium program has been transformative. It gives us unprecedented visibility into the Trainium architecture and the ability to work directly at the hardware level. Access to Trainium has enabled our team to accelerate our watermarking workloads significantly, reducing iteration cycles and allowing us to explore more complex models and techniques. This depth of access allows our researchers to prototype new ideas, experiment with low-level optimizations, and push the boundaries of what watermarking systems can achieve on modern AI accelerators. "