Lightning Talk: Distributed AI Without the Infrastructure Tax - Yahav Biran & Maen Suleiman

In this session, Yaav Biran, a customer engineer from the Annapurna ML team at Amazon, addresses the challenges of deploying machine learning models at scale. He focuses on the complexities of infrastructure, particularly when transitioning from research to production. The speaker discusses the issues surrounding build-time dependency conflicts and hardware complexities, as well as the necessity of scaling across multiple instances. Biran introduces deep learning containers (DLC) and the concept of dynamic resource allocation (DRA) within Kubernetes, highlighting how these tools can abstract away much of the underlying complexity. He emphasizes the importance of close proximity between neuron devices and network fabrics for efficient resource utilization. Overall, the talk provides insights into minimizing infrastructure overhead in AI deployments while enhancing operational efficiency.

How many people are deploying serving models um at scale? And I'm talking about scale, not not experiment, like hundreds of users. Okay. And um and uh how many times So, there is one, so I'm going to just going to talk to you, so how many times did you have to uh work through the process of deploying it from the point where the researcher um it's ready until

it runs uh efficiently. Exactly. Exactly. I mean, normally, I would uh I would like Maybe maybe in few years, many people will raise their hands, and they will be able to share more pain. But uh but but essentially um as an experienced model uh uh uh um building model serving uh internally, we faced uh uh issues, challenges, specifically around the that uh that point, and what I'm

going to share with you today how we are think we are trying to uh to solve it. Again, for those who have just joined, my name is Yaav Biran, and I'm a customer engineer in the Annapurna ML team over at Amazon. So, in the next 10 minutes, uh I will show you how to minimize the infrastructure tax uh this in AI deployments. And distributed AI is a

mouthful uh uh um statement, but the point is how you take a single serving, a single um unit single compute unit that serve the model, and expand it multiple nodes across network, etc., with and and this is specifically on the infrastructure, right? I'm focusing on the infrastructure, not so many on the platform, whether it's VLLM and other. Okay? So, the talks that I'm referring to starts at

build time dependency conflict. It continues in deploy time with hardware complexity, and then it compounds when you need to scale out to multiple instances. And and the reason that it's I'm talking about the scale out, I'm referring to where you want to take a tensor that sits whether it's By the way, everything I say here can can be valid also to GPUs. Uh but but I'm in

in in the neuron context, every bit that you have on the neuron core that you want to transfer to a remote neuron core, that will require some configuration. uh through today. The solution will be DLC, deep learning containers. It's a project that we built for build time. Neuron DRA, DRA stands for device resource allocation. This is a Kubernetes construct. It's a new construct Kubernetes community developed for

all accelerated compute, GPUs included, and and and then and other network components. And so and and then yeah, so we're going to talk about that in a second. Okay? So, these are This is the problem, to the solution, that's the that's the the problem. The first part is the dependency complexity. For those who tried to use neuron, but also to be fair, other with other accelerator on

on GPU, there is a lot of moving parts in terms packages. In our case, we have the Neuron SDK, the Neuron back back end, the PyTorch, and other libraries. They all need to work together. For those who participate in the previous session, we are trying to minimize that friction by going native. with the PyTorch community, but it's not always the case, and for that, we added this

this part, right? So, if you have a version version, sorry, mismatch, that will cause the unexpected issues during your build, and and so on. Second one is the hardware complexity. We are now at this today, we have Inferentia 2, we have Trainium 1, we have Trainium 2, we have Trainium 3. Same same story you can tell about GPUs with L4, L4TS, P. Uh so so there is

a lot of hardware complexity that is that is coming, and how do you know how to schedule the work the specific work to a specific hardware, especially when you when you run the infrastructure from the guy the guy that runs the infrastructure will not know every specific option. And then the third one is the network, right? So, when you want to expand whatever you have on a

single box, how do you expand it on on network, and this is not just a standard scale out exercise. It requires a little bit more details. So, the results of this problem is that ML practitioners like yourself are spending time understanding those topologies and instead of serving those models and uh benefit So, the the solution here is a is a three-layer. I'm I'm starting from here. We

have a project called the Neuron deep learning containers, right? That we are managing with that we are managing ourselves and our partners over at the uh and that are here today. And and and we add all the all the packages and basically we are creating a tested contract with with you. So, um so you don't have to start building everything from the from the beginning. In the

middle, if you is you we are adopting a like other vendors, including Nvidia, a dynamic resource allocation that abstracts away the hardware complexity. The Neuron DRA drivers enables the Kubernetes native scheduling where where are defining those templates and I'm going to show you an example of that template, right? And then basically you are as the as a as the practitioner, the scientist, the engineers are basically just

reference them and and use them. So, we are abstracting away those those uh semantics. And at the top, we have EFA. EFA stands for elastic fabric adapter, which is the Amazon way of saying InfiniBand, okay? For those who use other platform. And so, this InfiniBand requires also an extra configuration that makes you are uh checked. So, let's look at each This is just an example. This is

our DLC, right? And this is just an example of what we what we have. It's a It's a It's It's It's a lot It's It's a cognitive load that you need to take care of that. And and you don't need to, right? You don't need to. What you need to do basically is just pull the the choose the more the DLC that the deep learning containers from

a GitHub, use it, and uh extend it. We are also open-sourcing the Dockerfile, so you can you can customize it yourself. The key insight here is that you don't you don't need to debug compatibility and start deploying the models with the with the specific asset. On Now, I'm going I'm I'm double-clicking on the on the second problem, which is abstracting away the hardware complexity. This is what

you have before the array. So, for those who take the picture, this is not a good example. It just show you that how difficult it is. I'm just for for those who are the Kubernetes construct, you need to specify node selector, anti-affinity, and affinity, basically telling the the scheduler what to do and what not to do. Okay. And so, we're moving away from that. So, what we

want to do instead is you the ML ops will abstractions of each instance, right? So, for example, here we are saying, "I'm going to use four four four devices, right?" These four devices, I'm not going to tell I'm not ask the the the developer to know about all the details that I showed before. It allow you to specify the the attributes that we are publishing and you

as the scientist, all you need to do is just specify the template, right? And And why is that? It's It's important because that if I'm If I'm the ML ops, I'm going to de- This is the common pattern. I will define four types of claims, small, medium, large, extra large. So, you don't have to worry about all the weird things that I wrote here, okay? All you

need to do is just specify it. And this is the the DLC as an example. That's all. The The third thing, which is super important, if you want to have an EFA, right? If you want to do disaggregate inference, for example, right? And you want to move a tensor value from one machine to another machine, you need to make sure that the neuron device or device are

in with close proximity to the fabric that is going to run. This is also valid to other to other platforms, right? So, if you don't do it, it will find a a different device, but do a cross NUMA traffic and that will cost a that will cost in throughput. So, instead, what you have to do, same approach with the claim template, we're extending it and you basically

add another a another request here. So, and and what we are doing, what the EF what the DRA will do, the component that runs within Kubernetes, it will make sure neurons that you ask here will be in close proximity with the four EFA, basically on the same NUMA. Okay. And I don't think I have time for for for this thing, but that's basically the Kubernetes the Kubernetes

flow. Treats and the with LC 3D LC as the tested contract, right? So, don't have to don't be don't walk too much on the on the build complexity. Use the array for the right abstraction boundary with your ML ops. Large extra large t-shirts t-shirts that you have also the network DRA in place to to have this colocation correctly. Thank you very much.