Lightning Talk: Jigsaw: Domain and Tensor Parallelism for High-Resolution Inp... Deifilia Kieckhefen

This talk presents Jigsaw, a domain and tensor parallel technique aimed at optimizing the training of large models with high-resolution input images. The speaker, a doctoral researcher, highlights the challenges posed by GPU memory limitations, particularly when dealing with large data samples, such as those found in image processing and scientific analysis. Jigsaw effectively circumvents these limitations by using domain parallelism to shard data samples and intermediate activations and tensor parallelism to distribute model weights. By implementing distributed matrix multiplication and maintaining a user-friendly interface similar to PyTorch, Jigsaw allows developers to train significantly larger models efficiently. The speaker shares insights into practical implementations, benchmarks, and the framework's successful integration with existing parallelism methods. Jigsaw, still in development, shows promise for enhancing memory usage and training speed in machine learning applications.

Hi everyone, thanks for coming. Uh my name is Defilia. I am a doctoral researcher, PhD student at the Cosgrow Institute of Technology. Um and here or today I'm here to present Jigsaw which is a domain and tensor uh parallel technique which is specifically designed for a high high resolution input images into your models. So basically we want to distribute both the data and the weights and be

able to do this in order to train models where the input data samples are um on the order of gigabytes large. And the reason we want to do this is because complex tasks de uh complex tasks demand large models. We all know that state-of-the-art models for LLMs for instance have on the order of hundreds or on the order of uh tens of billions to trillions of parameters.

And the simple reason for why we need this is because we found empirically bigger model means better predictive performance which means that it is advantageous if we train a model that is as large as we can possibly fit onto our GPUs. Of course, this leads us to the very critical problem of how do we actually train such models? Because again, if we want a model, we have

to train them on modern hardware. For now, GPUs, TPUs, you heard the previous talk. But these have limited memory. We're talking about about 80 to 100 GB of memory per GPU. And that's if you're lucky enough to get your hands on such great computational resources. The thing is when we're training models, our GPU memory has to hold a lot of different pieces of information at once. It

has to hold of course the model weights and the data samples that you're training on. It also has to hold the optimizer state and what gets um easily forgotten are all of the intermediate activations. So these are the pieces of information you save between the forward and the backward path to make training faster. Of course, you have approaches such as FSTP0 or pipeline parallelism to offload some

of this computation partially to the CPU um or to allow you to train larger models by chaining together or having several GPUs work together. Um but this is still basically a fundamental limitation. So the GPU memory limits how big our models can be. And of course you have different uh different types of parallelism that you can use to circumvent such limitations. First one is of course data

parallelism. This is standard has to be used has to be done. In data parallelism you basically take the entire data set that you have subdivide it into disjoint subsets and then each GPU you have will then train its local copy of the model on the local subset of the data. You perform some overduces to make sure that the gradients stay synchronized. And thankfully PyTorch has the DDP

wrapper which performs this very efficiently. So solved problem we have to use it. What we can also do is tensor parallelism. So techniques such as Megatron LM for instance uses this. And in tensor parallelism we also sub we also split or shard the model weights across several processes. This means that now when you look at GPUs one and two they each hold half of the model weights.

Thus of course allowing you to train a larger model. But I'd like you to note one thing is that in in this schematic in this form of tensor parallelism, both of these two GPUs still need to load the same data item because its weights need to see the entire data sample of course. But what happens if the data that we're looking at is large? So I'm talking

about applications in image processing, scientific analysis, uh met uh satellite image classification or what I personally work in uh weather forecasting. In these cases, it's quite easy for a single data sample to be gigabytes large. Let's say 2 gigabytes. That's already a significant fra fraction of the 80 gigabytes of memory that you have. And the thing is, if you have a large data sample, if you really

count the number of pixels you have, you could have a sequence length that's on the order of millions. What this means is that the data samples and the intermediate activations that you have to hold in your GPUs explode. Practically what this means is that you leave very little room for your model weights and optimizer states. And that's why when we train models for imagebased applications on the

same hardware as if we were to train an LLM, we can train we can only train models that are relatively small. So I'm talking 1 billion parameter versus an LLM that could train on the same hardware with tens of billions of parameters. So of course we wanted to uh implement or develop a strategy to circumvent this. And with this and what Jigsaw contributes is domain parallelism. So

now what you see here are these two pairs of GPUs. So one and two which are now tensor parallel. They each hold half of the model weights but they're also domain parallel in that they only need to hold or reload from storage. Uh half of the individual data sample. This could be half of the spatial image you have or half of the channels. Either works. And what's

really nice about this is that you really have now zero memory redundancy across the two processes. So you have two GPUs that need to load completely disjoint uh fractions of the data and hold completely different disjoint subsets of the weights which allows for very efficient memory usage. So that's what uh Jigsaw is. Uh Jigsaw is a framework that is designed to be very user-friendly, build upon PyTorch,

be very PyTorch like and allows uh developers to train models with both domain and tensor parallelism. What we now have is uh through tensor parallelism, we shard the weights and the optimizer states. Through domain parallelism, we shard the data samples and the intermediate activations, which effectively means if you have two GPUs, you have twice the amount of memory to work with. you can really train much larger

models at all points. We avoid doing any types of all gathers to actually collect everything on one process. So the model can really be kind of twice as big as the original version. Um and how did we do this you may ask? Well, we looked at the fundamental component of the neural network and I would argue that that is the linear layer. Um everything's based on the

linear layer. Uh thankfully the linear layer is very easy. It is a matrix multiplication. you multiply some data X with some weights W. Okay, done. GPUs are very good at this. Um the thing is now when we want to apply domain parallelism, we now shard the data X over here across multiple processes and with tensor parallelism, we shard the weights W over multiple processes. And now we've

broken down the problem into a distributed matrix multiplication problem. In the world of distributed computation, this is a solved problem. Yay us. Um so what we were able to do is now implement uh distributed matrix multiplication kernels very simply. Um so they're very simple algorithms all in Python implemented via PyTorch distributed and of course we tried to overlap as much communication and computation as much as possible.

We use these distributed kernels to overwrite the autograd functions within uh to allow for the forward backward passes of a linear layer. And this now allows us to build distributed MLPS because all an MLP is are two distributed linear layers. And apart from the difference in the initialization, you'll notice that the forward pass looks exactly the same as vanilla PyTorch, which is very convenient. You might think,

okay, yeah, who cares about NMLP? Um, but we also realize that we can do this with attention. Uh, so attention, uh, I have the equation for attention up here. Um for those of that you for those of you that don't know um you have a softmax function and then you uh multiply a series of matrices together. Q times the transpose of K uh divided by a scaler

softmax that times a values matrix. So now we have a couple of distributed components or components that we would have to consider in the distributed setting. The first is a QKB mapping which is a linear layer. We have this the second we do have to implement a distributed softmax can be done. Um, and the rest of it, you'll notice in the forward pass, it's essentially just a

series of distributed matrix multiplications again, which again allows us to have a an interface that's very similar to vanilla pietorrch. Um, of course, we tried this uh it's a simple graph, but Jigsaw, we really found that Jigsaw allows for faster model training. So, in our case, we have until now uh kept the model parallelism up to within one node um because that's where the fastest communication happens.

So we par we basically trained a weather model with an with input data with about 400 megabytes per sample. So fairly big. And then we partitioned this data and trained as large of a model as we can fit on up to eight GPUs. On one GPU this was about a billion parameters. On 8 GPUs this was about 4 billion. Um and we found that this scales extremely

well because by having more GPUs then we can basically make the problem size bigger which means the GPUs will uh will over will communicate will compute things very efficiently and effectively. Um and all this goes to show is that we found this to be a very effective method um scaling method in order to train larger models. Of course, we also tried this in combination with other forms

of parallelism such as mentioned earlier with uh data parallelism and we have tried this on a fairly in a fairly robust manner on several different HPC systems with training models with the different layers of parallelism on up to hundreds or even close to thousands of GPUs and found this works very robustly. Um so here's a summary. Uh Jigsaw is a framework that allows you to stay very

close to vanilla pietorrch um uh syntax and it allows for distributed I discuss here training because it's a bit more interesting. It would also work with inference and the thing about Jigsaw is that it builds very cleanly on PyTorch such that it still takes advantage of PyTorch's auto differentiation and the CUDA back end for very efficient computations. Um it is a parallelism scheme that allows for for

highly efficient memory usage and we have tried this. It has seamless integration with other forms of parallelism such as DDP pipeline um FSTP and so on. Uh this is a work in progress. This project has been in development for about a year and a half now grown quickly. Um so feedback is welcome and I'm happy to talk with anyone uh who has questions. Thank you.