DeepSpeed – Efficient Training Scalability for Deep Learning Models - Olatunji Ruwase, SnowFlake

Key Takeaways

Thanks Matt. Hi everyone. Uh today I'm going to talk about deep speed. Uh it's um a library for efficient deep learning and scalability. Um this is a swore from a lot of contributors. So uh yeah indeed when we created it we the goal was to create this open source library um for distributed uh uh training uh because we recognized early on that uh this will be important and the goals that we have then was uh basically to make uh distrib training easy um because model scientists are the ones who will primarily be doing this um wanted to make it efficient so as you scale uh it's very expensive from a compute and resource perspective. So we wanted to ensure that the workloads were running as efficiently as possible. So you're getting you know the biggest bang for your investments and wanted to make it effective as well. So this is more from a data efficiency perspective and you know on on those points so like from efficiency perspective we've seen that you can run with high performance whether you're running a single GPU or you're running you know over 100 thousand GPUs. Um we think the adoption really has been driven by the fact that it's easy actually easy to sort of get running out of the box once you install it and democratization was also something that we also kind of highly value which is basically you know some uh the cost of uh state-of-the-art models can be quite because of the hardware investment. So we also work on techniques that help to lower this uh the the budget requirements for you to be able to use state-of-the-art AI. Um, so it's on GitHub and that's just the screenshot of our page. So why did we create speed? It was based on the observation that scalability is what drives deep learning, right? The the industry is driven by scalability. And here on this plot, I'm highlighting two dimensions and there are others, but I think the two primary dimensions of scalability. On the y- axis you have the model size um which is basically you know in the in in the billions of parameters going up and on the x- axis you have the tokens that you use to train your model right so in the billions of tokens and so if you can sort of point if I point your attention to down there on the lower uh uh left that's we could think of that as being the ground a zero for you know the current wave of deep learning right the model then was just about 300 million parameters trained on um um I think some hundreds hundreds of billions of tokens um at that time. And if you can I can then point you to the top right hand corner which is sort of where we are today. You know lava 4 we just talked about earlier um you know there we're we're basically at hundreds of billions of parameters model train and then you know trillions of tokens um uh tens of trillions of tokens training. So that's basically how things are scaled. Um and so because of all of the scalability challenges basically they uh create a lot of systems to put a lot of pressure on the system resources. So on the compute side, on the memory side, on the communication side, on IO side, on data and and and so in other words, what what I'm showing here is that um the advancements we've had in deep learning have really been held up by innovations along this this four pillars that I'm showing and and in particular, deep speed. At deep speed, we've been working on all of these areas and creating solutions along the way. And um looking at the models up there, I'm sure you recognize most if not all of them. The ones I've highlighted are the models that have been trained using Deep Speed. And so even though the project started out at Microsoft, some of these models are external models, not necessarily Microsoft only. And and this is because Deep Speed has been open source from day zero, right? So we've had this opportunity to collaborate with the community and and together we've been able to build the solutions and they've been able to build the models. Um so I I'll talk about the real key challenge that we tackled in the early days and sort of what kind of I think put speed on the map is like the GPU memory wall problem and and this is shown here where here on the y-axis I'm showing you know scaling in terms of the model size. Um this plot is very popular. It's from UC Berkeley and and X-axis is showing time and it's basically comparing the size of the model to what the acceleration HBM memory that's available at that time. And and you can see there there's that point where it just sort of crosses over and then your model cannot fit on a single GPU. I think that was roughly around the 2020ish Megatron LN time. And because the models were growing at about 100x every year, the HBM is kind of like, you know, just growing 2x every two years. And even if we project forward till today and the projections that show up there with like the B300, we're still just going to end up with about 280 GB of HPM. So it's still, you know, not not sufficiently large enough. And so we sort of this is the first problem that we tackled. And we did that using two key solutions that we just sort of summarize here. The first one is zero and zero is this optimization that we designed for the data parallel training just uh where data parallel training paradigm where traditionally data parallelism has the model states replicated on all the ranks and with zero we uh came up with this idea of instead of uh replicating why don't we uh partition and you know you duplicate across and so then that helps us to reduce the memory footprint per GPU and allows us to fit in much larger models. Um then the other direction we went in was like well you know your servers come with other iterogenous memory like your DRAM and your local NVME. Uh we could also leverage that memory as well. And those are typically like orders of magnitude bigger than your HVM. And so the zero comes in like three different flavors uh you know for the various kind of partitioning schemes that are in there. Um and that helps like you know fit extremely large models um onto onto your GPU. 3D parallelism kind of like a subset of what uh previous uh speaker spoke about where you're able to use auto parallelism strategies to sort of strategic your model state. So tens of parallel pipeline parallel and then you know it's zero in parallel but the the main takeaway here is that plot over there where it's sort of still building on that GPU memory wall plot I showed earlier but now it's showing what deep speed when you bring in deep speed systems capabilities like if you see that there and you see that deep speed actually is scaling even faster than the models are growing so you know roughly like uh just looking at the impact we provided system capability ability to train truly empire models as far back as 2021. So this is well ahead of the model sciences um and and this the techniques which I've listed here have inspired other frameworks like FSDP and closely AI meatron who've all sort of adopted the zero uh optimizations. So um I want to spend the rest of the time to just talk about some of the new things we've been looking at in dspe uh you know uh and and there are basically four things that I'll highlight today. So the first one is universal checkpoint. Um and the motivation for this is that now uh you know distributed training states are distributed whether you're doing pipeline parallel tensor parallel or zero. And the problem that comes from that is that when you create your model checkpoints they're basically uh coupled with your hardware uh that you created them on. So in other words you cannot resume training unless you have exactly the same number of GPUs. And that's just because the distributed checkpoint is tightly coupled with your hardware configuration. Um, and so this creates a lot of like inflexibility challenges with resource uh capacity management. And so the solution here is that we decouple those two things. So your checkpoint and your hardware parallelism or or your even your parallelism strategies like tensor parallel or pipeline parallel, we decouple those two things and that brings some practical benefits. One is for fall tolerance, right? So if you have a failure, you can continue training with the healthy nodes. you don't have to kind of wait until you get, you know, get a replacement or you can continue training. So that's really great. And the other is faster training. So basically, if you have elastic resources become available, you can take advantage of them to speed up your training just opportunistically um once you have sort of give them up and this work will actually appear at ATC later on this year. Um and so the overview here, I'm just going to just touch on them very high level. Glad to chat later on. So the first thing is we recognize that the different parallelism strategies exist and we provide a way for you to reconfigure that. So you started out with some kind of TP BP whatever uh we provide a mechanism that allows you to reshape your parallelism strategy. Um so you can go from any source strategy to any target strategy. Um and then uh and this reconfiguration is pattern based which means that it makes it portable. We provide a language where you can express your parallelism strategy and so you can rewrite you know with that you can rewrite what you want your target uh strategy to be. So it makes it easy to adapt and apply to other uh in in various scenarios and then we've also provided optimizations for this reconfiguration. Um so here I just highlight some uh highle results. uh one is the reshaping and this is just to show that you know when you do reshape your strategy your training is not affected your convergence kind of remains the same and so the plot up on the on the left hand side is showing that you you started out training with one particular strategy and now you want to reshape to multiple different strategies and that's kind of just showing that you know the loss curve is the same. The next plot to it is kind of showing the inverse where you've started out with you know different strategies and we're showing that you can go into a new strategy without you know uh suffering of your uh in your convergence and then the lower uh the lower uh uh plots are just specific end to end uh results and natural training both for dance and sparse I'll draw your attention to the middle one the GPT 176B this is the big science model training that happened a few years ago and we actually used this uh because at that time uh the the project was going to lose half of its GPUs somewhere midway through the training and that's when we actually built this technology to enable it to continue training on half of the GPUs that it started with. Yep. And so you know technique has been adopted in many different places like the be science like I mentioned was sort of like you know the first one and so um yeah now I'll move on to the uh next work which is more of a communication domino communication free training. So a lot of the parallelism strategies like tensor and pipeline parallel they do I mean they partition but then you pay the price of some communication. Um and so here uh the goal of this project is to see how can we reduce that communication cost hopefully to like zero. Um and we started out by looking at tensor parallelism. So the plot up there the top left is showing when you do tens of parallel training of a GPT 13B uh size model and it's sort of breaking down the time between computation and communication. As you can expect as you go across nodes your communication time starts to dominate up to like 40% of your end to end time. And so the solution here is we introduce some fine grain overlapping of your compute and your communication. And we can do that either within your transformer layers or across transformer layers. Um and the plot over there showing up there is showing like a megatron lm you know sort of how you know if you do tensor parallelism your transformer computation got split into computation which is sort of you can think of it as a upper side and then communication at the lower side. Um if you use Domino you can see down there we're now sort of doing some fine grain overlapping so you don't have those gaps where your you know GPU is idle or your uh communication is idle and that helps you improve uh performance. Um so high level results here we're showing here that on the both for the forward pass and backward pass is just the profiler shots showing that we actually on H21 H200 we're able to completely hide the communication intensor parallels completely um uh on a single node um and then we have results comparing against uh Megatron LM uh just you know single node and multi-node and we get you know there's about 16% speed up compared to like Megatron LM Okay. Uh and then ongoing work uh we are integrating with hogenface natron. Um we have some good results there. We're achieving 75% communication hiding at this point. Um the goal is to get to like 100% communication hiding. We're extending this work toe and inference um and and long sequence training and some quantizations for the communication as well. Uh the next work is deep compile. And so the goal here really is that distributed training optimization is hard. uh it requires uh uh you know developer uh work and but then in particular tuning them to get peak performance for different models in different hardware is really difficult right and so the idea here is we want to automate this by using a compiler technology and so last year you know PyTorch released the touch to compile awesome technology that allows you to basically rewrite your computational graph in a flexible way and so we kind of leveraged that and so the two key ideas there is that we will be doing uh optimization using compiler transformations right um and so the first part is we will automatically uh apply the distributed techniques whether it's TP PP whatever we will use the compiler to actually do that to the graph and then we will do something which is kind of classic compiler tricks of profile guided optimizations we'll allow the training to run for a while we'll profile it in this runtime environment and then we'll tune the optim the parallelism techniques to get the peak performance out from that scenario and like I said it's built on touch compile. Uh the plot over there is just kind of high level view you know we take a model we run it through touch compile we get the graph from the graph we apply the tech the train uh the transformations for optimization and parallelism we run for a while we profile while we're running figure out what are the optimal settings for the various uh parallelism techniques and then we recompile and we do this over and over again since training takes a really long time you know within a few unread iterations we'll be able to get this done and then the rest of your time can be spent running at peak performance uh we've implemented three different optimizations at this point and we're working on more in collaboration with the community. Um and I'll show some uh so so so let me sort of drill into the 031 to sort of give you a bit more flavor here. So there are two aspects to it like I said is so here on the left hand side assuming you have like a single GPU implementation of your model with a parameter and some operators 03 basically involves sharding the parameters. Um and so here we can see with the compiler we're able to automatically apply the sharding uh instead of having the users uh have to write it and then the profile guided part uh when you shard uh uh parameters you pay the price of communication um that's the price um but then prefetching is sort of like the standard trick to sort of hide that but then the problem is how much should you prefetch when should you prefetch the profile guide it's showing there is showing the memory consump profile of your training. So you see that at the beginning of the forward pass, for example, you do have a lot of memory. So you can very aggressive with prefetching, but not so much in the backward pass. And so with the profile guided optimization, we're able to detect this pattern and then recompile the model uh to to to sort of take advantage of that. So in terms of results uh yeah we're seeing great results speed ups for 03 here both on the dense model and on the sparse the mixtural model uh you know we're seeing up to like 50% speed up. Um another optimization we did was 03 with offloading. So basically here we're able to also just compare to like baseline zero or zero using just pure touch compile. We're seeing you know over 5x speed up uh for the 03 offloading for zero stage one which is just optimizer sharding uh we're seeing about uh you know 30 over 30% uh speed up as well and I'll jump to the this is the last project deep the motivation here is that we recognize that um IO is becoming a key bottleneck as we continue you know deep learning scaling um and and it shows up in a few ways one in a few ways because uh uh where you know uh the persistent storage is becoming important. So the data your data sets the data loading data prep-processing all involve interaction with just persistent storage model checkpointing is sort of a natural one too. You have you save your checkpoint and persistent storage and then sensor offloading is kind of becoming sort of uh important and so the key idea here is that we note that there are a lot of innovations in the storage space and we want to build optimizations on top of that. So things like local NVME SSDs, a synchronous IO stack that like Linux now provides and then things like Nvidia GPU direct storage that allows direct communication between storage and HPM. And so there we've built different kinds of optimizations there. And um and so the results on there is just showing that you know on the PCI gen 5 like sort of the latest greatest um Azure VMs we can hit almost 50 GB per second disc reads there uh about 25 GB uh rights and and and importantly I'm showing there that um as we scale the number of NVMemes your performance is scaling linearly. So it's kind of cheap to scale. Um here are some just application um highlights and I'll just to wrap up. So here the one on the left is just basically showing uh uh you know NVME offloading which allows you to inference a very large model on a single GPU that you wouldn't have been able to do because it's not big enough. So we integrated this into into SG lang and so we can you know uh you can at least use this for like true based offline inference and over there we're showing about a 20x faster model checkpointing compared to pytor same um yeah uh speed is a community project it's been that from the get-go uh we have over 400 contributors a number of organizations supporting us and you know breaking news today uh it's now part of the Python foundation which is super exciting So uh just invite you to join us and let's uh make state of the state of the art AI available to everyone.