How to Analyze Performance on AWS Trainium | Amazon Web Services

Key Takeaways

All right. Hey folks, Emily here. Uh, today you are going to learn how to analyze performance on AWS Traium. Let's dive in. All right. So, fundamentally performance is what it's all about. This is why we're here. This is why we have accelerators. This is what enables you to train models. This is what enables you to host models. This is what enables you to consume models. You probably wouldn't like working with a language model if it was too slow. You probably wouldn't like it if it was too expensive. You probably wouldn't like it if you couldn't fit all of the data into that language model that you wanted to fit it into. One thing makes all of this possible and that is performance. And so today uh we are here to learn about performance and how to think about analyzing performance on tranium. So fundamentally the reality is that there are just many ways to think about compute performance within the space of a IML. In the first video we learned about some of the hardware factors that drive performance. So those are the low-level specs about the accelerators, about the HPM bandwidth capacity, about how quickly you can move data, about your networking. There are hardware factors that drive performance and that serve as the upper bound. We're always trying to hit the maximum performance that we can achieve on our hardware. Additionally, there are theoretical frameworks that you can use to think about performance and analyze compute performance. This is called the roofline model. This the roofline model is not the only theoretical framework. There are a couple different out there. But we will use the roofline model regularly as a way to analyze what's going on in our workloads. Try and figure out how we can improve performance. A third way to think about compute performance is just practical goals. Like compute performance really is about just practically avoiding extra work. So if if your system is doing too much extra work, you want to minimize that. And you also want to identify spare cycles in your compute that you can just use to get extra stuff done. So that there are practical tips and techniques that you can use for compute performance. And then the fourth way to think about compute performance is endtoend model KPIs. So once we have your entire model constructed and brought up on tranium, we want to look at the endtoend KPIs or key performance indicators of that model, things like time to first token, things like tokens per second, etc. that you can use to understand how well your overall model is behaving on the hardware. So let's break these down step by step. Okay, so the first one, hardware factors that drive performance. So really there are four key hardware factors that you want to understand about the system itself. And remember these are the raw specs that are available from the hardware. You want to understand what those raw specs are because you're going to constantly drive towards achieving those in your applications. So tflops or teraflops per second, pedlops per second fundamentally these are floating point operations. It's a count of the overall floatingpoint operations that are available on your hardware at a given moment in time. In the context of language models and machine learning, you're going to care about TFlops when you are in a phase called context encoding. So when you send a document, for example, to your LLM and that the document runs into the LLM, that first phase is called context encoding. Context encoding is what we call computebound which we'll learn about in the next slide. And we care about Tflops in terms of how well we are able to do context encoding. We also care about T-flops in terms of training the model. So when we're training a model, we want to understand how many of those T-flops we are able to achieve during those phases. The second uh hardware factor that drives performance is memory bandwidth. So remember memory bandwidth is how quickly we can move data across those HBM banks and across ESPOF uh which is the onchip memory cache which is silicon accelerator itself. So memory bandwidth is a relevant hardware factor when we are in the workload phase called token generation. So you sent your document to the LLM. The LLM reads the context. It encodes the context and then it starts to generate tokens in order to formulate a response. Remember it's sampling from a probability distribution token by token by token. So token generation when we are in the phase of token generation when the LLM is actually responding to the document that you put into the LLM and maybe the question that you asked your prompt we need to consider the memory bandwidth of of the hardware because normally token generation is what we call memory bound. So the performance optimization goal in the case of token generation is trying to hit and consume more of that memory bandwidth. Normally that is the most common phase. The third hardware factor that drives performance is memory capacity. Remember memory capacity is the total size of those HBM banks or the total size of memory that we're able to store on the chip itself on the silicon itself on the accelerator and memory capacity is a relevant hardware factor when we are in a workload of model training with pipeline parallelism. So for example when the each layer of the neural network in your model is sharded across multiple accelerators which is pipeline parallelism memory capacity is a relevant hardware factor. Normally we're going to try to drive to improve memory capacity during that stage. The fourth hardware factor that drives performance is the interconnect or is the network capacity across nodes. And so that's how quickly we are able to move data between ultraservers in the case of project reineer or between uh tranium 2 instances in the case of say a a hyperpod distributed training context. So we care about that interconnect because it serves as a uh aspect of how quickly we can train our models. So those are some of the hardware factors that drive compute performance. And as we are becoming hardware engineers, as we're becoming performance engineers, we want to really study these. We want to really own the hardware. We want to understand what the maximum capabilities of the hardware are because everything in the rest of our model training process and our model hosting process is going to drive towards achieving those. All right, the next way to think about compute performance are theoretical frameworks. And here we call this the roofline model. Why do you think we call this the roofline model? Any guesses? Well, I hope you can see that it's because this actually looks like a roof. So that's why this is called the roof line model. Let's try to understand this graph a little bit. So on the x- axis down here we have something called arithmetic intensity. Arithmetic intensity is an indicator for how intense a given workload is. So the number of operations you are attempting to execute per bite of data. So as you are attempting to execute more operations such as by a dense uh language model or such as by a more complex language model or such as by increasing the batch size in your computational process that arithmetic intensity increases. As you increase arithmetic intensity, you're going to eventually cross over this ridge point. The ridge point is available from the hardware actually. So the hardware though remember those performance specs that we looked at in that first slide. So hardware comes with this ridge point at which a workload on the left hand side of that of that ridge point is what's called memory bound or data movement bound. And that means that the workload isn't really saturating the hosts. It's not really saturating the nodes. It's still just barely able to consume memory. And what we want to try to do is we want to try to move that workload to the right. We want to try to increase the arithmetic intensity of that workload to move it from that red stage which is why this whole area is red. We want to move that into green so we can convert it into something that's computebound. Why do we want to do this? We want to do this because when workloads are in that green computebound, they're just more efficient. It means you can scale it much more efficiently. It means you can add more host to it and the overall performance is going to increase. It means you can do much more at a much faster clip and at a much better price point when you're in the computebound phase. Uh and many workloads are strictly in the data movement phase. And so I should explain the y- axis here. So the y- axis here uh is the performance. And so that's operations per second uh that you're able to achieve. And so these are some important ways to think about your model and your workload. So when we are aspiring performance engineers uh and we are aspiring machine learning engineers and what we're trying to do is improve the performance of our model on train them and in Frenchia we can actually use the roof line model to study the operations in our model. So if we're doing a deep learning model, we're doing a ResNet, we're doing a language model, we can look at each of the components and each of the operations of that language model and consider where that specific operation sits on this graph. Actually, we can think about each operation if it's computebound or if that operation is memory bound. What we want to try to do is convert as many of those operations as we can into that green computebound range. So this is one theoretical framework that we can use to think about performance. All right. So the third way to think about compute performance are practical goals. Again just practically we're going to try to avoid extra work and we're going to try to use the spare cycles that we can. So when you are thinking about your model and again you're going through operation by operation you're studying the roof line performance of this you're looking at the compute behavior of that you're trying to understand how the data are moving you're trying to understand how you can better improve your compute utilization and specifically for each of the instruction and the set of instructions I want you to ask yourself can I reduce this or can I remove this work? Is this work necessary? What you'll find is that actually there are a good amount of bottlenecks that emerge because it's it's just not necessary. Like the precision of your model might be overly high. The layout of your tensors just might be resulting in extra transposes that you're doing in your matrices. Your memory movement might be redundant. You might have spills. You might have fills. You might be reloading data. And so for each of those symptoms, there are common optimization approaches that you can use to reduce that extra work. You might do tiling. You might fuse the operations. You might try to more efficiently select better instructions. You might try to lower precision etc. So that whole category of optimization techniques. The way to think about it is can I stop doing extra work basically. So the second way to think about this is underutilization. Underutilization means your hardware is sitting there. You have all of these compute engines. You have all of these flops. You have all of these HBM banks that are just sitting there. And in this type of bottleneck, the problem is you're actually not using those. Your workload is doing some process. What you want to do is find a way to use and better utilize the spare compute that's just sitting there. So some example symptoms of this are things like underutilized or idle compute cores or compute engines. Your scaling might not be great across the cores. Uh you might have lots of semaphors or syncs or barriers where your engines are not in synchronization and so that's blocking your ability to compute efficiently. you might have really low memory etc. And so for each of these symptoms there are common optimization approaches like better scheduling, better sharding strategies, better allocation on the chip itself, again tiling, uh higher batch sizes, etc. So those are just some common ways to think about reducing the extra work and better utilizing the compute. So finding spare cycles and using those spare cycles. All right. And then again the fourth way to think about compute performance is the endtoend model KPIs. So you've gone through your model, you looked at operation by operation. You know where that is on the roofline model. You know how much of the hardware capabilities you're able to achieve. You have stopped doing extra work and you've used the rest of your compute utilization. And now you're going to put all of those pieces together to look at one cohesive model and study the endtoend model key performance indicators or KPIs. So generally there are four families of KPIs for endto-end model performance you want to consider. The first one is latency. And so latency really is about time. And latency normally is about time where you are waiting for something. So the latency is like the latent time between you and the thing that you want to happen. And so some of the KPIs in in latency space are things like time to first token or TTFT. So TTFT is the time from when you send data to your model and you're waiting for a response uh normally in milliseconds. So that's time to first token. Per token latency is when your model is already in the generation the token generation stage and you're measuring essentially how efficiently you can generate each token. So PTL per token latency and then the endtoend model latency is the full round trip. So the time from doing the context encoding uh to the token generation itself. And so you have the endto-end model latency as a KPI. Another family of KPIs for N10 model performance is throughput. And then again throughput is essentially is an indicator for how much stuff you can get done at a single point in time. And so so the pure throughput or the raw uh capability to just maximally get stuff done uh with the compute. So some of the KPIs in the throughput space are output tokens per second. So after you've sent the content to your model and now it's generating the tokens for you uh and you want to measure essentially how many output tokens you can receive per second. Uh other KPIs are uh TPS or tokens per second uh during the context encoding phase. So you can measure the throughput of context encoding. You can measure the throughput or tokens per second again for that token generation phase. So those are some of the throughput KPIs. Memory utilization tends to be a KPI that is relevant for model inference. Uh so particularly the token generation phase of language models tends to be highly uh memory bandwidth uh focused where we want to better utilize uh the memory bandwidth. And then in compute utilization a normal KPI for the endto-end model is what's called model flops utilization or MFU. When we're training your overall language model, we want to understand how well that's consuming the compute. Normally, we'll use the MFU KPI for this. So, those are the four ways to think about compute performance. All right. So, fortunately for you, we have many tools that you can use on Tranium to make this really easy. The neuron profiler is an excellent tool that you can use to extract lots of really raw and rich metrics. When you compile your model, actually when you use nxdt or nxdi to compile your model, even just pytorch xla, you create this nef, right? This neuron executable file format. The nef is incredibly rich with data. You have so much information in that nef about the expected runtime, the expected performance, the instruction set, the data movement, the collectives. You have so much information in the nef. you can use the neuron profiler tool to generate a new file from that and that's called NTFF. And so NTFF is a new file format that is matched with the nef. So the nef is the output of the neuron compiler and then the NTFF is the output of the neuron profiler. So you take the nef, you throw it into the neuron profiler and you get this NTFF. The NTFF and the NEF are both going to go back into the profiler and give you a beautiful UI. So you have a really nice UI and application that you can use to study uh the clock cycles to study the instruction set that's being utilized to pull all of the really really rich summary and metrics data about the capabilities of your workload on Tranium at that moment in time. And so we use the neuron profiler every day uh within anaperna to study the performance of models um and see how well they're behaving and find ways to improve their performance. So that's the neuron profiler. And then in terms of the endtoend model kpis neuronx distributed offers benchmarking capabilities that are out of the box. So when you're establishing support for your model with NXDT or NXDI, you can add flags to produce benchmarking statistics that will tell you your endto-end model KPIs, context encoding, and token generation KPIs as well. So with that, let's take a look at the demo. We can use the neuron profiler to understand the performance characteristic of a machine learning model in a very straightforward way. Uh there are many ways uh we can use to generate the profiles uh using neuron profiler. Uh the easiest way is by specifying the neuron runtime inspect flag as enabled and specify a directory for profiler artifacts when launching the workload. Uh that is exactly what we are doing for this workload. And uh once my job is done, we would be able to see the output directory that gets created with all the artifacts required by the neuron profiler. And now uh we can use the neuron profilers view command to view the profile uh which will render this profile in a browser. So let's have a look at it. We are looking at a profile for a transformer-based model for a single neuron core. Along the x-axis we have the total execution time while the y-axis highlights the utilization of five key compute engines. Uh the first engine is uh sync engine which is primarily responsible for DMA triggers. Uh then we have two tensor engines which handles the matrix multiplications. Uh then we have scalar and vector engines uh which manage uh general operations and then we have GPMD engine as well. You'll also notice we have two CC engines uh which are crucial for handling the collective communication operations. At the bottom uh we have the HFU graph uh which stands for hardware flops utilization. Hovering over this section provides insights into how efficiently the neuron core is being utilized. Of course, the higher the utilization, the better the performance. Uh this is one of the most critical indicators that helps us to understand our hardware utilization at any given time. If we want to inspect the profile at a granular level, we can simply zoom in uh just like we are doing here. Uh now we are viewing the specific section of the full profile and when needed we can just as easily zoom back out to return to the original view. Additionally uh we can also use uh annotations to locate the collective communication operations and analyze the performance based on that. As here uh we are locating two CC ops and then adding annotations for them. Uh we can also rename them however we want. uh as I'm calling them marker one and marker two just to make them uh these annotations more meaningful. Now once we have these two annotations in place uh we can now ask the profiler to focus only on the section between them. Uh now this will make it easier for us to analyze the parts of the profile uh between these two annotations in an easier manner. Here uh we are just uh looking at the HFU and other tens other uh engines activities between these two annotations. Another interesting feature of neuron profiler is the box selection feature which lets us focus on the performance of a specific engine. In this case the tensor engine. Uh if you go down we'll be able to see the metrics corresponding to the selected part of the profile for that selected engine. Also uh if you want to have more details uh we can switch to uh full view. Uh this expands the interface uh to show DMA activity as well uh with all the DMA engines uh conveniently listed on the left hand side. Now uh moving on to the tabular data which is available at the bottom half of the page. Uh we have the summary tab uh that offers a comprehensive view of a workload matrix. Uh the first section here is overall stats uh which gives us highle overview of the performance. Uh then we have data movement section uh which dives into DMA transfer details helping us uh to understand how data is flowing through the system. And then we have tensor engine details which gives us detailed look at flops utilization offering valuable insights into compute efficiencies. Each of these sections and others uh provides crucial information that supports in-depth performance profiling. With that, that's a wrap for our walkthrough of the neuron profiler. Thanks for watching. >> All right, Saddaf, thank you so much for the walkthrough of the neuron profiler. Great to see the capabilities. Great to see the performance. Just a couple calls to action here. So, first off, when you are trying to become a performance engineer, trying to improve the performance of your model, three things to know. Number one, know the hardware. Study the specs of the hardware. Study the raw capabilities that come from that hardware out of the box. Second thing to know is the model. Study the operations of that model. Know where they sit on the roof line. theoretical framework, study and find opportunities for fusion, for tiling, for better layout, etc. The third thing to select is the right goal. Uh it's surprising how challenging this can be to find sometimes, but at the beginning of your performance engineering journey, you want to identify a clear goal for which KPI to improve and by how much you want to try to improve that KPI. And again, aligning on that KPI is is actually the first step uh in becoming a performance engineer. And again, the neuron documentation has an incredibly rich set of resources for you to learn about this entire stack and all of the content. And the neuron profiler has many built-in statistics and summaries that you can use for this out of the box. Then again giant thank you to the team who makes this possible who develops the profiler who develops kernels who develops performance tooling on AWS tranium and neuron today. It would not be possible uh for us to be here without all of you. So giant thank you to the team and with that I will see you in the next video.