Accelerating AI Training and Inference with AWS Trainium2 with Ron Diamant - #720

0:00So, Antropic is an amazing customer to work with Project Rainier. we're building a gigantic training cluster that fully embraces the scaling laws. That cluster will have many hundreds of thousands of training to devices. That cluster is at least 5x larger than Anthropic's previous training cluster. And they're planning to train the largest and most intelligent frontier model on top of that cluster.

0:42All right, everyone. Welcome to another episode of the TwiML AI podcast. I am your host, Sam Charrington. Today, I'm joined by Ron Diamant. Ron is Chief Architect for Tranium at Amazon Web Services. Ron, welcome to the podcast. Thanks for having me, Sam. I'm super excited to have you on the show. I'm looking forward to our chat. You joined AWS as part of the Annapurna Labs acquisition, and we're connecting at a pretty big moment for the Annapurna team, including the 10th anniversary of that acquisition, the Tranium 2 launch at last year's reInvent, and some pretty big customer announcements with Apple and Anthropic.

1:19We'll be digging into all of these and more, and I'd love to have you get us started by sharing a bit about your background. Sure, yes. So first of all, thanks again for having me today. I joined Annapurna Labs in 2011, just a few weeks after it was founded, very early on. And since then, I've been fortunate to participate in all of the different product lines that we built in Annapurna Labs. That includes Nitro, where we basically offloaded the hypervisor from the host CPU to dedicated chips. And we reinvented how virtualization in the cloud happens. Then we built Graviton, our host CPUs that are ARM-based and provide the best performance and price performance in the cloud for a variety of cloud workloads.

2:07And in the past eight years or so, I've been running the architecture team for Trainium, where we build dedicated, highly performant and highly efficient AI ML accelerators for deep learning workloads and the training of frontier models. I was at the reInvent when those first chips were announced and even when the acquisition was announced. And it's hard to believe that it's been 10 years already. We've been amazingly lucky to get acquired by AWS and get integrated into that team. I can tell you personally that this is the second acquisition that I've been through in my career. And the first acquisition, I won't mention the names, but it wasn't as smooth as Anapurna Labs and AWS.

3:00But AWS's culture is like a sea of startups. And the startups are, on the one hand, very autonomous and progressing to their roadmaps and milestones with minimum dependencies on other teams. But on the other hand, they share roadmaps and share information extremely collaboratively and openly. So it's been a blast for us at Annapurna building chips for internal customers in AWS that tell us exactly where the pain points are and what they're trying to achieve. And at the same time, keeping our startup culture and moving very quickly to build devices and solutions in new grounds. So how many chips has it been all told since the acquisition?

3:47So since the acquisition, it's been a decade and we've built more than 10 chips, so more than a chip per year. It's across three product lines, as I mentioned before, Nitro, Graviton and Tranium. And each one of these product lines is shipping in many millions and in some cases, tens of millions of parts. and almost all of our chips got to these volumes in the very first tape out, which is quite unique in the semiconductor industry. Tape out is like when a design is released to manufacturing? Oh yeah, exactly. Yes, that's what it is. And in the chip world, we basically work extremely hard before we release the design to manufacturing because at that point you trigger an activity, the manufacturing activity that costs tens of millions of dollars.

4:35So we work extremely hard to make sure that the device is healthy and well tested and can support scaling into volumes before we trigger manufacturing. One of the things that's always struck me in my conversations with folks about hardware in general and chips in particular is really as a result of those lead times being so long and how important it is to, you know, get a lot of things right, kind of catch the market. And we've seen kind of waves of accelerator startups come over the past, you know, 10 years or more. And it strikes me that we're in a really interesting point in time right now because we're finally seeing like a convergence or, you know, critical mass at scale of a particular type of workload in the form of LLMs and transformers.

5:29Whereas before, it struck me that there were a lot of companies trying to innovate around a more generalized, you know, work set of workloads in deep learning. You know, so some folks would specialize on convolutional architectures, but there wasn't enough volume there to really, you know, scale it the way we're seeing with LLMs. Does that resonate with you, first of all? And then, like, how do you think about, you know, what you're doing there in that context? It absolutely does. So I think we're seeing the convergence of two market forces roughly at the same time. So first of all, the transformer architecture took the industry by storm.

6:14And many, many workloads are mapping towards some variance of the transformer architecture, which means that we can dedicate our efforts more towards optimizing that architecture as long as we keep enough generality. And we should probably touch on that more as we talk through the details. But the other force that we're seeing across the industry is the force of scaling laws. So we now have about a decade or slightly more than a decade where the scaling laws are empirically proving themselves again and again. And what the scaling laws mean is that the more compute we put into training of a model, the better intelligence or accuracy we get on the other side.

7:00Now, if I think about the convergent of these two forces, it means that there are a few companies out there or a few players out there that are spending a lot of compute in order to train frontier models. And most of them are deploying some form of a transformer architecture with a few tweaks. Everyone has their own secret sauce. Now, what that means is that instead of kind of spreading our efforts in supporting a wide variety of workloads, we can dedicate our efforts, our engineering efforts, into being extremely efficient in these massive scale workloads. And by doing that, we can enable the folks that are training for frontier models to be even more efficient in their deployments and thus spend more compute during training and get more intelligent models overall.

7:59Now, I'm simplifying things to some extent here because the last thing that I think we want to do is build a chip that can only do Transformer. I don't think that's wise by any stretch of the imagination because new innovation will come through and none of us can anticipate what new operators will get invented in the next couple of years. But we can focus our engineering efforts where it matters the most because the majority of the industry is driven by a few workloads and not thousands of workloads. And so when you think about architecting a chip, how do you balance, you know, where the workload is today, projecting forward into time to where the workload and the customer will be tomorrow and manage the requirements for the chip that you're trying to build?

8:50Yeah, definitely. So I think we start from a concept that may sound very simple, but I think it's very powerful. It's actually a quote from Jeff Bezos from back in the days where people spend a lot of time trying to anticipate how the world is going to change in the next five years, but they don't spend enough time thinking about what's not going to change in the next five years. And I think it's quite powerful to think about what's actually not going to change. So, for example, we know that folks are going to want more compute performance. We know that they're going to want better cost structures.

9:29We know that power efficiency is going to be a big deal in training and in ML in general. And we also know that folks will want the flexibility to innovate on top of our platform because they don't want to be confined to just one or five workloads, but rather they want the ability to invent the new secret sauce that will make their models perform better than others. So the first thing that we do is to make sure that we keep our design very efficient across these different dimensions. And we can go on a very long tangent on what each of these dimensions mean, but I'll try to keep it short. For example, performance doesn't only mean compute performance.

10:20Performance is actually driven by quite a few dimensions of the design. So by getting to know the workload, we see that some workloads are dominated or bound by our compute performance. So we make sure to pack a lot of compute in our chips. But other workloads are bound by memory bandwidth. For example, when we interact with a chatbot, it processes your input request, your prompt, and that portion tends to be very compute bound. But after that, it enters what we call an autoregressive stage, where it generates one token or one word after the other. And that portion of the workload tends to be very memory bandwidth bound.

11:04So we pack a lot of memory bandwidth into our chips. We also pack a lot of memory capacity and a lot of network bandwidth into our chips because other workloads, and I'll kind of skip through the details for now, other workloads tend to be either memory capacity bound or network bound. So we make sure that we balance the design across all these dimensions. But then there's also a second level of design considerations that is probably what you were referring to with your question, which is what actual building blocks do we put into our chips? And it's funny because for the very first chip that we built, we started building it before transformers existed.

11:51So we didn't even know about the building blocks of a transformer when we designed that chip. But what we did do back in the days is to break a set of popular workloads into primitives and then the first goal or the first attempt was to make sure that we support all the different primitives that serve as building blocks to supporting all these different workloads but then the second step after that was to kind of zoom out and try to generalize each one of these primitives such that if any new operator will get invented there will be a high chance that our generalized primitives would be able to support it.

12:32And then what we do after that, just two more sentences on this topic. What we do after that is we take a few workloads that we didn't see before and test our hypothesis on these workloads. We basically try to see whether, for example, if we design the chip to support language workloads and vision workloads, we test ourselves on an audio workload. and we try to see whether new kernels emerge that we don't support yet or we actually support our generalized set of kernels do support all the different workloads. And one of our kind of fun experiences during the design of Tranium is that when we were building the first chip, which was named Inferentia back in the days, as I said, transformer didn't exist.

13:20And when the transformer paper came out, We basically mapped it to our generalized set of kernels. We saw that it's fully supported by the building blocks that we designed. And it was actually more performant than other workloads that we designed before and that we designed the building blocks based on. So that was a happy moment during the growth of this training architecture that we've been working on. How should we think about what is a kernel in this context? Is it a set of functionality for common operations like multiply, add, accumulate, that kind of thing? Or is it a higher level abstraction?

13:58That's a good question. I actually tend to intermix three terms and we might want to distinguish between them a little bit. One term is what we call an instruction. Then there's another term that is an operator. And the third term is a kernel. And so as a hardware architect, we care a lot about, or I care a lot about the instruction set architecture. These are basically the building blocks of the hardware, or you can think about it as the software hardware contract. This could be a matrix multiplication or a nonlinear activation function, or maybe addition of two tensors to one another. So that's the language that the hardware operates with.

14:51And any software that operates the hardware needs to map higher level workloads to that instruction set, that language. Then an operator is a framework level computation that is agnostic of the hardware underneath it. So that could be a layer normalization operator, a fully connected operator, which is a matrix multiplication followed by a nonlinear activation layer. And there's a bunch of other operators. And a kernel tends to be the mapping between an operator and the instruction set. So a kernel basically is a function that a software programmer wrote that implements an operator or a set of operators and maps them to the instruction set of the hardware.

15:44So maybe before we go further, talk a little bit about the Tranium architecture as a whole relative to GPUs and where are the similarities and distinctions? Yeah, definitely. So Tranium is a different architecture from a GPU, and I'll touch on the differences in a second, but we also see some form of convergence between the two architectures. So it would be interesting to compare and contrast there. So to begin with, a GPU architecture is a massively parallel architecture where we have thousands of different cores that are working together to run the same program. The programming model is typically called SIMD, a single instruction, multiple data.

16:37Basically, we have a single instruction set that is running on multiple cores, and each core is processing a different portion of the data. The training architecture is built around a small number of very powerful cores and very power-efficient cores. We built them around what we call a systolic array, which is a hardware architecture that optimizes for data reuse and energy efficiency, such that we read the data into one set of compute elements. They perform a computation and move the data to the next set of compute elements. so it minimizes the amount of memory touches that we need to do when running a workload on Tranium and that in turn improves the energy efficiency quite dramatically.

17:35Now having said that, we do see that the two architectures are somewhat converging, at least partially because Tranium is now introducing multiple cores, a multiple powerful cores but it but we are growing from two cores to four cores to eight cores and even more in the future and the gpu landscape starting from a the volta 100 generation by nvidia started introducing what they call tensor cores which are systolic array based compute units that try to take advantage of data reuse in order to improve energy efficiency so even though there There are a lot of differences between the architectures.

18:23There's also some similarities and they're increasing over time. Talk about how the user needs to think about these architectural differences. Does the user need to be very aware of it or are the differences managed by whatever frameworks they're using? So that's a nuanced answer to that question. So it depends on the type of user. We have lots of users that want their models to run, be decently optimized, but they're not necessarily trying to squeeze out every last percentage of hardware performance. So for these users, they actually don't need to be aware of the architectural differences. They can implement their model on PyTorch or JAX or any framework of choice.

19:16and then we have what we call the neuron compiler that maps between the framework level implementation and the instruction set of the hardware and optimizes the execution as it lowers the abstraction level. For example, if you go to Hugging Face today and you look at the, you have thousands of models there that natively map to Tranium and they just run efficiently. Specifically with Hugging Face, they introduce what they call the Neuron Optimum library that basically connects between the Transformer library and Neuron. So you just get automatic optimizations under the hood. You don't need to do much.

19:58That's one. And Neuron is the SDK for Tranium. Oh, yeah, I should have said so. Thanks, yes. And then there's another type of customers. Typically, the folks that are training very large foundation models. and for them, every last percentage of performance matters. If you can improve performance by three or 4%, it makes a world of a difference. I think that was one of the interesting things that we saw out of the DeepSeek team is that they talked about going underneath CUDA and writing assembly code to optimize their models or something along those lines. Like it's for teams operating at a certain degree of scale, it becomes very important to be able to do that.

20:43Exactly right. And that's exactly the type of examples that I'm referring to. And for these customers, we recently introduced the Neuron Kernel Interface, NKI, or we like to call it NICI. And NICI delivers basically instruction set architecture interface to any programmer. So basically you can come as an expert customer and say, I don't want to be reliant on the entire software stack, but rather I want to get direct access to the bare metal hardware along with its instruction set. And I want to program my operators or kernels and squeeze the maximum performance out of them. And we've had quite a few examples of customers, including Antropic and Poolside, and a few internal customers in AWS as well, that use Nikkei in order to really squeeze performance from the devices.

21:43How does CUDA compare to what you offer with Nikkei for Tranium 2? So first of all, I think NVIDIA did a fantastic job with CUDA. So props to them, and they built a very large user base on top of CUDA, which I think was a significant and important achievement in the industry. and at the same time cuda is somewhat hard to use because you need to map an algorithm which you think of as an atomic computation and you need to break it into hundreds or sometimes even thousands of cores and handle all the synchronization that is needed between the cores and that's one of the reasons that folks in the industry try to build abstraction levels on top of CUDA, including the Triton language by OpenAI.

22:40What we are doing with Niki is to provide with the same level of bare metal access, but in a way that we think is much easier to program. With the Niki AISA, I didn't mention, but Niki actually delivers two namespaces or abstraction levels to the programmer. The first one is called Niki AISA. You can literally import Niki.ISA and get that namespace. And the second namespace is Niki.lang. And we basically are trying to give customers the ability to choose the abstraction level in which they want to program the devices. With the Niki AISA abstraction level, the name is pretty telling. You get full control over the instruction set.

23:34So you get to program at the instruction level, and we believe that it's easier to program at the NICI-ISA level than in the CUDA level because you know exactly what the hardware is doing, and you also are interacting with a small number of large cores, so you have less of a communication overhead that you need to manage. with the nikki lang abstraction level we basically raise the abstraction a bit higher and we provide numpy like semantics which scientists are used to and and enjoy in their day-to-day life and that's that's an abstraction level that is useful for quick experimentation so think about trying to invent a new operator you probably don't want to squeeze every bit of performance in the first go.

24:24You just want it to run relatively quickly, see if the operator does what you want it to do. And if it's important, you'll optimize it later on. So that's what the Niki-Lang abstraction level is for. So if I kind of try to summarize a bit, I think CUDA is a phenomenal programming environment. We're trying to provide with more capabilities and easier interface with the Neuron Kernel Interface or NICI language that we provide to customers. And we're actually collaborating with quite a few universities that are now teaching the Neuron Kernel Interface as part of their high-performance programming courses.

25:03And we're seeing really good feedback from students so far. To what degree today in 2025 do you think that CUDA still represents this moat for NVIDIA that companies can't overcome? Or do you think that, you know, Nikkei or other technologies or, you know, the pace of innovation, like some set of factors has created an environment where CUDA no longer represents this unassailable moat for NVIDIA? So I think CUDA is a phenomenal programming environment and it's been very useful. and we as an industry built a large customer base around it. But I think the world has changed quite a bit in the last five to 10 years.

25:52And what we're seeing today is that the transformer architecture is extremely popular. And in order to build a transformer, all you need is five to 10 operators. We don't need to implement 100 or 200 different operators. And on top of that, folks that are training frontier models are spending lots of dollars on compute infrastructure. So if you take these two considerations together, it's absolutely worthwhile for leading AI labs to spend time in optimizing five to 10 operators to have their transformer running as efficiently as possible. and they'd be absolutely willing to do it in an alternative programming environment to CUDA, which I think is a good tailwind behind the neural kernel interface these days.

26:53So the folks that really care about performance and would go to the level of CUDA or even lower than CUDA, they just want something that's going to allow them to do what they need to do and gives them the most bang for the buck from a performance perspective. And the folks that aren't as performance sensitive, they can do whatever they want without going to that level just because of the framework compatibility. Exactly. So just to restate this in slightly different words, I think the folks that are training frontier models are going to chase the best compute efficiency, even if it requires them to do extra engineering work.

27:36And we're seeing that very clearly. and that's the driving force behind folks migrating to Tranium and also using the Neuron Kernel Interface and the folks that are not spending as much on compute will want to stay on the framework level they will not want to code neither in the CUDA level nor in the Neuron Kernel Interface level and that's why we built an XLA compiler that allows you to take a model from PyTorch or from JAX and seamlessly map it to training. Is the XLA compiler part of the neuron interface that you mentioned at Hugging Face? Is that under the covers? Or if I'm using a PyTorch model, will I need to run that XLA compiler myself?

28:24Both PyTorch and JAX have integrated plugins into the XLA compiler. The JAX framework was actually built on top of the XLA compiler. And the PyTorch framework introduced it over time with a shared project by Google, Meta, and us in AWS and a few others. What the XLA compiler does is it basically aims to be a just-in-time compiler, which means that it's pretty seamless to you. You run the model from the framework without any know-how that an XLA compiler is working under the hood. And then for the very first time that the framework encounters an uncompiled code, it will trigger the XLA compiler and map the framework level graph to the hardware.

29:18And every subsequent time that you call the same computation graph, it will map automatically to the previously done compilation. When does Train EM make sense? And for which types of users? Is it primarily targeting the users with these large workloads? Is it meant to ultimately be an alternative to the GPU for kind of your everyday generative AI workloads? How do you think about the user base in that way? So we generally build Trainium to be a general purpose AI acceleration platform, which can serve lots of needs across the AI ML community. But we do think along the lines that you mentioned before, and we have the following mental model.

30:07For the folks that are training frontier models, Trainium, and especially Trainium 2, is highly beneficial because it provides the highest performance and highest performance efficiency, whether you're denominated by energy or by cost. So folks that really need to build large compute clusters can benefit a ton from migrating to training. On top of that, we partner with a few providers, including Meta and HuggingFace and Ray and a few others, to provide with a list of already supported architectures, including transformers and the diffusion transformers and stable diffusions models. So folks that are using already existing model architectures can very well ramp on Trinium 2 with minimal efforts on their side.

31:03They run the model at the framework level and just trust the Neuron software stack to do the rest and map it to the device. And I would also say that the folks that might want to wait a little bit before migrating to Trinium are the folks that are inventing new models. but are not using large clusters. So for those folks, you might not want to take the burden of migrating to a new architecture for optimizing for a relatively low level of cost. So these folks might want to wait a little bit until the trinium infrastructure matures a bit and the migration becomes simpler and simpler over time. So in the near term, if I'm in a research environment, for example, experimenting with novel architectures, the burden to ensure that those architectures work efficiently on Tranium would fall on me since I'm not able to take advantage of a popular abstraction.

32:03And that might be someone who, you know, won't reap the full benefit. Even though there are caveats to that, that we might want to discuss. So Tranium does come with a couple of capabilities that just don't exist in other hardware out there. One of these capabilities, for example, is 4 to 16 sparsity. So an interesting sparsity optimization that allows you to scale, to speed up your compute by up to 4x. So if you want to try to make, take advantage of this capability and 4x speed up is something that seems quite compelling to you, you might want to try Tranium as well. But because you're trying to break new grounds, the way to take advantage of this capability would be through the neuron kernel interface or NICI programming environment that we described before.