PyTorch Symmetric Memory: A New Programming Paradigm for Distributed AI - Ke Wen & Chien-Chin Huang

Key Takeaways

Hello everyone, welcome to our talk. I'm Kwan. I'm super happy to be here with my colleague Chinchin and Natalia to present you Pyos symmetric memory, a new way of doing distributed programming. We have plenty of time, so let's start with a couple polls. Do you think it's possible to program your system as a single gigantic GPU? Can you raise your hand if you believe that's possible? Can you raise your hand if you think that's not possible? Okay, second poll. This is how we write our distributed program today. everything on the host side. We call some touch ops and then we call the stru to do the computation. Do you believe one day we can move all this code into a device kernel and let it run inside a single device kernel and let it do communication again. Can you raise your hand if you believe that's possible? Okay, we have more people now. Can you raise your hand if you think that's not possible? Great. So, if you have the same imagination as we do, I hope you will find the next 25 minutes worth your time because we're going to bring you a one-stop shop for authoring your own multiGPU kernel. So, what is symmetric memory in PyTorch? We can go back to the two basic communication mechanisms between processes. One way is called message passing. It is similar to packing all the information into a box. Take the box to a post office and ask the post office to ship it out for you. And because of the round trip to the post office, you will want to pack as much information as possible into the box, right? And hence this limits your flexibility of communication. Share memory is another way where two processes just open up some of its own memory space to be visible to the other process and the communication the data assets can be made arbitrarily at a very final grain. So if it is that great how should PyTorch enable it? Next, next let's welcome Natalia to walk through it. >> Hey, so uh let me tell you a little bit more how PyTorch enables symmetric memory on the technical level. Did we have the slide in between? Sorry. So, uh there should have been the slide. That's we need to correctly allocate this uh memory uh on all the GPUs that uh can be used for interprocess communication. But unfortunately, correctly allocating it requires some less common knowledge. It uses APIs from different uh vendors that are not particularly well documented and it's really easy to go haywire trying to do this. In fact, even when we were doing this, we fixed the memory leak pretty late in the game and it was present for a long time because it's hard to know which APIs to call in which order. Uh then uh enabling a pure memory access is also kind of tedious. So if you are left needing to do it yourself every time you want to experiment with some distributed APIs, this would uh create quite a lot of developer friction for you. And uh finally, even if you opened up those memory regions for the process to communicate, it still doesn't uh mean that any process can barge in and write in any address space at any moment it wants. They need to agree between themselves when uh particular memory is safe to write to when it's safe to read from. Are all the reads and writes visible the way we expect them to be visible? Memory spaces in GPUs are a complicated topic and getting it correctly is hard. So pyro symmetric memory is uh solving all those uh problems for you. And uh the way it's solving those problems is you are partitioning the address space in each of your devices. You designate some areas of that address space as accessible to all your peers allocated as a symmetric memory. And from then on uh you can uh write your own uh communication kernels or implement whatever communication patterns that you want to implement. Uh the agenda of our today's talk is a simple example of uh how you can use existing uh kernels and existing pytor symmetric memory. Then uh we'll move on to how you can write your own kernels with symmetric memory and finally we'll touch upon going to larger uh scale than just intra node or in link domain. Simple example of using symmetric memory in PyTorch. you are allocating symmetric memory in a very similar way like you would allocate any other tensor in PyTorch. Then you call a randu on uh this tensor to make sure that all your peers know how to communicate with this uh region in your memory. And from then on you can uh use it in your kernel or in other communication uh patterns. The simple way to do it is for example to uh call the symmetric memory kernels that we already have implemented in PyTorch uh like oneshot or reduce. If you want to do uh something uh more complicated uh an example on the right shows you the very bare minimum of what you can use to write your own kernels. You are getting pointers to your uh peer devices that you can uh use for communication. Typically you would start this kernel uh with synchronizing and making sure that all the participants are ready to start and the memory is ready and you would do the same thing in the end because otherwise your kernels might get out of sync but in the middle it's just the regular CUDA code that you would otherwise be writing and uh here we are showing the results for the oneshot symmetric memory kernel uh mentioned on the previous page and in latency bound regimes where you are all reducing small messages. The uh latency of this direct memory access kernel is clearly much better than what uh nickel or the twoot implementation provides. So what capabilities does symmetric memory unlock? It allows you to implement arbitrary communication patterns. For example, if you want only a subset of your devices to participate in some collective communication, that's pretty hard to do today with nickel. But with symmetric memory, you just write your own kernel that uh calls only those uh peers that you want without creating a new process group and it will work because you as a developer control it. Uh it enables as K mentioned in the beginning uh communication computation fusion uh within a single kernel and apparently a lot of you guys are already doing it but we are just hoping to make this experimentation easier and more accessible to everyone. And finally, it enables low latency remote access. So in latency boundary regimes where you just want to access your pure memory as fast as possible, this is your ticket. And next uh K will talk about uh some custom uh communication patterns that are enabled by symmetric memory. >> Thank you Natalia. Um for arbitrary communication patterns, I'm going to give two examples today. One is called mixture of aspert and second is tile tensors communication. Mixture of asper is an architecture that shuffles token between the attention layer and the aspert layers. And the key part of it is that there is a router that runs on each GPU to determine which token goes to which aspert. And because of the decision is dynamically made, we are faced with a issue that when we look at the input buffer to the auto shuffle, the splits of the asper chunks are dynamic determined and the information lives on the GPU. So if we were to use a traditional nickel API to perform this all to all one issue will be that we'll have to sync the splick decision back to host and then wait for the host to issue the next GPU operation which is all to all and because of this host to device sync the GPU will be waiting for more instructions to do and hence a store in the GPU timeline as you can see here and that uh hurts the performance with symmetric memory everything can be inside the kernel as Natalia just mentioned on the left hand side is a device function that does the split exchange here I'm telling my peer hey how much data I want to send to you and also I'm telling my peer hey this is the offset in I send buffer that you should get the data from. And on the right hand side is another device function that my peer read those meta information and use that meta information to perform the actual data fetch from my symmetric memory buffer. So with this we avoid uh D2H sync. And what's the speed up? 4x if we want uh particularly inference workloads. If you look at the trace uh you can see that we totally enimilate the gap on the GPU side and every kernel no matter it's compute or communication they got packed up together and because of the uh no sync computation we can apply kuda graph which is really a key uh for how you run your de application fast. The second example is uh from tiling. Here is a big tensor. Every rank has this tensor just with different data. And I'm cutting the tensor into four quadraters. And I want to reduce each quadure to a different rank to perform some data processing there. So that I can paralyze the data processing on this big tensor as a whole. This is a very common pattern in um high order optimizers that does blockwise optimization because you want different rank to optimize different part of the parameter and with again with traditional APIs like nickel this is not possible because nicl cannot support memory that's nonontiguous and a tile is exactly a data layout that's non-ontiguous. But with symmetric memory we can easily control where the CUDA thread should go look for the data. We can tell hey this is the uh starting point of a contiguous row and that is the ending point of a contiguous row. So with that you can write your own communication kernel and that's what we do in PyTorch symmetric memory package. Uh we uh surface a opore multi-root tile reduce. So you just need to provide a bunch of tiles like this one, four tiles and then specify a bunch of roots one, two, three, four corresponding to each tile and then the um the op will end up as uh output tile at each of the rank. So these are just some examples of we doing things but we want you to be the painter of the picture. Next, my colleague Chining will tell you how you can write your own kernel easily. >> Uh, thank you, K. So, Natalia and K already discussed how powerful symmetric memory is and also show us some interesting communication kernel. So, you might want to write your own communication kernel. So, how can you do that? So, try Trident is a very popular Python DSL developed by OpenAI for writing GPU kernel or like a kernel for other back ends. So, can you use Triton with symmetric memory? Sure, you can. This slice shows an communication skeleton written in Triton and the input to this Triton kernel are symmetric memory buffers, signal pad and also local data. We first use Triton's load and store to load the data and put it into symmetric memory buffer which will be then accessible by the remote rec. Then we use again use load instruction to load from remote rankers buffer. Then you can do whatever you want and whatever computation you want for uh with this loader data. All of this can be done with native Triton primitive. Uh but as Natalia mentioned we need some synchronization. How do we know if the buffer remote buffer is ready or not? So this is the part that requires some low-level PTX implementation which will wrap inside an API that you can or use. So with this skeleton in mind, let's see how we can actually apply Triton and a symmetric memory to a real use case. So this is uh this co snippet extracted from a real model and you first perform all reduce and then add a bias to the result and finally perform RMS no a pretty simple function. So when we use this model for inference the input can be as small as 10 kilobyte. So this poses several issue. First, Nicole reduce is not designed for low latency use case and also because the input is so small the three kernel launch overhead can be an problem. So let's see how we can use Triton and a symmetric memory to solve both issue. So for the first low latency issue uh there's a well-known uh or reduce algorithm one or reduce the algorithm is pretty simple. It's just like each rank receive the data directly from all other rank and perform the reduction locally. So we actually implement the core logic of this algorithm on the right side. There are two loop. The outer loop just iterate through all the data block and the inner loop iterate through all the ranks whether it's local rank or remote rank and we calculate the uh data address of each block on each rank and again we use load instruction to load the data whether it's local data or remote data it's the same and we load this data and perform the computation which is the reduction and we finish the one shot or reduce this implementation. So we solved the first issue. How about the second issue? The colon launch overhead. Can we extend this implementation? That's what we are going to do. So on the left side, this is our final implementation. On the left side, we simply just extend our for loop to include the bias addition because bias addition and also reduction, they are all wise operation. So they can easily be put into the same for loop. On the right side we have RMS known implementation in Triton and we put all of them in the same kernel. So we only have one kernel launch. So finally we also implement like two channel or reduce and we experiment uh on a a H100 machines and we can see that for small input size our uh symmetric memory implementation actually outperforms twice faster than Niko communication and when we increase the input size the gap becomes smaller and eventually yes we need to switch back to Nikico's implementation But this showcase a very important use case of symmetric memory which allow you to write your own custom communication kernel to address some issue that the traditional uh communication library doesn't uh cover and we put all of these APIs and also fuse uh kernel at Kraken and if you are interesting there are many more interesting kernels at a Kraken kernel and you can stop by Kraken's poster this afternoon. Next, K is going to discuss how to scale out symmetric memory. Thank you, Chini. Yeah, we really don't want to stop there. So, we want your kernel to be even able to go across your SS memory. How do we do that? Let's look at how the traditional way work. uh it actually is a three-party game. First, the GPU will have to prepare the data and signal the CPU that the data is ready and then the CP the CPU as a helper would uh prepare metadata for the network transaction and ring the doorbell at the nick. Lastly, the nick will fetch data from GPU and send it out. So this three-step round chip has a lot of overhead and hurts the performance of the network uh assets. But we want to let GPUs instead of CPUs talk. How do we do that? Nvidia have this great library called IBGDA which essentially allows GPU to directly initialize network transaction with the nick and as and and thus it has much shorter initialization time also as we know CUDA block has many threads so that we can prepare all the messages together concurrently and have mass amount of threads uh pumping messages into the network pipe. This help improve the ramp up of the network bandwidth. And lastly, RDMA is the core mechanism that allows the nick to access the uh memory space of a remote node. Just now, Chinshin told you how we can program symmetric memory with Triton. So, how can we let Triton go across no as well? This is how Triton calling enrichment man plugin. You just write your regular Triton kernel with Triton.jit and inside the kernel you can call enwishman functionalities in a Pythonic way. Enishman.put put send some data to some destination and the Pythonic API even observe uh tensor semantic where you don't actually need to give number of bytes you can just give number of elements that's how you operate a tensor and for triton to recognize that this is the external plugin it's very simple with torch you just need to add this requires enishman decorator and triton will bring in the assembly when it compiles the code from uh both the kernel itself as well as the enrichment library on the whole side just regular business allocate the symmetric tensor render it and then call the triton kernel uh with the symmetric tensor this allows the kernel to go across node so what will happen in the future Natalia >> yeah So we are not going to stop here and we will work on uh both uh including new uh interesting operations in PyTorch itself and we hope that you as a community will use symmetric memory to uh realize some of your developer plans. We also want symmetric memory to be actively used in our DSLs such as Triton and Helion and uh people are already actively working on using symmetric memory in Helen. Again, some of the details can be seen in the Kraken poster today. And finally, we want to uh do better localization and reporting. Uh and uh unfortunately, I'm out of time here. So, the main takeaway I want you to have from uh this talk is that you shouldn't be afraid of actually writing your own communication kernel. Similarly to how Triton uh recently enabled a lot of people who are not expert CUDA programmers to write their CUDA kernels, you should be e able to write your communication kernels also without doing a lot of extra prep work. Thanks to our awesome colleagues who uh worked on this project. [Applause] Uh we're ready for questions if there's any. >> I think Taval >> what is your recommendation? How can we performance issues with this library? Uh I can explain >> oh there's a there's a mic coming. Sorry. >> Thank you. Hi. Uh what is your recommendation? How can we debug performance issues, communication issues with uh symmetric memory? When we use symmetric memory, I can explain uh when we use PyTorch, we can use for instance PyTorch profiler and see clear nickel calls and see how how long it would take with this library. as I can see on the uh your slides. Uh there is a kernel that combines computation and communication and it takes some time and it also as you said heavily depends on the hardware on the nick on the IBG uh on GDA in general on GPU director sync and the following. So how uh if it takes for some reason if it takes a long time how we can debug this issue and how to uh who we should complain to about the performance of symmetric memory on different hardware. >> You want to take this >> K. Do you want me to take it or do you want to answer? >> Yeah. Yeah. So the uh first thing with all communication kernels is uh to determine if it's a really kernel issue or if you are just facing some uh struggler in your workload and you can use PyTorch profiler to help you do that. After you've determined that yes actually there is no struggler it's a kernel taking a long time. If you are using uh Triton, Triton provides profiling tools that allow you to mark the regions of the kernel and they would tell you for example in their profiling tools how long the computation part of the kernel uh takes place and how long the communication takes place. And if you've localized your failure to uh communication that's not getting as much bandwidth as you would expect uh then at this point since we are relying on env I I would talk to env people about what's going >> and I think that would be a great question for Nvidia too. How can you look inside a kernel to see what the instructions is doing? I mean maybe inside compute have some kind of functionality like that. Uh, second question. Sorry. >> Oh, >> I think uh the mic is at the back. Sorry. >> So, I'm new to NVISM. Um, I've seen the uh put and get constructs in your slides. Are they blocking constructs? >> Uh, that's a great question. Uh there are both blocking version which is just put or get and there's a non-blocking version which is put MBI and get MBI. >> Okay. So in the non-blocking version how do you manage the synchronization? >> Uh so you must pair it with a enrichment.quire call to make sure the transaction finish. >> So would they be similar to acquire lock sort of a thing in the um glue on TLX level or the synchronization primitives even very higher level? So, so what is the level of synchronization primitives uh in the nvishwim? I think that that's the question. >> Um not familiar with gluon so sorry. >> Okay. >> Yeah. >> Last last one. >> Oh um hi Sean McGovern from Red Hat. I just had a question in the off node context. Um so if we're thinking of this partition global address space is kind of a flat thing. Uh how do we think about data locality and what uh sort of challenges might there be? Uh another great question. So I hope you're only exchanging data you need to exchange like yeah there are a lot of like way of preserving data locality right so for example pipelining you you put the entire model way on one single GPU that's one way of preserving data locality and uh so yeah it's actually the programmer's own wisdom rather than a burden of this programming model. >> Thank you.