This is a one-page ASPLOS 2024 keynote abstract rather than a full paper — Nafea Bshara of Amazon recounts the journey of building three generations of AWS Trainium, the company's custom machine-learning training hardware, and draws one central lesson, that because the bottleneck in training a large model keeps moving between the chip, its instruction set, the server, the compiler, and the network, owning and co-designing all of those layers together beats optimizing any single one in isolation.
The argument in one glance
-
Problem
A fast training chip is worthless if any other layer starves it
The abstract opens by naming the challenge: machine-learning accelerators pose design problems across chip architecture, instruction set, server design, compiler, and both inter- and intra-chip connectivity. These are entangled. A chip with enormous compute, sitting in a server that cannot feed it, on a network that stalls during communication, programmed by a compiler that cannot schedule it well, is slow overall. Making one layer fast while another lags does not remove the bottleneck; it moves it.
-
What people did
Assemble a cluster from the best merchant parts
The common path is to buy general-purpose accelerators from a vendor, wire them into servers with standard networking, and program them through a general machine-learning software stack. Each layer is designed by a different party to a fixed interface.
-
Why it fell short
When no one owns the whole stack, the moving bottleneck is nobody's to chase
Fixed interfaces between separately owned layers mean each is optimized in isolation. But the true bottleneck in training a large model shifts over time and across workloads — now compute, now on-chip memory, now the compiler's schedule, now the network during collective communication. No single-layer owner can follow it across a boundary they do not control, so accelerators sit stranded waiting on data.
-
This paper's idea
Own the whole stack and co-design it, generation after generation
The keynote's thesis is that AWS's end-to-end ownership — chip, server, network, compilers, and runtime — let one organization co-design and optimize across all the layers at once, deciding for each problem which layer should absorb it, and folding the lessons from real customer workloads into the next chip, instruction set, and compiler across three product generations, all while treating simplicity and ease of use as a first-class goal.
Why this problem is worth solving
Training a large model is one of the most expensive computations organizations routinely run, and the expense is why the design of the hardware underneath it matters so much. A single frontier model can occupy thousands of accelerators for days or weeks, and every percent of that fleet left idle — waiting on memory, on the compiler, on the network — is real money and real time. That waiting is the crux. Full-stack co-design is the practice of designing the layers of a computing system — here the accelerator chip, the instruction set that exposes it, the server that houses many chips, the compiler and runtime that program it, and the interconnect that links the chips — together, as one problem, rather than each in isolation behind fixed interfaces. Its whole justification is that the layers interact: the best instruction set depends on what the compiler can exploit, the right amount of on-chip memory depends on how the network delivers data, and so on. When one team owns and can change all the layers, it can move a design decision to whichever layer solves it most cheaply, instead of accepting a boundary drawn by a vendor’s product line.
The layer people most often forget is the one connecting the chips, and at scale it can dominate. Collective communication refers to the group data-exchange operations that distributed training depends on, the best known being AllReduce, in which every accelerator’s partial result (for example, the gradients it computed) is summed with every other’s and the total handed back to all of them, so the many chips training one model stay in agreement. These collectives happen after each step of training, over and over, and while they run the expensive compute engines can sit idle waiting for data to arrive. As models and clusters grow, the fraction of time spent in communication grows with them, which is why the network is not a detail bolted on after the chip but a layer that has to be designed with the chip’s compute and memory in view — exactly the kind of cross-layer decision that full-stack ownership exists to make.
What people did before, and where it broke
- Buy general-purpose accelerators and assemble clusters. The default, and a reasonable one: merchant GPUs are capable, widely supported, and someone else carries the design risk. Where it broke. Generality is paid for in power, cost, and fixed interfaces; because the chip, its compiler, and the network are owned by different parties, a bottleneck that lives at a boundary between them is no one’s to remove.
- Design just the chip, buy the rest. Building a custom accelerator ASIC captures the biggest, most visible cost. Where it broke. A superb chip in a weak server, on a generic network, driven by a compiler that does not know its internals, is starved — the effort moves the bottleneck to the layers left untouched rather than removing it.
- Rely on a generic machine-learning compiler. Standard frameworks and compilers let any hardware run models without bespoke software. Where it broke. A compiler that cannot see a chip’s exact memory sizes, data paths, and collectives cannot schedule work to keep it fed, so hardware built to be driven by a generic stack leaves much of its peak on the table.
- Scale out with off-the-shelf networking. Ordinary datacenter networking connects the accelerators and is well understood. Where it broke. Distributed training is dominated by collective communication; a network not co-designed with the chip’s collective throughput leaves accelerators idle during every AllReduce, and the idle fraction grows with scale.
- Treat each hardware generation as a fresh design. Redesigning cleanly each time can incorporate the newest ideas. Where it broke. Without deliberately carrying lessons forward, each generation risks repeating the last one’s mistakes; the compounding advantage comes from iterating the whole stack against what real workloads actually did.
- The gap. Every one of these is a point solution owning a slice of the stack. None can chase a bottleneck that lives between slices. Only an owner of the entire stack can decide, per problem, which layer should absorb the cost — and keep making that decision better across generations.
How the layers fit together
The layers a full-stack ML training system spans
-
-
new
-
new
-
new
-
new
-
new
The idea
Owning every layer lets you put each bottleneck where it is cheapest to remove. The recurring frustration of assembling a system from separately owned parts is that the interfaces are fixed by other people’s products, so when a bottleneck appears at a boundary — the chip could go faster but the compiler cannot express it, or the memory is ample but the network cannot fill it — no one is positioned to move the cost to the layer best able to bear it. End-to-end ownership dissolves that constraint: if a problem is awkward to solve in silicon, perhaps the compiler can, and if the compiler cannot, perhaps the instruction set should change to let it, and if the network is the limit, the chip’s data layout can adapt to the network rather than the reverse. The design decision is made once, globally, for the whole stack, instead of locally within each vendor’s slice. Choosing jointly across layers that would otherwise be optimized separately — deciding, for instance, how on-chip memory and communication should trade against each other as one question rather than two — is the same co-exploration the accelerator memory-versus-communication design work covered elsewhere on this site formalizes, searching the layers together because the best point in one depends on the choice in the other.
The software is what makes the hardware usable, so ease of use is a design goal, not an afterthought. A fast accelerator that is painful to program is slow in practice, because real workloads never reach a peak that only a heroic hand-tuning effort could touch. The abstract singles out simplicity and ease of use as things AWS emphasized, which is a claim about where value is actually captured: in the compiler and runtime that let an ordinary user’s model reach the hardware’s potential without becoming an expert in its internals. This only works when the same owner who shapes the instruction set also builds the compiler, so the two are designed to fit — the hardware exposes what the compiler needs, and the compiler hides what the user should not have to think about. Making a diverse, awkward piece of hardware approachable by putting the burden on a well-designed software stack, so that using it does not require understanding it, is the same division of labour the computing-in-memory compilation-stack work covered elsewhere on this site is built around, absorbing the hardware’s complexity into the compiler so the person above does not have to carry it.
Distributed training is a communication problem as much as a compute one, so the network is a co-designed layer. It is tempting to think of a training accelerator as fundamentally a compute device with networking attached, but at the scale where large models are trained, the collective communication between chips can be the thing that decides overall speed, because the compute engines stall whenever the exchange of gradients lags. Treating the interconnect as a first-class part of the design — sized to the chip’s throughput, shaped so communication overlaps computation, exposed to the compiler so it can schedule around it — is precisely the kind of cross-layer decision that full-stack ownership makes possible and that buying the network separately makes hard. Co-designing the paths that carry collective operations so the accelerators are not left idle waiting on each other is the same concern the GPU-cluster communication-path work covered elsewhere on this site takes up, treating how the collectives traverse the links as something to be designed deliberately rather than left to a default.
An illustrative co-design decision
How a single training bottleneck can be pushed across the layers
This traces the general reasoning that full-stack ownership enables, using textbook accelerator design. It is an illustration of the co-design principle the keynote describes, not a disclosed detail of any specific Trainium design.
Step 1 of 5
Notice the tension
During a training step, the chip's tensor engine can consume weights and activations faster than off-chip memory, or the network during gradient exchange, can supply them. Left alone, the expensive compute sits idle waiting for data.
This is the universal shape of the accelerator problem: arithmetic is cheap and plentiful, feeding it is the hard part, and the layer that limits the feeding is the one that sets the real speed.
Why it works
- The bottleneck moves, and only a whole-stack owner can follow it. In training, the limiting layer shifts between compute, memory, compiler, and network across time and workloads → and a party owning only one layer cannot chase a bottleneck across a boundary it does not control → so end-to-end ownership is what lets the cost be placed wherever it is cheapest to remove.
- An explicit instruction set pays off only with a matched compiler. Exposing on-chip memory and data movement gives software room to keep the compute fed → but that room is usable only if the compiler understands the hardware intimately → so co-owning the instruction set and the compiler is what turns exposed hardware into sustained throughput rather than idle silicon.
- Co-designing the network with the chip removes idle time at scale. Distributed training stalls the compute whenever collective communication lags → but sizing and shaping the interconnect to the chip’s gradient throughput lets communication overlap computation → so the accelerators keep working through the exchange instead of waiting on it.
- Ease of use is what converts peak capability into delivered performance. Real workloads only reach hardware built to be approachable → because a stack that demands deep expertise is left un-exploited by ordinary users → so treating simplicity as a design goal is what makes the co-design pay off in practice rather than only on a datasheet.
- Generational iteration compounds the advantage. Each product generation is measured against what real customer workloads actually did → and those lessons feed the next chip, instruction set, and compiler together → so the whole stack improves against reality repeatedly, which a one-off design cannot match.
What it costs (trade-offs and scope)
- This is a keynote abstract, so there is nothing to verify. The one-page source publishes no architecture, no numbers, and no evaluation; the reasoning above is general background on why full-stack co-design helps, not evidence that Trainium achieves any particular result. The honest reading is an orientation to the problem, with the specifics deferred to the talk itself and to AWS’s separate technical documentation.
- Full-stack ownership is available to almost no one. It requires a single organization that owns the chip, the server, the network, the compiler, and a large customer base to justify the investment — a strategy open to a few hyperscalers, not to most teams, who will rightly keep buying merchant hardware.
- Fit is bought with generality. A stack co-designed for one owner’s workloads may serve others poorly; merchant accelerators remain more flexible across diverse and unforeseen uses, and that flexibility is a real value the co-designed approach gives up.
- Vertical integration concentrates risk. When one team owns every layer, a misstep in any of them is theirs alone to fix, and the layers must advance together across generations — a coordination burden that separately owned stacks distribute.
- The ease-of-use claim is asserted, not shown here. Simplicity is named as a goal, but the abstract gives no way to judge how well it was reached, against what baseline, or at what cost — that evidence, if it exists, lives outside this source.
- The “three generations of lessons” are promised, not recorded. The abstract says the talk will illustrate the design principles, tradeoffs, and lessons learned; the written artifact does not contain them, so this explainer cannot report what they were.
- Co-design can become customer lock-in. A custom compiler and runtime that make a stack fast can also make workloads harder to move elsewhere — a cost borne by users of the platform, and one a vendor keynote has little reason to dwell on.
The evaluation
Three ways to build a machine-learning training system
| Property | Assemble merchant accelerators | Design only a custom chip | Co-design the full stack |
|---|---|---|---|
| Who owns the boundaries between layers | |||
| What happens when the bottleneck moves | |||
| Flexibility across diverse workloads | |||
| Who can realistically pursue it | |||
| How well the compiler knows the chip |
Check yourself
Check your understanding
1. What kind of source is this catalogue entry, and what does that mean for its contents?
2. What is the central argument the abstract makes about designing ML training hardware?
3. Why does the abstract emphasize simplicity and ease of use as design goals?
4. Which limitation of the full-stack co-design strategy is most important to keep in mind?
Sources
The item is “AWS Trainium: The Journey for Designing and Optimization Full Stack ML Hardware” — a keynote abstract by Nafea Bshara (Amazon), published at ASPLOS 2024, Volume 3, as a one-page entry, DOI 10.1145/3620666.3655592.
This explainer is written from that one-page abstract, which is the only text available. The abstract states the talk’s theme — that AWS used its end-to-end ownership of chip, server, network, compilers, and runtime tools to co-design and optimize across all layers, emphasizing simplicity and ease of use, across three generations of ML products, and that the talk would illustrate the design principles, tradeoffs, and lessons learned. Because the abstract contains no architecture details, measurements, or evaluation, the concrete explanations above — of systolic-array compute, collective communication, instruction-set and compiler co-design, and interconnect sizing — are general, publicly known background on machine-learning accelerator design offered to make the keynote’s theme legible, not disclosures or results from this source. Nothing here should be read as a specific claim about Trainium’s internals or performance; the frontmatter is marked draft accordingly, and readers wanting technical specifics should consult AWS’s separate Trainium and Neuron documentation.