AWS Trainium Competition: Co-designing Models and Kernels for AI Chips
Exploring model and kernel design for custom AI chips on AWS Trainium.
1Modern AI Models and Hardware Relationship
Modern large language models (LLMs) have evolved within a specific hardware family. The shape of attention mechanisms, the structure of multilayer perceptrons (MLPs), numerical format choices, and even the granularity of parallelism strategies have been shaped by hardware constraints. AWS Trainium, a custom AI chip, redefines the efficient frontier of model architectures by altering these constraints.
2Unique Features of AWS Trainium
AWS Trainium is a purpose-built AI chip. It offers features like more on-chip SRAM (SBUF), explicit software control over data movement, and energy-efficient systolic matrix multiplication. This changes the memory bandwidth-to-TFLOPs ratio, potentially turning memory-bound operations on conventional chips into compute-bound operations on Trainium.
3Competition Goals and Participants
The AWS Trainium Frontier competition aims to enable academic and industry labs to explore model architectures on this new hardware. Participants train language models from scratch on Trainium, investigating the full design space from model architecture to custom kernel design. The competition rewards full-stack thinking under a fixed time and compute budget.
4Two Phases of the Competition
The competition consists of two phases. In the first phase, each team receives a single Trainium2 chip and a 30-minute training budget. In the second phase, the top 10 teams gain access to a full Trainium2 server with a four-hour budget. This phase allows techniques like distributed parallelism and communication-aware model shaping.
5Evaluation Criteria
The competition evaluates model training efficiency and inference performance. In the first phase, a single number, validation bits-per-byte (val_bpb), is used. In the second phase, the CORE score, which spans reasoning, comprehension, and world knowledge tasks, is added. The final score is an equal-weighted combination of these two axes.
6Flexibility for Participants
Participants can use different strategies to improve the model. A better learning rate schedule can be as important as a faster kernel. This rewards full-stack thinking: a novel attention mechanism is only as fast as the kernel that runs it, and the fastest kernel only matters if the architecture knows how to use it.
7Model Size and Time Constraint
Model size is uncapped; the constraint is not parameters but the time spent on the chip. Participants choose the model architecture that maximizes capability within a fixed training window. This inversion of the traditional scaling paradigm makes it a more challenging research question.
8Competition Contribution and Rewards
The competition contributes to the discovery of model architectures designed for custom AI chips. Winners can pioneer a new area of machine learning research. Academic teams can obtain AWS Promotional Credits for Trainium compute and Amazon Bedrock access. The competition enables hardware-native model designs on AI chips.
Original article: Amazon Science — AWS Trainium Yarışması: AI Çipleri için Model ve Kernel Tasarımı