Zero Code Change Acceleration: familiar interfaces and high performance
The interfaces defined by libraries like Numpy, pandas or scikit-learn are the defacto standard APIs in each library's domain. Data scientists use these libraries directly as well as indirectly through libraries that depend on them.
This talk will look at the different approaches that recent efforts have taken to give users both a familiar interface and GPU acceleration. This means users do not have to rewrite their code, learn a new library and benefit from acceleration when using existing libraries.
The cuml team built a scikit-learn accelerator by diving deep into the import system of Python. By hooking into the import system you can replace the result of import sklearn with a library that uses cuml where possible and falls back to scikit-learn where necessary.
The scikit-learn team is adding experimental support to handle PyTorch and CuPy inputs by using the array API standard. Instead of using the Numpy API to perform array computations, scikit-learn is switching to using the array API. This is a subset of the Numpy API that is supported by several other array libraries. The API of Numpy and PyTorch is similar but not exactly the same, this makes writing code that works with both hard. The array API addresses this problem by providing a unified API. Users can accelerate their scikit-learn code by passing in a CuPy or PyTorch array instead of a Numpy array.
This session took place in track PyData & Scientific Libraries Stack and was classified suitable for novice domain / intermediate python by the speaker.
Transcript (auto)
Auto-generated from the recording utilizing Open-Source AI. Speaker labels (Speaker 1, Speaker 2) reflect diarization, not identity. Timestamps refer to the recording.
Speaker 1 [00:08]
Oh, it's very loud. My name is Tim, and I'm very happy that you have made it to the last session of what feels like a very long first day. At least I'm very happy that I've made it here. So, like the session chair said, I'll talk about familiar interfaces and high performance. and the whole question for this talk is how can we keep the interfaces for PyData libraries, so APIs of these libraries the same? Think of that as a constraint. We want to keep them the same and yet improve their performance. Because, for example, there are things you can do in libraries like Polars that you just cannot do in Pandas. but Polus comes with a new API so the question is if we keep the interface the same what can we do in terms of performance and now take a moment to think about what your favorite part of Python is so for me the best thing ever about all of Python is the REPL. It comes in lots of different shapes and forms you can have the old school one that's on the slide here, you can have Jupyter Notebook, you can have IPython, things like this and the reason I think this is the best thing ever in Python is that I can type a little bit of code and I can press return and it will run it and it will tell me the answer or very often it will tell me that I made a mistake and then I can improve it and then I can run it again and you know it can really quickly iterate and get towards something that works whereas when I work in something like C++ I have to like save the file and then like figure out how to call clang or gcc and then it goes and thinks about things and I have to find a new terminal then I run the executable and then it crashes and I go back to my you know you get the idea so I really like this thing that you can iterate very quickly in the REPL and that's something I think is interesting if you're doing like data science kind of things you maybe got a science problem or a data problem in your head and you want to iterate really quickly. So the Python data ecosystem is ginormous. Here's a few libraries that exist in it and it's often considered like a stack. now everything builds on top of Python and there's a library like numpy and then which builds on top of Python and there's libraries like scikit-learn that build on top of numpy or there's libraries like scikit-image that build on top of numpy and you can use matplotlib with all of this and really the cool thing is that you can put your data into a numpy array and then you can give it to psychic image and you can also give it to psychic learn and there's no reason why it would have to be that way now it could be that psychic learn has a data set class that you use and you put your data in there and psychic image has an image class and you've got to put your images there and matplotlib maybe has its own array thing so it's actually quite a you know it's not an accident people who built this ecosystem made it like this on purpose and it allows you to mix and match all these libraries and that's another reason i think constraining yourself to familiar interfaces which are supported by all these libraries is a good thing make it slightly more concrete you're a ice cream manufacturer or you run an ice dealer here in Darmstadt and in lots of other places and a fantastically sunny day like today. So and you're analyzing your data and maybe in your Darmstadt data you you know it's not it's not very big so you can do it all in your notebook you use pandas because it lives it's like tabular data use seaborn to plot it and everything is great because the data is is small it takes a few seconds to run your whole notebook and then you spot some weird trend towards the end of April 2025 and you wonder is that like just Darmstadt or is it in all my locations all across Germany now you load the data from all of those places and now running your notebook takes one minute or two minutes or five minutes and that really slows you down now if if have to wait for something to run for five minutes I'll go and read the news or look at cat images or ask chat GPT to make me cat images but the one thing that is certain when I come back to the notebook I've 80% forgotten what I was doing now so keeping the the cycle time really low is I think super important so some things to think about we want faster duration not trying to make work that is almost you know takes a few minutes making it interactive that is really useful taking something that takes 24 hours to run and making it run in 12 hours is also interesting but that's not what I'm interested in today familiar interfaces now it's nice because we can keep using the whole of PI Data Ecosystem and the PI Data Ecosystem is like this interconnected web of libraries. So how can we do this? There's two things I can think of that we can do and that I will unsurprisingly talk about today. So one way to make this potentially impossible thing possible is to modify the import system. So when you import pandas we actually give you something that isn't pandas at all, but something that pretends to be pandas. And then the other thing we can do, which is the more obvious thing to do, is you can start rewriting libraries to make them faster. So I'll talk about how to modify the import system first, and then I'll talk about work that's going on in scikit-learn to rewrite it. So the idea is, and this is like the small concession that we have to make on the left hand side you have some random Jupyter notebook cell that takes about a minute to run on my on my machine at home and then we want to make as few modifications to it as possible to make it run faster. So the small concession you have to make is to insert a cell at the very top of your notebook that loads a magic extension that does this modification to the import system and because this is something that actually exists and not just you know a random thought it's called what you have to put there is qml.xl because we would use a library that we created at nvidia which is called qml which is implements a lot of the algorithms which are in scikit-learn so in addition to the magic import system modification but you can see you don't have to change the rest of your code no that's quite nice so if you want to try it out take a picture now it has the URL you can also google for it to explain how you can mess around with the import system we have to think about a little bit how Python actually imports things so the code that you see is roughly for educational purposes what gets executed when you type import sklearn so we figure out the absolute name of the module that you're trying to import and then the first thing we do is we check if you've already imported it. That's the sys.modules brackets thing and after that we go through a list of finders to find the module and when we find one when one of the finder says it's found your module we use that finder to give us a loader that then loads the module and we store it in sys.modules and off we go so there's two opportunities here or two things we need to do if we want to mess around with the import system we need to write a finder and we need to write a loader what we will do is we will write a finder that when you import scikit-learn says yep I found it and then we load something which pretends to be scikit-learn but isn't actually scikit-learn so the module finder looks quite complicated but fairly straightforward the thing that we need to do is or what we will do, our strategy is to say yes we found it and then we will remove ourselves from the list of potential finders and then we will ask all the other finders where is this thing that we just claimed we found. And then we will get the original thing and we will modify how it's loaded and insert our own loader into that. So that's a nice way of making sure that when you install our magic we preserve the original behavior by basically cheating. and then the loader does basically the same thing we use the original loader that we found when we remove ourselves and after we ask it to do all the hard work then we modify the module so in this case for example at the very end we go through everything that's defined in the module and if there's a random forest class then we replace it with our like wrapper that tries to pretend to be the scikit-learn random forest but when possible accelerate things and this little wrapper class basically looks like this it holds on to a reference to the QML implementation and the scikit-learn implementation when you called fit we try and figure out should we accelerate this or not depending on the hyper parameters maybe on the data set size things like this if yes then we give the work to QML if no we give the work to scikit-learn and then when we call predict or when you access other attributes we forward it to the respective thing that we gave the actual work to also of course if you want to make something like this for talk it needs to fit on some slides if you want to do this I don't know professionally so that it works for lots and lots of users the code is unfortunately much longer than this but it's the basic idea of how this magic works and what it allows you to do is you put if you're not using a Jupyter notebook two lines at the top of your code and the rest of the code looks exactly like normal scikit-learn code. So how much do we actually accelerate? It's a subset of the algorithms that are in scikit-learn because scikit-learn has a lot of algorithms so we focus on the algorithms that can actually be sensibly accelerated on a GPU. Now there's no point in accelerating algorithms for which the GPU implementation is not 10 20 times faster. The nice thing is we also accelerate UMAP and HDBSCAN. So these are things which don't actually live in in scikit-learn. And yeah. Okay so final part rewrite all the libraries. Now this is the other thing we can do which can modify scikit-learn and make it faster. So rewrite all the libraries is too big a task, so I'll just talk about evolving scikit-learn. And if you've been around the PyData ecosystem for a while, you might know scikit-learn as a library that basically says GPU support no, never, over my dead body. And if you read the FAQ entry from a version from a few years ago, they have a more polite way of saying this but that's basically what it says and the reason for that is not some like irrational hate of GPUs but because installing GPU software drivers and so on like it has a lot of complexity and the one thing that CycleLearn is good at is making it very easy for people to install it. However recently we got around to updating our FAQ and now there's a longer answer and it's basically yes it's coming and it's been coming since 2023 so we've been quietly working on this for a few years and there is experimental support in scikit-learn now which it would be great if you try it out and give us feedback on and the familiar scikit-learn api looks like this we import a class we give it some data and the output of calling In this case transform is a numpy array and what you need to change if you want to make this work for example with PyTorch arrays is You need to set Set some scikit-learn configuration to enable this Experimental feature because we turn it off by default at the moment then you create PyTorch tensors Then everything else stays the same and the really cool thing is if you put your PyTorch tensors On the GPU, whether that's a GPU in your MacBook or a GPU that you bought from NVIDIA, it will use the GPU to do all the computations that PCA needs to do, and it will be significantly faster. And that's another way we can keep the familiar API and get speedups from your GPU. you. So there are many, many libraries that implement arrays. Now there's KooPy, there's PyTorch, there's NumPy, there's DAS, there's JAX, X-Array, the list goes on. And the problem with them is that they all have APIs which is very similar to NumPy, but subtly different. so that's why we have to do work in scikit-learn, that's why it doesn't just work and luckily for us there's the consortium for python data API standards and they are working on creating something called the array API and the array API is a subset of all the a subset of the numpy API that all the array implementing libraries can agree to. So there's some things that NumPy does that are extremely difficult to implement in JAX, for example, and there's also things which NumPy does which it just does because NumPy was invented a very long time ago and we now realize were a mistake. So there's something called the Array API, which is a subset of the NumPy API that all these libraries implement. And this is what allows us to modify Scikit-learn. And how does this look? So this is a function that somehow normalizes an array of NumPy values. It computes the mean, it computes the standard deviation, then it modifies the array and it returns it to you. If you rewrite it to use the array API, the biggest change is that we interrogate the input array and ask it what kind of array are you. Give us your namespace and then we use all the functions that live inside that namespace to operate on it. So in practice if you look at scikit-learn code, what you see a lot is that np. changes into xp. because the x is a placeholder. And there are some subtle differences and lots of little warts that need dealing with for example there are things which NumPy implements which are not in the array API then you need to figure out how do we deal with this but in general it's a fairly mechanical task so we've been slowly working on this and this is a list of almost all the things in scikit-learn which now support the array API in this experimental setting you can see it's a long list but unfortunately if you made a list of all the things which exist in scikit-learn it would be even longer so a thing for you to do if you're interested in this is to try it out but also if you are somehow interested in contributing to open source projects we would be very very happy if you came and helped convert some of these things and the reason the metrics here are so far ahead is that the metrics are quite easy to convert so everybody who comes to a sprint for example manages to convert a metric and then we've got them hooked and we try and convert an estimate and that turns out usually to be much harder Okay, so hopefully I could convince you that having standardized APIs which are familiar to us humans, because we've used them a lot, is a useful thing. Of course, to reach the ultimate performance, you often need a new API. For example, Poloz. There's things that Poloz does that are just impossible to do with the Pandas API, and makes total sense to invent something like Polars. The problem with using something like Polars is that libraries like Seaborn now don't work with your Polars data frame and then maybe you need to write your own visualization and things like that. However, it's nice that you can get pretty good performance. So often now we're talking about 10 or 20 times faster code, something that used to take a minute now takes a few seconds even if we restrict ourselves to to these familiar pi data apis and yeah both within rapids which is the team in nvidia that works on this there's efforts to do this for pandas and for scikit-learn and in scikit-learn there's also an effort to evolve scikit-learn itself so that it can use your GPU if you have one and that's the end thank you very much
Speaker 2 [20:29]
So, now we'll go ahead and I can ask the questions. There's also a pointer for the people that didn't yet. You can still ask your questions on Slido. We still have a bit of time. So, I will go ahead with the first question. Are there any tools to help convert from NumPy arrays to the array API standard?
Speaker 1 [20:51]
Not really the best tool is volunteers with a human brain. It's not it the diffs often look like it's quite mechanical, but For example when it comes to non not a number handling often you have to write your own function that implements whatever NumPy is non min or non mean or whatever does so So yeah, it's unfortunately not automatable at least, or maybe somebody like is good friends with an LLM and can try, like, I don't think anybody's tried that yet, but.
Speaker 2 [21:33]
I think the second question is more of a design question. So you showed a normalized function as an example for using the XP dot. Why is it not feasible to call methods of the supplied array, methods of the supplied array, instead of functions of the module, like array.mean instead of xpy.mean?
Speaker 1 [21:52]
I think it's mostly how the array API standard was designed, that the functions live in the namespace instead of on the array, which is how a lot of things in NumPy work now. So that's just a difference. I think it makes it easier because, of course, all the libraries like JAX and CUPAI and so on, they did not implement all of the array API standard on day one. So there's actually a library we use which is like a compatibility library. And it's much easier to put things in the namespace and modify them there, instead of modifying things that live on the array object itself.
Speaker 2 [22:42]
So when it comes to GPU standards, what do you personally prefer, PyTorch or QPy or even JAX for GPU arrays?
Speaker 1 [22:52]
I don't think I have a personal opinion. I think I use KuPi and PyTorch a lot. I don't use JAX a lot.
Speaker 2 [23:04]
So then we are also finished with all the questions except someone still wants to ask one on Slido. I will give 10 seconds Yeah, there's still one question about Typing arrays is complex and typing for the generalized array API even more complicated Should we set aside typing in numeric contexts for now?
Speaker 1 [23:38]
I don't have an easy answer. Like, I'm so old and such a long-time Python user that I often still wonder, like, why do we need typing? Which probably tells you what my answer to the complicated problem of how to do typing with numerical arrays is. Like, I'm like, if you want to spend your time working on this, like, be my guest. I'll go and do something else. It's difficult, yeah.
Speaker 2 [24:05]
okay okay but then yeah i want to thank you for your talk tim
Speaker 1 [24:05]
Okay.
Speaker 2 [24:09]
thank you and thank you for answering the questions yeah thank you for this nice topic
Speaker 1 [24:15]
Thank you.