Optimizing PyTorch on CPU-GPU Coherent Platforms - Matthias Jouanneaux, Nvidia

In this session, Matthias, a software engineer for PyTorch at Nvidia, discusses the optimization of PyTorch for CPU-GPU coherent systems. He provides an overview of systems utilizing coherent memory, distinguishing between those with physically separate and unified memory. The talk focuses on optimizing CPU-GPU transfers within PyTorch, highlighting the use of pinned memory and custom tensor classes to enhance performance. The speaker also elaborates on offloading strategies to improve throughput by utilizing interconnect bandwidth effectively. He concludes by outlining key considerations for evaluating offloading's benefits in various workloads, emphasizing the need for a thorough memory breakdown and potential bandwidth analysis.

So, my name's Matthias. Uh, I'm German. Uh, I have kind of partially grown up close to um to here actually. So, so I do speak French uh in case you want to talk to me. >> [snorts] >> And I'm a software engineer for PyTorch at Nvidia. Uh, in particular we're working on kind new hardware features and and how to map them to to PyTorch. Um, so today

I'll talk about um CPU GPU coherent systems. And um uh yeah, uh what what you can do to kind of optimize PyTorch for for those platforms. Um, so first of all, quick overview of systems that have been recently introduced and are using some sort of coherent memory. Um, you have basically all the sort of flagship systems from Nvidia and uh AMD as well. So, the top row

here um represents data center platforms. And then um we also have recently like a boom I would say in um uh unified memory in the prosumer or like yeah, um embedded um uh you know, uh stacks where um they will let you essentially utilize the entire um or they they just have like a a physically um unified memory. And so, this is also how we can kind

of separate those. It's going to be very important for optimizing obviously. Um, some of these systems in particular all the you know, Grace Hopper Grace Black systems from Nvidia, they have a physically separate CPU and GPU memory. Um, whereas all the other systems and in particular those consumer prosumer um you know, both from Nvidia and Apple and and AMD have a physically unified memory. Um, so you

as you might imagine this this changes things if you want to optimize. Um, and then in terms of the you know, what we're going to look at, um I I wanted to keep it a bit low level because talking about specific use cases would be would just take way too long. there's sort of four things that you could do. You could have independent work running on the

CPU and the GPU on on independent data. But there's no real change with when you're using unified memory in that case, right? Um, so we'll just ignore that. Um, makes our life easier. Then [snorts] the second thing is you could have um actual um concurrent work on the same data between the CPU and GPU. Um, that's something that we we have seen use cases from customers at

Nvidia um for this uh in like physics simulations for example, but in PyTorch I think it's quite rare. So, I will also skip this for now, but you know, please reach out if you you can always chat me up. Uh, and we can discuss this um more. Um, just for simplicity I'll leave it out of this talk. >> And then finally, we have sort of explicit you

know, CPU GPU handoffs or transfers where you will you know, take some you have some CPU processing and then that data needs to move to the GPU which which continues processing or just the other way around. And [snorts] luckily, those are very similar. There are some subtle differences between um moving data from the CPU to the GPU and moving them from the GPU to the CPU. But

I'll again kind of ignore those subtle differences and we'll just focus on CPU GPU transfer. So, this [snorts] is how this would kind of look like in code in PyTorch. Um, the important part is essentially just the last three lines which are um the ones that that are doing the compute. So, the first one of those is doing like this torch add um with a uh just

a CPU tensor and then we output into our kind of transfer tensor um which is in this case allocated using pinned memory. Then there's an explicit transfer to the GPU which you can see in green. So, that pinned memory trans tensor is going to be explicitly copied essentially to the GPU's memory. So, here we use non-blocking true. Um, that has some benefits um but there's a great

tutorial which I've referenced which discusses all the intricacies of using you know, pinned memory as well as non-blocking. So, I'll just defer to And [snorts] and then finally, what we get out of this two operation is obviously a CUDA tensor which we can input to a CUDA or GPU based operation, right? That's the last line. The last line will be executed on the GPU and it's both

of its inputs are in GPU memory and then its output is also in Um, now imagine you want to run this on DGX Spark or similar platform with physically unified >> So, there's a bit of a problem because um that explicit transfer it effectively just duplicates memory, right? So, you have the same physical memory for CPU So, if you if you just do this explicit transfer, you're

going to um simply duplicate memory. Now, it often isn't such a big deal. And as we'll see in in in the later part, it's not such a big deal because many traditional applications are actually written with very poor kind of interconnects between CPU and GPU in mind. And because of this, most traditional applications will sort of treat the interconnect as the devil. Like you should never transfer

anything um ideally between the CPU and the GPU uh for best performance. And of course, that means in many cases um these transfers are pretty small. But if you do have an application that actually needs to transfer a lot of data between CPU and GPU, um then this becomes an issue, right? Because you have you you're going to duplicate a lot of memory in this case. now

there is um you know, one way you could solve this is like let's say we just remove this green part where we're duplicating and we basically just input a pinned memory tensor or or just a regular CPU tensor to the to the last line the torch add. Um, you'll see that you'll get uh an error today in PyTorch because there is a bit of an issue with

dispatch in this case. Um, if you're inputting uh GPU tensor and a CPU tensor to um to one you know, operation, it's not completely clear where to even execute that. So, you would then have to if if PyTorch would accept this, you you'd basically have to sort of solve that problem first of you know, how do you actually like on on which device do you actually want

to execute things um depending on on the inputs, maybe also depending on size constraints and and so on. It gets quite complicated. Um, and to just avoid this, right? PyTorch just throws an error essentially. >> [clears throat] >> we could on um something like Spark, we could basically um sort of hack a um a view to the tensor um using the same and [snorts] um and just

using some CPU tensor and then just kind of hack a view to it um as seen by the GPU. Um, that would work. It would be functional, but there's still some issue with this because um uh I'll I'll also mention this on the next slide. So, uh when you just use regular um sort of system allocated memory as we call it, um so something like lib C

malloc, um that's often not ideal for GPU processing. So, if you were to just take a CPU tensor and then you sort of explicitly tell PyTorch okay, this is actually a GPU um tensor but backed by CPU data. We input it uh to a to this GPU operation. Um, it gets executed on a GPU, but you might run into issues like um page size um problems um

and and also actually alignment which might not be um respected by your CPU allocator So, what you can do um I'm going to refer to to a little code snippet that I that I published uh recently. Um, you can kind of fix these issues with a custom tensor class. Um, so here on the right, we create a custom class um called transfer tensor in that case. And

it basically sort of fixes all of the problems. For one, um it uses the torch function interface so that you can um easily um kind of fix the dispatch um problem. Uh, it then uh has also um an explicit sort of um backing memory you know, either uh pinned memory or managed memory. And it also because it's a custom tensor class, we can basically also override autograd

semantics. So, it will even work if you want to propagate the tensor um between the transfers essentially. Um, so link to the um to that um is provided. Um, I'll I'll um provide the slides uh through through through the interface later. so there's one thing I wanted to mention. There is actually a difference on platforms like DGX Spark between pinned memory and managed memory um because uh

this is kind of the way that um Nvidia has decided to expose different coherency protocols. So, pinned memory is non-coherent or non-cache coherent on the GPU which means that it's not cached in GPU L2. Um, but it does have because of this, it doesn't have overheads from from the coherency protocol. Meaning that it's it's actually better in performance for streaming accesses. So, if you're doing something like

just copying the data or like in this case, um [snorts] uh the first operation is like an addition which is typically memory bound. It's a streaming type access. You should be using pinned memory on on Spark. If GPU L2 caching is really important, then um you might want to use managed memory instead. Um, in general the I think the sort of my my recommendation would be start

by using pin memory. It's it's usually the better option and then if you're sure that um, that you know, you're you're running to like L2 cache issues then use managed Okay, so moving on to Grace Hopper and Blackwell um, platforms. here the the story is is quite different. So, let's look back at our original example of of this sort of CPU to GPU track um, yeah, transfer

um, you know, traditional uh, implementation. Uh, with Grace Hopper this just works. Um, you have physically separate CPU and GPU memory anyway, so there isn't really a significant reason to um, you know, to to to make um, so sort of to avoid the duplication. You can avoid it because Grace Hopper and Grace Blackwell provide you this coherency protocol. Um, and that might become interesting at the point

where these transfer tensors are, you know, it and in sizes in the order of magnitude of your combined CPU and GPU memory. So, these tensors would have to be really really large. Um, in in order to um, um, to to, you know, for this duplication to really matter on Grace Hopper and Blackwell. Um, so for almost all applications it won't matter. Uh, and on top of that

even if you do have tensors that are really really large, what you can typically do is to just basically um, you know, split down yeah, you do a sort of divide and conquer strategy, right? You just split um, your your transfer into um, many smaller parts and then that means your actual transfer tensor that you that you copy over becomes very small and you just make, you

know, you take a size that basically saturates the um, interconnect bandwidth and that's pretty much it. You're good to go. So, um, you can always basically make sure that you come back down to such an example with a pretty small tensor. And the advantage is that these platforms in order to kind of give you the illusion of coherency, they typically have very high interconnect bandwidth. So, here's

a on the right you can see an example for Grace Hopper compared to x86 and Hopper. Um, it's uh, you know, one thing that you can see is significantly um, better bandwidth. Um, the other thing that we've also shown here is that on these platforms it doesn't matter whether you use pin memory or pageable memory because it's essentially the same thing. but at the same time there

isn't really a reason to rewrite your code to use pageable memory. So, if you're already that's probably fine. There are some subtle differences between, you know, when you should be using pin memory or or pageable memory. I'm not going to go into those. Um, uh, but uh, all the, you know, all the edge cases where it really matters that you're using pin memory, those would be the

same anyway. Like whether you're using a x86 based system or or Grace Hopper or or anything. Um, so there isn't really a good reason not to use pin memory, so I would just recommend doing that. and of course you you get significant improvements in in bandwidth over the interconnect. >> So, here's a bit of a recap um, of the different systems and then I'm also showing kind

of the ratio. So, the important column is on the very right. Basically, I'm showing the the ratio of peak um, you know, DRAM video or GPU DRAM um, bandwidth compared to uh, the interconnect bandwidth. And what you can see is I mean of of course a system like which has physically unified memory basically has a sort of one-to-one ratio. So, in that case you should basically never

be using the interconnect, right? So, you should always avoid uh, sending anything. And the traditional systems, they have a typically a pretty high ratio, um, which means that you should also generally it it's it's, you know, a bit hand-wavy, but um, in many cases um, it it doesn't really make sense to to transfer a lot of data. But then these systems like Grace Hopper and Grace Blackwell,

they have significantly lower ratios, which which means that sometimes it does actually make sense to transfer data explicitly in order to improve your performance. But it becomes a bit of a chicken and egg problem because as you can see like both of the other types of systems, they there you actually want to avoid transferring anything essentially. So, um, this is a bit special to those systems. Like

as as you kind of move the ratio towards one-to-one you you want to transfer more and more and then at some point if it's like actually one-to-one or or very close to one-to-one then then you basically um, want to just treat it as one unified memory instead. so uh, I wanted to also um, talk a little bit about CPU offloading um, where you know, you can use

the interconnect bandwidth essentially to to improve throughputs or or improve other measures. Um, and this is an example from a GTC talk which I gave actually 2 years ago already um, on Grace Hopper. Um, this is for Llama 2 70B. I guess it was the um, the latest Llama model at the time. I know it's not anymore, but um, this uh, you know, applies um, in very

similar ways to to recent models as well. Um, it's just that there are so many models that I could potentially pick now that um, I didn't see a a reason to to to um, to pick uh, something else. Um, back then we showed that you can um, improve throughput uh, sort of per GPU throughput by up to 60 or 67%. And this is compared to a DGX

platform. So, this is basically taking kind of um, you know, one single Grace Hopper system comparing to a DGX platform. So, obviously it's not exactly um, apples to apples in some sense because with the DGX since you can use up to eight GPUs, um, you uh, you have a overall you're going to have a lower runtime obviously in absolute terms, but the per GPU throughput is going

to be higher with the single Grace Hopper system using offloading. And then even if you use multiple Grace Hopper systems, so here we we went up to four, um, we we still were able to get um, you know, more than 40% um, improved throughput uh, through offloading. So, this is something that can be used and so I wanted to show a little bit for one like how

you can actually do this and secondly how you can decide whether offloading like so that that you can have sort of the tools to decide whether offloading makes sense in your case or not. >> first how can you do this? Typically um, all implementations will use uh, multiple streams. So, you basically just use um, so you have your compute stream anyway to for for PyTorch applications and

you do use additional streams, one for offloading and one for prefetching at least. In this particular case, uh, I'm showing an example where we're doing both offloading of weights and activations. So, then I'm I'm also using additional stream um, for for, you know, offloading the weights and offloading the activations. and this is something so this kind of implementation, it's already available in Transformer Engine for example. So,

if you're able to use um, this through Transformer Engine, you can just um, use it as is. And there is also a a PR um, open on on the PyTorch repo or under discussion at least for um, supporting this in in FSDP directly. yeah, so in the next slides I wanted to to show a little bit um, how how you can decide whether offloading makes sense. And

sometimes this I think um, is a little bit counterintuitive. So, so I I I'm using a bit of a hypothetical scenario here to kind of show all the different things that that you might see on your applications. So, this is using still the assumption of a Grace Hopper system. You can apply this the same way to um, Grace Blackwell obviously just changing the the assumptions. So, here

we have some, you know, number of uh, um, gigabytes per second for our GPU bandwidth. Have our interconnect bandwidth. And then what's very important to start with uh, is to make a sort of memory breakdown, right? So, you you want to know um, what your model where your model basically spends um, memory. So, in this case we assume there's, you know, most of the um, gigabytes are

used by the weights, then some activations and there is very little unused memory. Um, the next thing is you want to know what's the iteration time, right? So, you you want to benchmark this model. In this case for example, we'll we'll just assume that we just do a forward pass just to you know, make it a little easier, but um, you could also um, do this with

like a forward and backward pass of course. Um, but you will want to know exactly how much time is spent. Um, so that's very simple benchmarking of your model. And from this you can already determine a bandwidth that your model takes, right? So, you can basically say uh, I know my weights are typically read once during the iterations and my activations here's a little bit more complicated.

Um, you would have to know exactly which activations go or stay in L2 or how much activation stay in L2 or not. Um, you know, if if you don't know that just go with all the activations go to DRAM. It's it's a a pretty decent um, approximation. Here I'm just going with like 15 gigabytes out of the 21s are going to DRAM. And we multiply that by

two because we activation we typically have to write them and read them back um, when we uh, when we use them. And um, then you divide that by the runtime and you get some some bandwidth, right? In this case, it's uh, about a terabyte per second. This is very important because if this is already very close to your peak bandwidth, then offloading will never be beneficial because

you won't be basically use the bandwidth for both, you know, the offloading and your model at the same time, right? Like you cannot uh so barring some some special cases, you you cannot um use both the interconnect and your compute to um move data in and out of the DRAM, right? So, if this was like, let's say, already above 80% or 90% of your peak bandwidth, um

then offloading is likely not going to bene- be beneficial. Here, it's significantly below, so we're fine. The next thing is we can determine what's the bandwidth that we need for like a naive offloading strategy. So, for example, if we were to offload the weights here, we need to move 72 GB both in and out of the GPU. Um you can divide that by our iteration time, that's

how much we time time we have to do that, and we'll get something like, you know, 1,500 um GB per second in this case, and you can see this is much higher than the 450 that our interconnect provides. So, at first, you might think, okay, it's hopeless, we cannot gain anything. But this is where the next step comes in. It's very important to figure out what is

actually what you can actually potentially gain um by uh by doing offload, right? So, sometimes you cannot gain anything, and then yes, indeed, it's uh it's going to be hopeless here. But let's assume um so at the bottom here, let's assume if we can increase our micro batch size by factor four, we notice that the run time only increases by factor of two. You might think, okay,

how can I even know this? Um well, typically, a lot of models are uh decomposed into many layers, so what you can do is to just time a single layer, and then for that single layer, you you run again with different micro batch sizes, for example, there you should be using much less memory, so you won't go out of capacity, right? And and you can then determine

what how the run time evolves. So, if our run time is, you know, um increased by 2x from 4x improve- uh 4x increase of micro batch um size, then that means we basically get a 2x improvement of throughput. You can do a sort of the math. I'm not going to go through through um through everything here because the run time is increased by 2x, it means that

ultimately, instead of 1,500 GB per second, we only need need 750 GB per second in reality um for the offloading. That's still higher, but the throughput improvements are sort of larger than what we um than what we lose essentially in this case from being bottlenecked by the uh interconnect. So, here we are bottlenecked by the interconnect, but we can still gain throughput. Um but of course, we

also don't have to offload everything. We can also just offload offload, for example, a third of the weights in this case, and if we assume in that case that we get, you know, by improving the micro size by factor of two, we get 40% improvement of throughput, right? Then in that case, we can actually gain the entire throughput improvement you know, because we're not bottlenecked by interconnect

bandwidth at all. Um you might also wonder where these kind of numbers for throughput improvement come from. Um typically, that's just um because micro batch sizes are often very small uh in these applications, and um you might be able to significantly improve um the efficiency of kernels outside of the matrix multiplications, and the matrix multiplications themselves might be able to use different kind of algorithms, which improves

um throughput quite a bit. Um so, I'm going to I want to talk a little bit about the workloads. Um in many cases, inference is heavily uh memory bound, so in that case, as I mentioned, it's not um all that useful, but there are, of course, uh use cases in inference that benefit from offloading. Large-scale training depends a little bit. Uh it's typically compute bound, but um

the gains, the potential gains, are often a bit questionable. And in fine-tuning, uh we that's where we mostly see this uh being useful because first of all, um you're often kind of limited by the capacity that you that you essentially want to use for for the system, and um you uh you're still compute bound because it's a training-like workload. it's also always good to compare this against

other techniques to save memory capacity. Like in particular, activation checkpointing comes to mind. This is actually a little bit orthogonal, but, you know, you can do similar sort of math if you know your um forward and backward times for the uh models that you are or the modules that are you are you're checkpointing, you can pretty easily determine what's the throughput gain from checkpointing. Uh by the

way, here, I'm assuming again that we can sort of improve like we can increase the micro batch size by the same factors for checkpointing than for offloading. That's uh you know, not at all the case in reality. Um you um because the the the increase in micro batch size for checkpointing depends on basically how much intermediate activations you have uh for those modules that you're checkpointing. Um

but, you know, for simplicity, I'm just assuming this, and you know, here you can see at the bottom a sort of breakdown. Um and you you you'll get, you know, some some improvements um depending on on the technique. Yeah, so ultimately, you know, this is sort of the decision framework um should hopefully should should help you determine whether offloading is useful, right? For one, make a memory

breakdown. Um you should you should at least know which tensors are consuming memory and and which are not. Um determine the memory bandwidth of your model. Uh then also very important, evaluate what you can even gain um because if you can't gain anything, then um offloading is not going to be useful. And finally, um you can compare or even, you know, combine with other memory um reduction

techniques. So, there's activation checkpointing. In the example from from the GPT-C talk, we also basically compared it against tensor parallelism. You could also um include context parallelism. There's, you know, it it can get as complicated as you want to make it. Um and then of course, you can also um do all of this, you know, for just partial modules, so you can offload your model partially, you

can checkpoint partially, and so on. The possibilities are endless. Um and that's why I uh I excluded sort of any any combined, you know, analysis because that that would just go completely off the rails. Um yeah, so that's it from from me. sorry for taking a little bit longer, and uh if you have any questions, then then just uh yeah, see me um in the hallway or

here. >> [applause]