Squirrel - Efficient Data Loading for Large-Scale Deep Learning

Speaker 1 [00:03]

Thank you very much for the kind introduction. Good morning, everyone. We are very excited to present you Skrill today, a Python library that we just open sourced for efficient data loading for large-scale deep learning. And today, I will tell you how to use that library to load data fast from local and block storage for every deep learning framework and share it through a data mesh. How awesome is that? But first of all, I want to tell you a little bit about Merantics Momentum, because this is important to understand why we actually built Squirrel, the solution. I, myself, leading the engineers at Merantics Momentum. And we cope with lots of different AI use cases on a day-by-day basis. So we take the journey with our clients from various industries, from finding AI use cases with them, researching them, then also building them into an AI solution and bringing them into production and operating them. So end-to-end to full chain. And our focus is a lot on deep learning models, and we try to leverage the latest research actually in production systems, which is very challenging. And we do this across various industries and with very diverse data. So you can see NLP data here, computer vision data, structured data, you name it. And also different ML tasks, so from detection, classification, learning some embeddings, and in information retrieval, like the whole thing we do for our clients. And we built a very flexible stack to cope with all these data and ML challenges in our projects. And since we are building very state-of-the-art deep learning models, usually these models are large, they need a lot of data, so you have to feed these beasts of computation graphs with huge amounts of data very efficiently. So in a nutshell, you get some raw data, either from clients or from some public data sets. You prepare it so it's nice, juicy, and tasty for your model. And then you feed it to your model and let it train so it learns the task. And Squirrel is about the later step. So taking the training data and feeding it to your model training at maximum efficiency. So if you're doing this naively, and you might started training a new machine learning model on the deep learning framework of your choice in the cloud, you might experience some long training process. And maybe your cloud bill goes up pretty quickly. And you pop up your NVIDIA SME and have a look what your GPU is doing, what the training is doing, and you see a behavior like this, so your GPU utilization jumps around quite a bit. You cannot make use of the GPU hardware all the time. This should be the ultimate goal to keep that GPU busy because you are paying a lot of money to basically rent these machines in the cloud or even use them on-prem. So what might happen is you have these three very abstract steps of first loading, so loading it from a disk or downloading your data from a cloud, a specific sample, then preparing it for the model training, and then copying it over to the GPU memory. And if you are not investigating some parallelization, you would do this sequentially. But in an optimal scenario, you want to pack this together in an asynchronous way so the GPU is always busy. So let's imagine you have a use case from medical imaging. So this is histopathology images, a very standard HE-stained image. It's fairly large, so these images are like 200K by 200K pixels, so not your average, I don't know, NIST image that you have in your research project, but a fairly large image. And what you usually do is you take a sliding window, and to each of these window slides, you apply a neural network, which then predicts the class that you want, for example, background, or tumor, or non-tumor. And if you want to train such a model, you have to feed it one of these windows, and then it has to predict the task. So if you are doing this naively, you would slice up the image, and you would get billions and billions of small files or these windows. If you store them locally, good luck opening this folder in your file browser. Or if you're storing it in the cloud and you're trying to download it, for example, from a cloud bucket, then you already pay lots of money just for downloading all these small files because you pay per file downloaded. If you store them in charts or just leave them in one large image or in a format like ZAR, you have the issue that you might want to shuffle your data. But if you just take windows slide by slide by moving the window over the image, you have very correlated images. And since the image is so huge, you will not get a lot of randomness in here. Moreover, if you get new data, what do you do basically interleaving the new and old data set? Do you want to copy it? Lots of questions that you are facing, which leads to engineering efforts that you have to add to your project. And these challenges, they reoccur. and now your model is trained, and it might not work that well. So here you see this red marker, and usually models don't work that well in these regions because someone painted on the actual image and removed some of the important features. So you might want to debug the model and retrieve the specific patches that lead to model failure. But if you're using some data formats or slice everything up, then you have to find this particular crop again. And for that, you need metadata. Now the question is, how do you store the metadata? How do you link all this together? You can do this like in a data frame, store it to a disk or like a CSV file to a database. But all of this is custom and gets hacky pretty quickly, and you have to build that yourself. So lots of questions arise and there are different frameworks out there which are all great for their purpose and that try to ease parts of this process. So you have, for example, TensorFlow data or PyTorch data to load your data into your model basically when using these frameworks, but the speed and also the capabilities are very limited. And then there are frameworks that try to solve these challenges by either providing some additional abstractions or providing their own data format, like Active Loop or Hub from Active Loop web data sets or Hugging Face or FFCV. But then you have this challenge, you might want to share your data between your members of your project or between teams, and there are solutions by hugging phase or intake for that. If you want to pre-compute features, like Feast or Hopsworks, or if you want to have faster data transforms, like faster data augmentation, then you might want to use NVIDIA DALI. But all of these frameworks do one particular thing, and they do it very good. But if you have to cope with all these different ML use cases, cases, then using these frameworks together is pretty cumbersome, and they lack certain features that we would like to have in our projects. So our journey of actually finding the right framework was pretty intense. So we started with using raw files. This was obviously very slow, especially if you have the use case that I showed before. And the reusability of the data loading code was super low to non-existent. Then we switched to tf.records. But shortly later, we also switched to PyTorch. And PyTorch and tf.records, they work together, but it's not a very pleasuring experience. And since you don't have an index, this debugging of the ML models also gets very cumbersome. Then we move to ZAR. So ZAR is an excellent framework for storing tensors in chunks. But this leads to many, many files in the cloud, many connections and cloud costs. And we try to optimize the ZAR library to make this more efficient, but the engineering efforts, they just escalated. So this was not a fit for us as well. And then we tried out some more narrow libraries. We looked at web data sets, but it uses native file formats by default, and this brings its own challenges, and you have the painful ML debugging again. And then we tried Hub, good for images, but didn't work so well with text. Then we tried Hugging Face, which uses PyError under the hood, not so good for images, but for text. Then FFCV, again, good for images, not so good for text. We have large datasets, so we can't store them just on the disk. This is just not provided by FFCV. There was no solution for us until now. Also, we collected the key requirements that we need from our past experience. And we came up that we want something which is very efficient and doesn't produce a lot of cost just by downloading the data. We want to keep the GPU always utilized and load only the stuff from the remote place that is actually needed for training and nothing more. It should be easy to use, and we don't want this issue anymore of having unnecessarily connections. Moreover, we have these very different use cases and different data modalities, so the solution should be flexible. And we learned it the hard way. There's not one format that solves everything. So we need something which is flexible, can cope with different formats, and also easily can handle new formats if they arise. And multimodality should be a first citizen. So if you have a sample here, which is illustrated by dictionary, you want to have an image next to text, to metadata, but also your labels here, and don't need very specific data structures to store these different kinds of data. Moreover, the solution should not replace frameworks that are already great. It should integrate with these frameworks and make them even better. And it should be also very good for rapid prototyping and not be a beast that you have to get trained on for 10 years. And since we have lots of small teams that collaborate a lot, we also thought about data sharing between these teams. So if somebody prepared a data set, curated it, this person might want to share it with another team. And if you do this with a central hub, then you have a single point of failure. You have to invest in an ops team that runs this platform. And if it doesn't work, your company might stop. So we wanted to go for a centralized system and identified that there should not be a central point of failure. It should be very lightweight to operate the system. And since we want to have it decentralized, people should also operate their part of the infrastructure themselves in a self-serve fashion. And also, the agreement on the standards between the different teams should be also governed in a distributed way. So we didn't found a solution for us, not a single one and also not a combination of different frameworks that solved this for us. So we built it on our own. And here's Squirrel. It's now open source, and it's a library to share, load, and transform data in a collaborative, flexible, and efficient way. So to understand what this Squirrel is, Squirrel contains primarily three abstractions. So the foundation of everything are the iter streams. Iter streams, in a nutshell, are just chainable iterators with a neat API to do so. Then the next layer are drivers. Drivers let you read and write different data formats from local and remote places and implement different interfaces on them, for example, an IterStream interface. So you can load data from one specific data format and then get a standardized IterStream interface back. And now the sherry on top is the catalog. So the catalog is a dictionary-like data structure that maintains all the different data sources that you'll subscribe to in the data mesh and lets you access all the different instances of drivers and from there also the IterStreams. How does this now look in code? Let's start with the iter streams. Here, we have this list, one, two, three, and we put it into an iter Rebelsource object. Now, you have this very convenient dot-like API where you can chain with two, basically a callback onto it, which logs these values, in our case, to an MLflow server. Now, the next step would be adding plus one, like in a map function with this lambda that is here defined. The next step would be an async map. The same, you square the value, but now this is done in parallel, in different threads, basically. And the last step is filtering your pipeline by everything which is dividable by two. And now you can use the result and basically ingest from that. In our case here, a very simple example in a for loop, but imagine now you can pipe it to your deep learning framework. And these chainable iterators are very flexible. So now you have seen, first I might lock my throughput to an experiment tracking server, but then I do some transformations of my data. And here in the example that you see on the left, the parallelization is done in threads. But you can also leverage processes or DALI if you want to go to the GPU. Or if you have super crazy data augmentations, imagine you have 3D elastic deformations that you want to do online while training, you can offload this to a DAS cluster. And then you can chain other functionality on top of it, which have some short ends like filtering, batching, and all the standard stuff that you usually need. But be careful. This is a very powerful API, and the parameters, they matter. So here's an example. So you have an iterable source with 5,000 values. You shuffle these in an in-memory buffer, and then ingest from it. And if you look at the correlation of the values in the stream that you get back, then you can see that, especially in the beginning, the values are very correlated, but why is that the case? Yeah, you set the wrong parameters. Your initial buffer is too low, so you get some correlation. So be careful if you deviate from the standard parameters. It's powerful, but you should know what you're doing. Let's now move to the second part of Skrill, which are the drivers. So drivers let you load and write from different data formats, locally and from remote. Our preferred data format is MessagePack. It's very small and very fast to deserialize, which is pretty cool. And we use FS-Pack under the hood to enable loading from local and remote places. So here you can see three examples. Here in the first example, we take the message pack driver and from a local folder, we basically create an iter stream. And now we can use it like on the previous slide. But you can do the same with JSON-L. And since we are having FS-back, you can do this also from a remote place. But there it doesn't end. We also have integrations with other frameworks. For example, with Hugging Face, we can directly ingest from a Hugging Face data set, but on steroids. But what happens now if you have a data format that is not supported by scroll yet? We are using the amazing plug-in framework to let you extend scroll. So let me show you an example. Imagine you have a graph database. This is now very mocked up. And you implement a class which inherits from the iter driver interface. You give it a name so it can be later also addressed. And you implement this getIter method, which instantiates an iterable source, and in our case, takes random walks from our graph and returns this in an iterator. And now you're already done. To make this usable in Squirrel, you can either register it directly via the plug-in manager, or if you want to distribute your driver, you can register it in a classic pluggy fashion via your setupPy entry points. Now, people ask us all the time, yeah, this is all nice looking, but how fast is it actually? We did a small benchmark to show you, is this actually fast enough? Let's start with images. FFCV, for example, is a library that tries to be very fast at loading images from a local place, But there are also other libraries that try to solve similar things. Here, we are loading a very standard dataset, CIFAR-100, onto the GPU, and you can see that the samples that you have on your GPU per second with Squirrel are more or less the same as with FFCV. But now, what happens if you want to go remote, load from a remote GCP, or S3 bucket, or a database, whatever. This is not possible with FFCV, but with some other frameworks. But Skrull is much faster here. But what is now about text? Because we want to be able to load different kinds of modalities. So let's now load Wikitext 103. Again, locally, scroll really shines here. If you go remote, the same. You get really high speed for different modalities, and it does not really matter if you load locally or remote. In both cases, it is very fast. But what if this is not fast enough for you? You might have very expensive data transformations, so you can leverage different frameworks to make the transformations faster. You can use Torch Vision pipelines or NumPy. You can JIT compile your transformation, similar to what also FFCV does, or if you want to go full in, you can offload your computation to DALI and let it transform on the GPU, which is even more faster. So here you see an example of the very simple pipeline and the performance on different frameworks. Let me now come to the cherry on the top, the catalog, or the data sharing functionality. So here we figured out the concept that we actually want is a data mesh. So a team has a prepared data set, and it wants to offer this data set as a product to other teams. And you can think of it, they get kind of a capsule, and the interfaces in these capsules are standardized. It's more like a contract, how to load data and how to serve data to others. But in this capsule, you can do whatever you want. This is basically the product that you're building. your team IP, more or less. For example, transforming your data and then serving it to others. However, this is more like a concept, more like an organizational philosophy which was populated by ThoughtWorks and McKinsey and others. But there's no real opinionated open source implementation out there for machine learning that helps you actually realizing this mesh, but more like standard solutions that people try to plug together to bring up this vision. But there's an implementation now, which is Squirrel, and let me guide you through it, what it actually does. You can take this catalog object and retrieve a dictionary-like structure from it, from all the data sources that you already subscribe to. And here, you can, for example, access ImageNet if somebody else hosts it, and get the ITER stream directly from this ImageNet data set, and start training. And we thought that there should be three different ways how to share data. The first one is within a project. So you have a code base. So why not use the Python API? So you instantiate a new catalog object, which is empty. You describe your source. For example, which drivers do I want to use, what are basically the arguments, from where to load it, and some metadata, and design it to some ID with some version, so these sources are versionized. And now we implemented a set API on top of the catalog, so you can do different joins of catalogs. For example, here in this example, the default catalog that you're subscribed to and your custom catalog in your project, and you get a unified catalog to access all these data sets. If you want to share data between projects, we suggest to do this via a Python package. So we are already in the Python ecosystem. So why not use the deployment mechanism of Python itself? So for that, you have to implement a plug-in hook, like here in this example, which lists you all the sources that you want to distribute. And now you add the specific hook to your setup pile, like you always do with Pluggy, and distribute the whole thing as a Python package. And if somebody installs your Python package, for example, via pip install, then they already get all the data sources that you provided via your package. If you have a continuous delivery pipeline For machine learning, you might not want to generate Python code and pass it between steps. That's why you can also serialize your catalog as a YAML file and then scan specific locations for these YAML files and create this catalog object again from it. And this is especially interesting if you have a whole machine learning pipeline end-to-end, and you created a data set in one step and want to pass it to the next step, for example, and have this whole process documented, more or less, and versioned. So now you've seen these three API levels. And let me show you an end-to-end example so you see what is the full thing. Let's train the beast on ImageNet. The first step is to get your catalog from your standard subscriptions, access ImageNet, and get the driver object. This is loaded in a lazy loading fashion. Now you want to get an iter stream from it, and this driver defined several splits. we want to get the training split. And since we are training in a distributed way, you might have different nodes for training. But on each node, there might also be different workers that do the training in parallel. So instead of downloading the whole data set on all nodes and all workers, we provide hooks to split the data set already before downloading it to the specific machine. After comes the shuffling of the samples. So imagine you have now a stream of samples, and this stream gets now shuffled in a window of 100 samples. Next is applying some runtime transforms. So here we transform the image and return a tuple, basically. Finally, we want to pre-fetch this whole stream in an asynchronous way, and then batch it together in batches of 64. And now you can continue with what you do all the time. So here's some pseudo code. It might look different for the deep learning framework of your choice, but you have some trainer, and you can pass your just iterator to it and train your model. And this is basically it. This is an end-to-end example of how to use Squirrel in your training code. So why should you use Squirrel now? So first of all, we think the API is just fun to use, which is a great asset for us. Moreover, you have a lot of functionality to prevent GPU stalling. You have prefetching, you have this asynchronous transforms, you have these key hooks, and more, everything to keep your machine busy and get everything out of it. Moreover, you don't have to trash what you already built, but you can continue using your preferred frameworks, but now on steroids. Finally, we provide one of the first implementations of the data mesh concept for deep learning and are very excited for your feedback. So now it's your turn. If I sparked your interest, you can find Squirrel on PyPy and also on Conda. We provide a data mesh node with some Squirrel data sets already. Also, give us some love on GitHub. And if you need any support, then join our Slack channel and meet our amazing Scrooge developer team. And I guess now you have some questions.

Speaker 2 [33:03]

Thank you. Yes, you have lots of questions. Why is it much faster than other frameworks? What is the trick under the hood?

Speaker 1 [33:14]

There are many tricks I would say One one Interesting thing is message pack is quite fast and we do a lot of things in an asynchronous fashion Which makes it even faster?

Speaker 2 [33:33]

Thank you Loading data requires a combination of CPU and IO bounded tasks the former blocks of the GIL How does a squirrel mitigate this?

Speaker 1 [33:44]

The user can decide. So you can separate the CPU and IO-bound tasks so that you can utilize both in the best way and you are very freely to decide whatever fits your specific use case and if that is not enough then just Just offload it on a DAS cluster and just put a lot of compute on it and then it should work. Fingers crossed.

Speaker 2 [34:18]

Nice talk, thank you. When showing the benchmarks, I was wondering how you made sure that for the other frameworks you used the most efficient implementation.

Speaker 1 [34:29]

We didn't. We used the default parameters of the frameworks, so we didn't do hyperparameter optimizing for all these different frameworks. But since these are very preliminary benchmarks, I wouldn't put too much interpretation into it. We are right now preparing a more elaborate benchmark, so you can see all the different knobs that you can turn to make it even faster.

Speaker 2 [34:58]

Just out of curiosity, how did you come up with the name Squirrel?

Speaker 1 [35:03]

Very good question. So all of our software solutions at Mirantics Momentum have animal names. So we have Squirrel, but we also have, for example, Chameleon, which is our computer vision solution, or we have Parrot, which is our NLP solution, and so on. So we have different animal names, and Squirrel in particular, there are several reasons. But one punchline is it's a data infrastructure library, which doesn't drive your nuts

Speaker 2 [35:40]

Do you use a job queuing system or what is your approach to fully utilize all GPUs without idle time?

Speaker 1 [35:50]

I think I digged a little bit into it. You can use this splitting of CPU and IO-bound tasks. You can use a DAS cluster, multi-processing, multi-threading, depending on the individual step. And yeah, that's basically it. If you want more, go DALI. You can combine all these frameworks within Squirrel able to get the most performance out of your setup. And since we have very different ML use cases and data that we are coping, we also have different setups for different tasks. But since the API is very easy to use, people love putting together their pipelines for their specific ML problems.

Speaker 2 [36:40]

Thank you. Does Squirrel support splitting streams, for example, for n cross-field validation?

Speaker 1 [36:48]

Not by itself, but if you want, you can use this to method and plug in whatever callback that you want.

Speaker 2 [37:01]

Would you agree that most aspects of Squirrel are not limited to deep learning but might make sense also for other data processing flows?

Speaker 1 [37:10]

also a good question so we have heard this already a few times so we built this for actually deep learning training but of course you have this process or this problem or challenge of getting data onto your machine in a very fast way in also different scenarios and it's also possible to use scroll there because it's data or yeah it's a framework agnostic but this was not a intentional basically design but would be happy to see how you use it in in other creative use cases um yeah that would be cool to learn

Speaker 2 [37:53]

Thank you. Does it also work with multiple GPUs?

Speaker 1 [37:57]

Yes, of course. That is why we have these hooks, for example. If you want to distribute your training, then you can use the hooks to feed the samples that are, for example, for one specific GPU, just to this GPU. You do not have to download it for each GPU and then discard most of the data in memory.

Speaker 2 [38:28]

Thank you. Next question. Could you describe what was the issue with PyArrow? What was slow specifically? Did you reach the dev team?

Speaker 1 [38:39]

PyArrow itself is more like a columnar format. If you are not just reading one of these columns, but want to, in a very fast way, read rows of it, it does not provide you with lots of functionality here. If you have tensors, like images, also the deserialization was not that fast, as far as I remember.

Speaker 2 [39:10]

I think the next one is kind of a duplicate, but what exactly explains the higher throughput observed with Squirrel?

Speaker 1 [39:19]

Yeah, I think that's a duplicate.

Speaker 2 [39:20]

I already answered that. Have you been able to measure effective increased GPU usage through parallelization with ether streams?

Speaker 1 [39:30]

To be honest, right now these ITER streams are so fast that we can load much more data than we can actually process on the GPU, so it's just the maximum.

Speaker 2 [39:47]

Do you plan on evaluating bigger data sets?

Speaker 1 [39:51]

Yes, so this was just a very preliminary evaluation just to show you something, but a larger evaluation is in the pipeline.

Speaker 2 [40:02]

Someone wants to know if you're going to share the slides?

Speaker 1 [40:07]

No, but I think the talk will be online, right?

Speaker 2 [40:13]

The last question is, can you define a own way how to split the data set easily? I'm not sure I understand it, actually.

Speaker 1 [40:22]

I think it's again a duplicate. You can you can provide these key hooks to do that

Speaker 2 [40:29]

Thank you very much. Please, let's thank...