Heat: scaling the Python scientific stack to HPC systems

Speaker 1 [00:00]

Our next and last talk of this session is Scaling Python-Based Analysis to HPC Systems, presented to us by wonderful Thomas Zaupe. Welcome.

Speaker 2 [00:12]

Thank you for the introduction. And also thank you for being here at this late time. So yeah, I'm going to present HEAT, which is a Python library for distributed data analysis. It's an open source code that's developed at Forschungszentrum Jülich and KIT in Karlsruhe. So let's start with a bit of a motivational slide. We have some pretty pictures of people doing science with data. They somehow analyze some data, and a lot of the time when you do that, you use NumPy or the related libraries to do that. And for very good reason, because they're very robust, they work. We know how to use them at this point. But there's one fundamental limitation of NumPy and the related libraries, which is that they're limited to shared memory parallelization and CPUs. And that's unfortunately not really the landscape of the compute resources that we have these days. So maybe if you want to do some data analysis, you start coding on your laptop. And this has a few gigabytes of memory and a few CPU cores. But this is really limiting in the size of data that you can process in the end. Because you need to fit your data into memory. And then you also need to have the patience until all of the data is processed. So you want to use some larger compute resources. Maybe you have a big workstation available or access to some high-performance computing system. And then if you have written code in NumPy, then this will give you access to maybe a single node of the HPC system, which will maybe 10x the amount of memory that you have available and the CPU cores. But maybe this already has some GPUs that you then can't leverage with your NumPy code. And we now have a new system in Jülich called Jupyter. This is not actually an image of Jupyter, to be honest. It's not very photogenic. It's in some containers. So this is some old... I think at this point, decommissioned CPU system, but it is a system that has once been in Jülich. Like the new Jupyter system, anyways, I'm getting a bit sidetracked. There's like 6,000 compute nodes. So if you could scale to all of them, you could increase the size of data set that you can process by orders of magnitude. But you need the software to scale there. And this is sort of the point of heat. Let's try it again. Okay. So the point of heat is to work very similar to NumPy. We try to mirror the NumPy API as much as possible, and then you code on your laptop. Works fine. But then you just go to the big system, and it still works. But now you can process large data sets. And it does that by distributing the data. So here's an example of a 3D data set and the three different ways that heat can distribute this data because it only supports distribution on one axis. But yeah, hopefully this will work. Before we talk a bit more about the details of heat, I want to mention that there is a a lot of similar libraries available. So you probably heard this pitch before, like give us your NumPy code and we make it fast and distribute it in GPU and everything. And it's not like Heat is fundamentally better than any of these in some amazing way. But I think all of these libraries have a few advantages and disadvantages. Like in the feature set that they support and in their API, like they are a little bit different now than what you're used to from NumPy or they don't support all of the operations in a distributed fashion because some of them really are not trivial to parallelize. And I think this is where HEAT does have some selling points because it implements a lot of these linear algebra functions that are not trivial to parallelize. And maybe this is not the best fit for you, but maybe not. So feel free to check it out and see what's available. And by the way, if at any point you are convinced that this is something that you maybe want to try out, this QR code that's almost always there will lead you to the GitHub repository. All right, so now let's look a bit under the hood and see how it works. I took a simple screenshot of a little interactive session that I ran on my laptop to illustrate the very core concept of heat, which is the DND array, the distributed n-dimensional array that's kind of the equivalent of the NumPy array. So first, let's start with writing some NumPy code. So we... Oh, you can't see my cursor. So we import NumPy as we normally do, and then we define some NumPy data. and now this is what we want to accelerate with heat so we import heat instead and then we call exactly the same function on heat with some additional heat specific arguments what does this give us? first of all it gives us the same data and we can use this heat dnd array pretty much the same way as we would the numpy array so most of our favorite numpy functions have a heat equivalent that we can just plug this into and we're good to go, basically. But let's look at the additional arguments. So we have split equals zero here. The split tells Heat how you want to distribute the data. Earlier we've seen the 3D data where we had three options for distributing plus no distribution, I guess. This is 1D data, so there's only one way of distributing that along the first axis. And this is how you do that, very simple. And then the second argument, that NumPy also doesn't have is the device. Here I use MPS, which is somehow the accelerator in my laptop. He says this is a GPU, whatever. And then there's some additional information regarding the parallelization. So we have MPI rank zero and a local shape, which in this case is equal to the global shape because I execute this in serial. But if I ran this with multiple tasks, then we would have multiple MPI ranks and smaller local shape than the global shape, of course. Now, what's happening underneath, we can check out the local data on the task by just looking at .L array, so local array. And this is just a Torch tensor, actually. So he just built on Torch. Of course, we don't reimplement all of the things. Like whenever we call a serial operation, we just pass this on to Torch. which also supports all the devices. So that's good. Another heat-specific attribute of the DND array is the communicator.com. We have some abstraction around communicators, so this could be different parallelization frameworks, but for now it's just a wrapped MPI4Pi MPI communicator, which works on a lot of machines, I think. The DND array is really a torch tensor which supports all kinds of devices. They have, of course, good implementations for serial operations, which we can just leverage in heat, plus the MPI4Pi communicator, which I can use to parallelize on my laptop and also on big machines like Jupyter. It really runs on a lot of machines. It has a NumPy-like interface, which should make it very easy for you to get going if you're familiar with NumPy. Let's look now at a parallel example. So, I executed this script with four tasks now. We first define some data again. So, we just call heat arrange with 16 elements now. Here we haven't passed a split argument, so we define the same data on all tasks. Then we call a reshape operation, which we're familiar with from NumPy as well. And then we call a heat-specific function called resplit. And this just changes the distribution. It's kind of a dangerous function, maybe, because this will do a lot of communication underneath. And depending on the data set, you might run out of memory if you transpose it. But anyways, if you want to change the distribution, it's really very easy. Now we can print what we got. And you'll notice it looks quite pretty, right? So it's a single output on the first rank. There's no matching the output to MPI rank and print statement because the order is unexpected. So this should also make it easy to develop the code for people who are maybe not experts in parallelization. And our target audience is really people who are experts on analysing the data and not necessarily on parallel computing. But, of course, we can take a look under the hood and print, for instance, the rank of the communicator that's associated with the array and the local data. And then we see how the distribution works. So we actually have only a patch of the data on each task. And again, this is just a torch tensor. Now, I wrote here the distribution behaves sanely across operations, so maybe the distribution changes if you sum along the distributed axis, for instance. This axis is gone afterwards, and how should this be distributed? Probably it shouldn't be, and this is exactly what happens in heat. If, on the other hand, you sum along the non-distributed axis, you would expect it to still be distributed. And that's what he does, right? So ideally you don't have to worry. That's that's the point Okay, so so far if you're familiar with a bit of MPI this seemed very trivial but really the selling point of heat is that it implements functions that are not trivial to parallelize and I thought a good example maybe is QR factorization So here we take some matrix and we factorize this in an orthogonal matrix Q and an other triangular matrix R, which we can multiply together to recover the original matrix. And now in this example, we start again by defining some data, which is split along the first axis. And then we just call the functions just like a numpy, and it will give us these matrices, and they're all split also along axis zero. So if you have some code and you want to distribute it, you have to distribute really when you define the data or when you load the data, right? Like in practice, if you are doing some data analysis, you will load data. And during the load step, you just say split along this axis and then heat will just continue distributing along this axis. And also it will maintain the device. Okay, now let's look at some actual parallel scaling results, because this is of course the important thing. Like, is it at all faster? We start with a very basic operation, which is matrix multiplication. I show here what you have to communicate with this simple example on top, which is multiplication of two 2x2 matrices, which are split along the first axis. So the color denotes which element or which task the element started out on. And then in the result, you will see in the first line some blue things, which we had to communicate from the second task. And in the second line, some red things, which we had to communicate from the first task. So there's a lot of fine-grained communication going on in matrix multiplication. But even though we have to communicate a lot, we still get some decent speed up with heat. So I did hear some weak scaling, strong scaling combination runs on CPU. I have here the NumPy baseline, that is this black dashed line. And then the blue line is heat with one CPU. We see for tiny matrices, we have some overhead in the heat library. But for large matrices, we essentially spend all of the time in torch matrix multiplication, which is very similar to NumPy matrix multiplication performance-wise in this test. But then heat allows us to also distribute this. Then this yellow line is heat on two tasks now. We see for tiny matrices, we have just some communication overhead that is not worth doing. But that's not what heat is about. It's about processing larger data. And pretty soon we cross over the serial line And then we're faster than serial. And then here we increase the number of tasks and we gain even more speedup. This is the largest example that I had the patience to run with NumPy. And so for this example, I show the speedup plot on the right where the blue dots are reasonably close to the ideal line, I think. But of course, there's also the dots, which are like 100x speedups, which is GPU. because we base it on PyTorch rather than NumPy. We can just do it on GPU, and that's really much faster. So again, for tiny matrices, the GPU is not better, but pretty soon there's a crossover point where then the GPU is faster than CPU. Multi-GPU, you now need to do more compute to gain something by distribution on multi-GPU. But also here we eventually have a crossover point where heat with multi-GPU is faster than PyTorch with single GPU. Okay, but a very important metric is also the memory consumption, right? Because you need to fit the data set into memory. So here I have a different weak scaling test that I didn't record myself. So let's look at the top plot first, which is again time. And we see if we increase the number of GPUs, we have increased the data size even more so that we expect it to be slower if we increase the number of GPUs, which is maybe not ideal for a weak scaling test. But anyways, that's the test we ran here. And we tested different combinations of split matrices. You can have different split in both matrices that you multiply and you get different results. But anyways, if you split, for instance, both along x is zero, it's fine. This top dashed dotted line is the projection on a single GPU, and the dashed line is the ideal speedup that we expect, and we're somehow in the middle. So, okay, timing could maybe be better. But, yeah, like, if you've done some GPU programming, you know that the single GPU memory limit is really, can be very annoying, right? and at some point you need to distribute regardless of whether you gain time or not, like time speed up or not. And the maximum memory consumption per GPU that we measured during these runs is on the bottom plot. Here again, the projected single GPU memory usage is the dashed dotted line, which goes up a lot, very fast. And then the sort of ideal memory consumption is the dashed line. And now, going from single GPU to multi-GPU, we need a lot more memory on each GPU, because we do the communication, and then we have some additional data from other tasks now on the GPU. But once we are multi-GPU, it scales pretty well, and this allows us to fit a lot larger matrices. Yeah. Good. All right, so now also the QR factorization we looked at earlier, I recorded very similar scaling results, or scaling tests for this. And the result is similar, like the speedup is a bit worse, but that's not very surprising considering you need to do additional computation also in this parallel QR factorization because of the algorithm. But you do get speedup, which I think is good. And then we also have a very similar test with memory consumption. And here, again, going from single GPU to multi-GPU, we need a lot more memory. But once we're there, we're pretty flat. So that's good. And we could factorize even larger matrices still. And also here on the top plot, there's no different splits of the input matrix. but both are somewhere between the ideal scaling and the worst scaling. I think this doesn't look too bad. Okay, so a very nice feature of HEAT is also the distributed singular value decomposition or principal component analysis. Here I have an illustration from the paper that describes the algorithm that we implemented HEAT also on the right, which describes the algorithm a little bit. So we split the matrix along tasks, then we perform the local singular value decomposition, and then we merge the local SVDs and compute an SVD on the merged matrix, which is somehow faster because we can exploit some structure there. But because we need to merge this data, we cannot communicate everything because then we run out of memory. So typically, we limit the rank of these matrices. I should mention that this is kind of a simplified version where you merge all of them on a single task, but there's also a tree version of this, which is better. And I think in Heat, we also implement the tree version. So you don't have to truncate at super small matrices. But these speedup plots here on the left, they are now for a different truncation rank. I don't know, this might be a bit small, but the left plot is a very small truncation for rank 5, and this is a larger truncation with 500 ranks. Now we test against scikit-learn. That is the baseline here on a single compute node. This is now a hybrid parallelization with shared memory in PyTorch and distributed memory in HEAT. And yeah, so the single node, compute node, baseline is scikit-learn, here the green line. Then these three lines, these are heat with different number of tasks per compute node. And then x-axis is number of compute nodes. And we see for low rank, this scales really quite well. This is GPU, which is again much faster and also scales decently well. If we maintain a larger rank, we have to communicate more, so the scaling is a bit worse, but still we get pretty decent speedup over the serial baseline here, which is not bad. Heed also supports dynamic mode decomposition, which is kind of a funny algorithm, I think. You look at a time series, basically, and you say the next time point is related to the previous time point by a matrix multiplication, and then you try to approximate this matrix using this dynamic mode decomposition, and then here are some scaling results for heat when doing this in parallel. The strong scaling is actually quite poor, I will admit that. But the weak scaling, this looks decent, I think. The fitting stages, the fit stage is the approximation of this A matrix, and then the predict stage is to guess the next point in the time series, I guess, if you want. But regarding memory consumption, it scales pretty well, and again, you kind of need both. need short time and low memory consumption. So, here we have also a comparison with a Dask implementation. There's a paper last year written by one of our unfortunately former collaborators at this point, which has a lot more detail on this. So, I have to speed up a bit now. And I suggest, if you're interested, to read the paper. But I will very briefly highlight a few things. So loading data with Dask, it keeps taking longer, and that's bad because eventually you will run out of time, whereas heat, nice and flat. And then for most of these operations, heat on GPUs is faster than Dask, which is limited to CPUs in these tests. Very important also the memory footprint where we see the same thing as in the time to load data. So this is here the average memory per compute node, and the x-axis is number of nodes. And if the average memory goes up with the number of nodes, at some point we hit the limit of every node. And this is exactly what happens with Dask. Here We are at like 300 gigabytes already, which the typical node has not a lot more, I would say. Heat, on the other hand, is nice and flat once we reach a full node. So here you could continue to scale further with heat, but with Dask, maybe not so much. All right, and then it seems I was fast enough. To summarize, heat is supposed to be a sort of plug-in back-end for NumPy code. I'm supposed to avoid the phrase porting code to heat because that sounds scary, so you're supposed to use heat as a back-end for your NumPy code. It makes it very easy to distribute the code and also to make use of accelerators. We have to do that because that's the compute power that we have. We have no choice but to do this. If you're interested in HEAT, reach out to us. Again, scan the QR code for the GitHub repository. If you're looking for a feature that is missing, open an issue. If you're just interested in anything such as is HEAT the right code for you, open a discussion maybe. I think discussions are a good feature for that. And our development really is user-driven. We recently published version 1.8 where we added a few missing NumPy things that some users requested. We're happy to help you out if you want. Thank you.

Speaker 1 [25:05]

Oh, that's a nice photo. Thank you, Thomas, for this excellent talk. Quick reminder, you can ask questions on talks.pycon.de as well as upvote existing questions, so they're more likely to be asked. And we already have the first question. Is heat compatible for all types NVIDIA GPUs, or is there any limitation in terms of hardware?

Speaker 2 [25:34]

Well, to be honest, I'm not 100% sure. But essentially, we just use PyTorch for all of these operations. So my impression is that PyTorch supports a lot of devices that you would normally use. There probably are some niche GPUs, maybe, that it doesn't support. But again, it even supports the accelerator in my laptop. So my guess is probably it does support the GPU that you're looking for.

Speaker 1 [26:05]

The next question is, what kind of inter-node connection did you have for your benchmarks?

Speaker 2 [26:15]

I used two different machines, actually. I can't speak to these memory benchmarks, unfortunately, because I didn't record them. They were on a DLR machine, actually, and I assume they used some InfiniBand connection because that seems to be very popular in the machines that I work on anyways. But here I used, for the GPU tests, the Jules Booster machine. Actually, I used only a single node, so these are connected with NVLink, which is quite fast. And for the CPU tests, I used a very, like, it's a smallish machine called Yusuf, and this is connected with some InfiniBand connection, yeah, I hope that answers the question sufficiently.

Speaker 1 [27:11]

And the next question, a little bit confrontational, also refers to the beginning of your talk. Chunked, distributed, and the arrays have been attempted multiple times. Dusk, czar, blusk, blusk2.

Speaker 2 [27:24]

2 cubed

Speaker 1 [27:25]

What is heat better at?

Speaker 2 [27:29]

Yeah, so first of all, I have to admit that I am pretty new to the Heat development team and so I don't have a lot of experience with the competitors. But I think at the end of the day, it really boils down to which features are implemented in the library that you might want. So again, the selling point from my point of view of Heat is really if you are looking for a parallel implementation of some linear algebra that's not very easy to do, maybe Heat has implemented this. And if so, then that's great. And also the NumPy-like interface of Heat is good. I know many other libraries also have that. I'm not going to say that Heat is better than them. But I think there's a niche of people that maybe are not coding experts and they maybe don't have access to a lot of stuff and then maybe we can take care of them. We take the time. And maybe what they're looking for is already implemented in Heat. But yeah, I mean, there's plenty of competition and I'm not going to tell you that Heat is better than the competition. Maybe it's right for you, maybe it isn't. But feel free to try it out.

Speaker 1 [28:51]

Exactly right. The next question is, which limitations have you faced with data or projects using heat?

Speaker 2 [29:00]

Yeah, unfortunately, I really haven't because I'm too new in the development team, I guess. Unfortunately, I can't really answer that question. I'm sorry.

Speaker 1 [29:19]

Last question. Have you considered adding implementing an x-array backend? A lot of scientists prefer label data.

Speaker 2 [29:30]

I don't think that's going to come any time soon. I can only speculate, but I haven't heard that in any internal discussions, so I wouldn't hold my breath for that, I guess.

Speaker 1 [29:50]

Thank you, Thomas, again for your wonderful talk. Thank you. Check out the project. Thank you all.