PyTorch Symmetric Memory + NCCL Device APIs: A New Path Towards Multi-GP... Ke Wen & Sylvain Jeaugey

This talk focuses on the integration of PyTorch symmetric memory with Nico device APIs, presenting a collaborative approach to optimize programming for multi-GPU kernels. The speaker discusses the motivation behind symmetric memory, emphasizing the need to enhance performance in distributed workloads, particularly concerning small batch sizes in inference and the complexities introduced by Mixture of Experts (MoE) architectures. They elaborate on symmetric memory's ability to allow one GPU to access the memory of another, thereby facilitating efficient data transfer and reducing latency. The presentation introduces the NCCL device API, which aids in writing optimized multi-GPU compute kernels, and outlines various modules like LSA and GEN designed for efficient memory access and networking. Finally, the speaker highlights the advantages of collective operations enabled by symmetric memory, such as CE all gather, which can lead to significant performance improvements in PyTorch applications.

Hello everyone. Thanks for coming to our talk. My name is Kuan. I'm a developer of PyTorch distributed. Today I'm with my colleague Sylvain Okay. He's the lead engineer of Nico team. Our topic today is PyTorch symmetric memory plus Nico device APIs. I call it a second collaboration of these two libraries. And hopefully with today's talk, we can show you a new way of programming multi-GPU kernels. So,

why do we need symmetric memory? There are a couple challenge today in our distributed workloads. First, inference is becoming the GPU burner. And in the decoding phase of inference, the batch size is very small, so every collective need to be latency optimized. Second, MOE architectures create new communication patterns. metadata, including where a message should go, how many tokens I should send to a destination, they all exist

on GPU memory rather than CPU. So, we need to consume the metadata directly in GPU. And of course, we need to provide an infrastructure for compute and communication fusion, which is our ultimate performance goal. This is our agenda today. Hopefully, they consist of exciting topics here. Um but before we go into details, I want to thank the uh open source community that supports PyTorch symmetric memory. They

span multiple companies like Meta, Nvidia, AMD, and many others. Thank you all. So, what is symmetric memory? I want to give a simplified view here. Um remember, the goal is actually for one GPU to be able to access the the memory of another GPU. So, on this left-hand side, we have gray boxes, which represent physical memory of say GPU one. And we want to let GPU zero

to access it. The mechanism is that GPU one will expose some of its physical memory as a handle for GPU zero to import. And it is getting mapped into the virtual memory of GPU zero. When GPU zero writes data into that virtual memory space, it gets automatically mapped into the physical memory of GPU one. So, this is the unicast scenario. We can also do it in multicast

manner. In this case, this has one more GPU on the right-hand side. And both of them can expose their physical memory as a handle to GPU zero. And when GPU zero writes data into this Nico uh the NVLink switch will take care of broadcasting the same data to both physical memory of the two GPUs. So, that's the underlying mechanism. So, how can a PyTorch user use it?

We provided the torch.distributed.symmetric_memory module uh as a entry point. And this module contains the flow on the host side and the device side. So, let me start from here. First, you can create a symmetric tensor using symman.empty. That's very much like torch.empty. In order to tell PyTorch that this tensor is actually connected uh among a communication group, you can call the render with API given a communication

group argument. After that, you can start to call collective off that supports symmetric memory. For example, one shot all reduce. On the device side, of course, we need a CUDA kernel that supports the all reduce, for example. here is a very simple sketch of how you write a multi-GPU communication kernel. First, we need to sync with remote peers to make sure their data their input are both

ready. And then we fetch data from the peer. We do some local compute, for example, sum them up. And conditionally, we may want to write result back to our remote peer. So, we only have a sketch. How can you write multi-GPU kernels efficiently? Here is NCCL device API coming to help. Right. So, I'm going to present the NCCL device API. It's something that we added to NCCL

in 228, then more into 229, and now we getting ready to release 230. So, it's pretty recent. Um yeah, I'm going to Right. So, we need to define a new couple of new concepts. Um in NCCL, most people are familiar with the notion of communicator. It's the group of GPUs that compute together. Uh and so, on the host side, the first step is to create a device

communicator. Uh it's the same thing as a communicator, but just for device code. Uh it designates a group of ranks. Um and it's equivalent to the communicator just on device side. The other important concept is the window. Um the window is related to symmetric memory. Uh it's just a group of buffers across all the ranks. Uh and that's uh the other notion you need to access memory

uh in that world. Uh and then you pass those to your kernel, and your kernel will have a toolbox to write uh communication or to write multi-GPU compute kernels, if you want to write all the things that we discussed before. And right now we have a few different modules to write the those kernel um we have first the LSA module which means load store accessible memory. So

functions to access other GPUs memory directly. Uh you can't always access other GPUs memory but for those which are connected through PCI or NVLink you can. And so usually that's the right way to do it and actually the only way to do it efficiently. Uh so that for that part is really targeted to NVLink sorry not NVLink uh or PCI. Then we also have the multi-mem functions

to use NVLink sharp. Uh so just what Kurt described where you can uh have stores that broadcast to multiple GPUs or reduce from multiple And finally we have GEN which stands for GPU initiated networking and that's the network API basically. Um so all those are very low level um they're uh basically our our first break towards building uh an API where you can build more interesting things

on top um and they're very specialized. So the our goal was really like to make it as optimized as you could. Uh GEN really maps as directly as we can to the uh and LSA multi-mem I'm going to explain how they work uh but they also the the simplest thing you can think of uh to access other GPUs memory through NVLink. So starting with LSA um we

don't actually have a lot of APIs there. Uh it's extremely minimal. Uh it's basically two things get pointer and a and a barrier a synchronization barrier. So as Kurt explained you you use you may need barriers to uh synchronize wait for the others to have the data ready uh and then once you know the data is ready because they reached that point you can start accessing it

or storing to it. So you just get a pointer and then you do loads and stores. Uh and the reason we didn't really go further than that for now as a first step is that actually you do a lot of things with that. You can do loads stores, but you can also use TMA, you can use all the mechanisms that could gives you to access the other

side's memory. multi-mem approach is the same thing for In-memory Shark. It's basically just get pointer again. And then on top of that, you can then implement everything you want. The all the the the multi-mem libraries that we have that expose APIs to use it. You can use feed that pointer into it and you're going to be able to use In-memory Shark that way. So that those two

are not giving a lot of features, but it's just a way to integrate that nicely into Nickel without having to do your mappings yourself. I mean the the the In-memory Shark mappings are fairly that's doable. The multi-mem mappings they're much more painful to do. So yeah, this is mostly a convenience at this point. There we don't have a lot of new a lot of hard that that

replace big pieces of code, but that could come later. The last part though is one that probably you don't want to implement and that's Jin. So Jin is the network API to be able to initiate network operation from GPU memory to other GPU And we have more notions. The main mechanism is a Jin put. It's a put operation remote moving data to a remote memory and you

can actually add some signal to it so that it generates a signal on the remote side and then remote side can call wait signal to wait for the data to be received and once wait signal returns, you can actually process the data that you just received. We also have a barrier just like for ASA. And we also have actually counters if you want to wait for local

operations to be complete. So for now that's what it is. An interesting aspect is a Jin context. So actually when you create a gene object, um the last argument you pass is the context index. And that allows you to have different CUDA CTAs use different channels. Um and that also means you can actually use different NICs. If you have more than one NIC per GPU, naturally, the

the traffic is going to be spread on all the Um and also use different uh connections, which can also in some cases improve performance. So, gene context management is an important part. I won't go into details, but uh that's that's something you should look into if you want to write efficient So, all that is the current state of the API. It doesn't mean it's going to stay

like that. We we're going to build on top of that and uh add more and more nice APIs. Uh but for now, with the just that, you should be able to do a lot of things already. And you don't have to write anything in C C++. Uh you can actually go and write that uh with uh cute because we also worked on mappings to be able to

integrate that into cute. effectively, that means that if you write your code and you just add net some communication calls in the middle, you might be able to even fuse everything and get all your communication for free. So, that's that's the ultimate goal of communication computation fusion that Cur referred to. Uh and that's something that we feel think is going to take a huge space in the

future. Uh and of of course, because it's all JIT compiled, it's also all JIT optimized, which means that the code can actually be optimized for that's very specific case you're compiling, which can go very far in terms of optimization. And all those kernels here, uh we expect that they will start showing up everywhere. Uh they can start from the user cur- the user code. You can write

your own, integrating that in your code. That works. We We will also see some appear in libraries and even in expose in PyTorch, because everything in the end relies on nickel. So, nickel is there. You can just use it, create your custom kernels, integrate them, and and benefit from the diffusion. All right. And speaking of one kernel, that's interesting. I let's continue on nickel EP. Exactly. So,

with the great tools we have, we are so tempted to create what will be more meaningful to user applications. Nickel EP is our pride example here. Um so, it's a counterpart to deep EP in terms of API and functionality. And this is a slide made by notebook LM. A lot of details, but I'm not going to go through it. But you can tell that it provide dispatch

and combine API for MOE token exchange. Um so, you may wonder why we develop nickel EP, right? There's deep EP there. There are two reasons. First, we want the program to have a single communication runtime. So, if you have been already using nickel for the traditional all reduce, users probably don't want to pay another price in terms of memory to allocate resource for another communication runtime, right?

And the second reason is that as Shuai mentioned, nickel gen actually provide better feature or finer grain control than what's being used in deep EP, which is Elisha men. it allows people to control how they want the network operation to be performed, for example, on the Q pair level. But it also remove some of the complexity of direct RDMA course. So, this is a sweet spot that

goes uh some balance between Elisha men and direct RDMA. This is a performance comparison with deep EP on the low latency mode. As you can tell, we are basically on par with the EP and as the scale goes bigger, there's a small gain. Uh This is the resource of Nico EP. Today, the source code is already open sourced in the Nico repo under the country folder. And

I want to emphasize that for Nico EP, we start to embrace the GitHub first approach. So, we welcome open source PRs being posted there and you're also welcome to request for features. Um if you want to know more details about Nico EP, there's a paper on archive. And integration into the framework is also in progress. For example, I'm just giving a very short sketch of how it

may look like in PyTorch distributed. Hopefully, we can have a token switch that handles the token exchange. Um next, I'm going to talk about some of the more general PyTorch symmetric memory features because we are all excited about PyTorch uh development here. we know that symmetric memory requires a rendezvous and that may have some cost. So, MemPool is actually the tool we built to remove the overhead.

Um there are couple methods to use the MemPool for backing symmetric memory. For example, if you're using Cemento empty, we are already using a implicit MemPool at the background to back it. Uh so, your experience will be much similar to just calling torch.empty. Um so, when a tensor dies and you want to create another tensor with similar shape, PyTorch will already automatically use the same storage that

has been previously allocated as well as the render wool state. So, the second time is actually uh almost zero overhead. We also may run into case that um we want to call a collective on a tensor, and that tensor is the result of previous computation. So, in this case, we provided uh c men get men pool utility, so that you can wrap the men pool around the

previous op. Here, the Y will be created on symmetric Last but not the least, uh we're also automatic torch.compile support. Uh for writing an op called full taking an input, you can register with torch.compile using this API telling torch.compile that the input tensor must be symmetric. And next, when you compile the program, torch.compile will automatically put the tensor in symmetric memory for you. We also have a

lot of exciting collectives available in uh PyTorch symmetric memory today. Uh they are all available through torch.ops.cmen namespace. Uh three categories, low latency ones, you can got all flavors of all reduce one shot two shot multi men. Um one sided uh operation, put and wait as you I mentioned, we also have uh PyTorch exposure. Um and then, another category is irregular patterns, right? For example, we have

some all to all flavor that's good for dispatch in MOE. And these are exactly the ops I mentioned before that can consume the metadata on device. We also have um ops that does reduce that in uh tile shape that goes for shampoo plus FSDP. Um some other ops like reduce scatter offset, they can do column wise reduce scatter. And last but not least, all to all permute

that can do permutation with auto on the fly. Your contribution is most welcome here because uh we we need to know the uh the demand of developer uh and you know the demand of your application. Uh the other feature enabled by symmetric memory is called copy engine collective. Um because symmetric memory can provide remote pointers, we can directly call CUDA memcopy API uh using the symmetric pointers.

This enables collective that doesn't require communication. Uh that doesn't require uh computation. Uh and this is a small sketch of how you enable CE all gather with symmetric with a micro benchmark, we can see that by overlapping the CE all gather and matmul, we actually have 14% gain compared to the non-CE case. And this has been integrated into uh FSDP with uh the colleagues in PyTorch. So,

you just need to call model.set_seman_for_com and then boom. The uh profile shows that you're already using CE for all gather. Uh so, please tell us what you need uh and uh also please feel free to contribute ops here. With that, uh we are happy to take questions. If you have any question, please raise your hand and uh our assistants will hand you the mic. So, there's one

question over here. Uh thank you for your speech. I have one question about the Nico EP. Uh could you briefly summarize the major difference and innovations when compared with DPEEP? And maybe give us some examples difference when applied into the real case. Thank you. Uh yeah, that's a good question. So, Nico EP is a project that we just started. So, today the status that we get functionality

on par and performance on par. But, later on we will definitely provide more APIs, for example, to fit better with other group chain requirements of say PyTorch. Um second, the DPEEP, right, was built at a time where Hopper GPU is most used. So, moving to future, we may have like Blackwell support, um maybe next generation Rubin support, Um and Nico team is committed to port the performance

of Nico EP across different GPU architectures. Hi. Uh super cool talk. My question is around like NCCL logs. Uh the NCCL log, especially the error logs, are quite hard to read if something goes wrong. So, and all we see is like a timeout. Do we see some improvements in logging? >> Do you want to start with this or No, I think that's your question. >> Sure, all

right. Yeah, okay. So, that's yeah, that's a very common question actually and a very good one that I answer a lot. Um yeah, so time outs usually uh don't even come from nickel. Uh they come from usually PyTorch timing out at some point and saying, "Oh, nickel time nickel operation timed out." Uh the problem is that nickel often has nothing to do with it. Uh it's just

like some GPU didn't show up and well, what can we do about it? We're just a communication library. Um so um but I agree though that's especially in the past the message was just like time Uh and I understand that can be confusing. Um and so yeah, what we're trying to do here is to improve just the message so that you understand better what happened. We won't

fix the fact that maybe some GPU died somewhere or stuck or got stuck for some reason. That could be extremely complex and that we don't even have visibility on. Uh but we try to first um make sure that the error message were properly reported. But when it's a hang, well, there is no error message. We're just waiting for someone to show up, which will probably never show

up because it's it's stuck. Um one thing we we're trying to do though is to actually improve the message to at least tell you which one did not show up. Uh we actually added um a feature in nickel called RAS uh that lets lets you get a status of where nickel is. So, in RAS if you get the the report, you should get something like all those

ranks are waiting in an operation, but that one over there is not there. So, at least you know which process and which node which process ID is stuck. You can even attach GDB to it to understand why it's stuck or maybe disappeared, in which case RAS will just tell you, "Well, I could not contact that rank, so it just disappeared. It's no longer there." Um although in

that case normally we should get an error and then it should just stop before, but in any case the rust system is kind of designed for that. We're trying to see if we can integrate that into PyTorch so that when at least when PyTorch times out it queries the rust report to at least print it so that you get that information and then you what happened and

and or what in which state we were. It could be that everything was there and just there's a bug in nickel. That's pretty rare though, but I hope. Uh, but yeah, most of the time you see that there is a mismatch some calls were some ranks were actually behind and did not call into the collective operation. And there's recently a blog on PyTorch about flight recorder. I

think that's another utility we can use to give you a better understanding of what's going on. Yeah, nickel can only see the things inside nickel and PyTorch is actually at an even better level to give you an understanding of what was happening. We can tell you we have that all reduce operation that communicator. But if you can't translate that into oh, that's like my TP dimension or

T or PP or EP it can be hard to understand so PyTorch being level above can give even more uh, even easier to understand explanation. Any other question? All right. Thank you for the time. Thank you. Enjoy the rest of the >> day.