Sponsored Session: Accelerating GenAI Inference: From AWS Deep Learning... - P. Nguyen & A. Zhao

Key Takeaways

Thank you for joining us for our talk from AWS deep learning containers to scaling Amazon Rufus on trainium. My name is Fi Nuen and um I'm joined by Adam Zhao uh and uh today we'll talk about uh open source uh AWS contribution to the open source community uh and Adam will talk about secrets of Amazon Rufus highly scalable and a performant LLM inference cluster. So this is yet another chart uh depicting how uh openw weight model uh closing the gap with frontier model. Uh and this is showing exactly that it takes about 7 months before you can uh expect an openweight model running on a single uh consumer GPU uh matching today's uh you know frontier model which is you know pretty pretty astounding um if you think about it. So at AWS uh and Amazon we've done a number of uh contributions to the open source community across many different level of the stack. Most notably in 2021 we started our partnership with Meta and PyTorch on several initiative such as uh torch serve torch elastic and the Amazon S3 plug-in. uh and you know today very excited to also support uh the VLM project uh by providing some cloud credits for the CI and uh also uh a number of contributions across Amazon and AWS not only on GPU performance but also um you know leading the neuon support uh in VLM um and you also we've made VLM part of the deep learning container uh so that it's easy for anyone on to get started using VLM on AWS. So in order to support this open-source innovation engine, AWS provides a modular and a flexible stack composed of building blocks across compute, storage and network with open interfaces. So this allow you to build frameworks for the constantly evolving LLM uh landscape uh for different accelerators and optionality uh in mind. So it might be worthwhile uh to pause and take a look at the current pipeline shift we're going through and see some of these implications at the infrastructure level. So I heard um across several talks and you know including the keynote that we are shifting from pre-training to inference and post-raining and within inference itself we are also seeing a shift between input heavy token to output heavy token right at the beginning of the LLM era um you know rag and stuffing the context window uh for summizations and and translation and then evolving to uh chatbot multi-turn conversation to today's uh reasoning models pairing those agents and a deep research type of use case right um so uh shifting from input heavy to output heavy uh and those reasoning models changing uh the pattern of compute and memory uh consumption right with a bigger emphasis on uh memory uh consumption And so you know in order to work around those memory limitations there has been many uh different innovations um across several frameworks uh and across several use cases right page attention from VLM uh to um you know how to train a model using multi-query attention uh MQA and GQA um and and maybe quantizations and you know all the different hardware uh supporting different numeric and data type to uh different strategies on how to offload your KB cache across different memory tier or or storage tier and finally disagregated architecture. So as you may know in a transformer model there are two phases or two stages. There is the prefield which is known to be compute bound uh and then the decoding which is known to be memory bound. And so you know initially when you have those two different stages on the same host uh when you uh run the decoding phase uh very often your GPU sits idle right and so there was a a a need to decouple those two different stages and so now you see uh VLM and SG lang supporting uh this type of architecture also known as disagregated architecture. So by allowing to allocate um you know resources for prefield and allocating more flops or high compute GPUs uh and then allocating more memory optimized for decoding uh it you know allows you to really scale those uh stages independently um and and more efficiently and also you have um you know you need another protocol to transfer the KB cache across those different clusters. Another trend is the rise of mixture of expert or moes. Inherently provide two to 4x better cost performance. So when you train a model you only need to train uh two out of the eight experts or update the weights of two out of eight expert. At the same time when you run inference you only need to activate 25% of the weight right. Uh there's however a trade-off when you deploy those MOS model there's a bigger memory footprint um and uh you know a more complex deployment in nature um and you know you also may be familiar with the 3D of parallelism right data tensor and pipeline parallelism and they all share similar characteristics in terms of communication patterns right they all have symmetric uh and data volume exchanges between those different nodes um And u Mes breaks those assumption right and so when you do expert parallelism you need to shard those different expert across different nodes and each expert needs to communicate to another expert. So this is also known as all to all communications. You can also expect a three to sixx more communication between uh those uh those experts compared to dense model. Um and a very hard to predict the volume between those different expert and so therefore um there is an asymmetric uh type of uh volume and communication in nature. So, Amazon and as provide uh the elastic fabric adapter to enable those advanced uh networking patterns uh also known as EFA and so EFA allow you to bypass the OS and allow you to do RDMA or direct GPU toGPU communication right uh and EFA uh relies on lip fabric lip fabric is an open-source library to allow you to offer those advanced communication pattern uh lip fabric is also integrated with other standards and different protocol such as open MPI or Nvidia Niko and Vishm or Nixon. So um it's not just on the networking layer but you can also use EFA on the storage layer. So you can use EFA to read from storage directly to GPU and um you know allow for better throughput uh and uh to saturate your GPU more quickly. So we touch on networking storage and finally compute. So VLM great platform to run over 100 model architecture across different accelerators. Right? We touch on how the prefield and decoding are becoming separate and require different hardware characteristics and so now you can allocate um you know GPU trainia AMD and other accelerators um you know to better uh serve those um task specifically and get a better utilization overall of your uh resources. And so VLM allows you to really uh benefit from uh the variety of those uh compute and allow you to you know build those uh use cases with optionality in mind as well. Um and um we've made the VLM framework part of the deep learning image and container so that it's very easy for anyone uh to run or build on top of VLM um on top of ads. Right? So we've encapsulated all the complexity of libraries, dependencies, optimize the runtime uh and different drivers so that you know you can deploy and uh be confident that it's going to run um you know VLM uh with the best experience. So you know whether you are innovating at the AI gateway or router uh using different framework and maybe needing to integrate with different inference engine such as VLM or SGLAN maybe you need to offer those advanced communication pattern using uh Nvidia Nixo or Nickel or maybe you want to you know do KV um offloading uh using LM cache and moonake uh or maybe doing fine grain GPU orchestration using Kai or Ray uh And finally, maybe you want to offer those custom kernels um using a variety of building blocks on AWS. And so as provides all those different building blocks uh with an open uh interfaces um so that you can evolve as uh those use cases are evolving. Okay. So Adam is now going to be talking about Amazon Rufus the shopping assistant and how uh they were able to to scale to 80,000 tranium chips uh during Prime Day. >> Uh thank you V. Uh so my name is Adam and I'm going to talk about uh some of the experience that we have uh within Amazon uh Amazon Rufus and how we built uh such application to run at Amazon scale. Um so what is Rufus? It's really a a conversational shopping assistant that developed within Amazon uh that can deal with a bunch of like shopping related tasks. So you can think about like as simple as the tasks like asking a factual question specific to a product or ask this assistant to give you some recommendations on some circumstances that you come up with like you need help from someone that you can buy something to fulfill the task. Um and what we want to do is really to build a AI powered um feature available on Amazon through both mobile app and also Amazon website that can help customers making informed decisions when they are shopping on Amazon. um and deals with a wider range of questions that powered by advanced generative AI technologies and pretty much uh what we have been doing here without guy people sitting down there like developing all these fancy algorithms and also model innovations uh we collectively build this innovative approach uh to conversational shopping experience that is non-existing before um so at early stage when we started building um a inference stack specifically or the foundation infrastructure to power the Amazon Rufus uh feature uh we faced a couple challenges. So there are major four challenges I want to list here and we will go through like why they are challenging in what sense. Um so first like you need to understand that we are building application that's running at Amazon scale and also it's need to be run reliably and not something that you can deal with today but it crashes tomorrow and uh the average uh compute that we need to handle is about three million tokens per minute. Um and you need to serve millions of customers around the globe at that scale with high availability and the reliability and second um you want throughput that's to serve like multiple like hundreds of customers but you also want speed. So latency is also important to us because people are very spoiled uh with the CPU based load nowadays that they are very used to responsiveness on like these fast applications. But you understand generative AI usually like it's a bit slow when you comput compared to CPU based workload uh and you need to like really invest heavily to make them running fast. And we want to still honor that the experience for all our customers where uh we want to put that constraint under one second so that we can stream back responses to customers as soon as we can and they are getting instantaneous like a responsive uh uh responsiveness from the application as possible. Uh and third uh we want to make sure this business is running with a sustainable cost and that helps us to like ensure that we can continue delivering new features to our customers and make the business sustainable. Um and uh we how do we achieve this is we want to build a scalable inference cluster with low overall cost at Amazon scale. And fourth I think around nowadays the trend is people are still building larger and larger model although that activated parameter size is reducing but you still see like 100 bit h 100red billions of parameter size models being developed and you need to make sure that these models can be loaded into the memory before you need you can serve them. Um and uh at one point you may no longer have an instance that has the total memory that's fitting a single model that we are building here. So we don't want to be bottlenecked by that and that's very important that how can we remove this constraint on model size. So the solution we landed um is really like a combination of the uh the the levers that we have and the tools that we have that u within Amazon and also from open source community. Um so we landed a solution to build the multi-node inference cluster uh with a training chips pietorch and also VLM. So there are two very important pillars uh uh like related to requirements that we thought about uh when we come up with this solution. So on one uh on one side uh we do have requirements around model performance. So we want to achieve the best uh the best speed and also the best throughput. So we want to maximize the compute and also memory utilization to achieve that. On the second side, um we also want to build a very robust infrastructure and also that's easy to uh transfer to like be as flexible as possible to be transferred to like dynamic uh infrastructure options. So what we need is a containerized abstraction of this multi-node topology that can do um like communications very efficiently and also maintain the consistency across nodes and we need this infrastructure to be flexible and scalable and that's means that we can pour this solution into any infrastructure that we are evolving into and be um like keep evolving for the uh next adventures that we have. um collaboratively. So let's look into a little bit more into the details on the uh architecture that we have. Um so this is a oversimplified version. Um so you can see there are two types of nodes. One is the leader node, one is the follower node. Um and theoretically ideally you really can scale this into like tens of dozens of nodes that doesn't need have to be just two nodes. Um but this this is just for simplicity here. So you can see that leader node and the follower node they are almost identical. Um leader node is just taking a little bit more responsibility here. So what the leader node does in addition is to do the request orchestration. So the leader node leverages VM for efficient dynamic traffic scheduling for follower nodes to execute. And what it does is to leverage like something like continuous batching available from VLM to broadcast the inputs and assign these requests to all these followers that we have and so that these for nodes can start executing and for sure the other nodes also have compute because it's also a GPU device so it can broadcast to itself as well and when the follower nodes receive those inputs then they will start execution. Um so during the input broadcasting we leverage PyTorch which is the torch distributed um library uh for the uh input uh data uh broadcasting and during the actual request processing uh there will be a lot of uh collective communications happening across nodes. uh and we leverage the neuron uh NXDI library which is the neuron distributed inference library and also the elastic fabric adapter which is is EFA uh to leverage the high bandwidth and the low latency internal communication. So which ensures that uh we can move data around efficiently and also we can minimize the networking um overhead and in order to achieve the best performance on a certain model and also with the certain infrastructure you have uh you tend to always like need to do a little bit more finetuning. So what we do in addition to like think about what architecture should be we also invested a little into the model parism strategies. So on one side we want to maximize the memory bandwidth uh which is MBU. On the other side we want to also maximize the flops utilization which is MF MFU. So one stands for the computer, one stands for the memory and that's exactly what we're talking about here is like throughput and latency. Um, and you also need to have a good balance uh because you're dealing with so many nodes at the same time. You don't want to overload a single node and that's that becomes a bottleneck to slow down everything. So you want to minimize the memory pressure on individual nodes uh when you have such architecture. Um and we use a hybrid uh par strategy meaning it's hybrid in the sense of prefill stage and decode stage. They use different uh strategies here. At the prefill stage, we leverage context parism uh which basically split the input of the same request into smaller sequences along the context and then that we can do the parallel computation for the prefill and also attention like uh computations across different nodes. And at the decode stage we use data parism so that we can shard or partition the data or requests along the batch dimension and we can deal with a subset of the request across different nodes. So there also help to balance the the workload across the nodes. Um and all these we uh what we just talked about uh we have a very well containerized solution built around it. So it's a nice abstraction that we can port it into a like certain infra infrastructure that's not necessarily like built by ourselves but it can also be leveraged by other customers if they want uh with native PyTorch VM and the neuro SDK support. Uh and also it's generalizable to other hardware types. So one important I think I want to echo is like if you talk about EFA support with GPUs but EFA also supports a variety of hardware types like trainium and infranchise is one example that we built uh with these solutions um within Amazon uh and also we're able to use this solution and scale for the Amazon traffic um with EC2 and ECS these are very scalable and robust services available for the compute and the deployment um and uh we also have these levers that we can leverage to minimize the cross node latency overhead. Um so on one side we do have the flexibility to optimize the node placement. So you you you you can think about like when you're dealing with data center you want to have the affinity of the nodes located physically closer so that they are physically uh lowering the network uh overhead and on the on the other side other hand we also leverage EFA again uh to achieve this low latency by leveraging the high bandwidth and uh for the cross node communication. Um and another thing you very important to ensure the robustness and the reliability availability for a system is you need to have a good telemetry and also monitoring system. Um and AWS has the cloud watch service that we also leverage uh to build that. So it gives us like alarms when there's errors happening and give us good telemetry and visibility into the service. And lastly, we leverage trrenium uh trrenium chips uh because it has a very uh good like advantage in terms of price performance. It's 30 to 40% better um than the comparable GPUs. That's also a public available data uh data that you can uh search from the Amazon official u posts. Uh and there's a link over there which is um uh an ad as blog post we posted earlier this year has a little bit more details if you want to learn uh and feel free to go there and check that out. Yeah. So um just a quick recap um what we mentioned like four challenges we um faced when we started building the inference infrastructure and stack for rufus application in Amazon. So we had a challenge in latency uh which we call it user perceived latency because um you need to send something meaningful visible to customers uh when they feel that you are responsive. Um and second um you need to make sure that uh your infrastructure is a reliable and scalable that can support uh customers as the business grow and also for certain volume traffic uh per second or per uh on a daily basis. And third, you want to have the low cost uh because you want to make sure the business is sustainable and also you want to make sure that you can continue delivering new features to customers um as long as you are running business efficiently. And lastly, you don't want your uh infrastructure to be a bottleneck on the model innovation. So you want to unblock yourself from deploying any size of the model and then you're really doing like distributing inference with multi node and uh the model can be sharded across multiple nodes and they're running collaboratively. Uh and lastly uh I think um is really want to like innovate together with all you guys and also open source community and uh to continue build and innovate on AWS. Um and we believe uh the next breakthrough in AI will come from someone sitting in this room or in this conference. And uh uh FE is here also like what we I think one of the mission that we want to uh fulfill is which is ongoing mission is to give you the building blocks uh for scalable and performant um AI solutions that we can offer uh from Amazon and make it happen. Um and there's I think uh on November 3rd there is another deep dive on meetups with Ma and Nvidia at at the ads loft which you can check more details by clicking the link below. Uh and welcome to find us uh at that booth on November 3rd. Yeah, I think that basically closes uh my presentation. Thank you. >> Thank you.