Scaling PyTorch Model Training With Minimal Code Changes

Key Takeaways

yeah hi everyone so today I want to talk about scaling pytorch model training with minimal code changes so I want to show you a few tips and tricks how we can make our pytouch models train much faster without having to put too much effort into changing the code so we will take an example here a vision Transformer that I trained from scratch and this Vision Transformer took about 60 minutes that's approximately an hour with a test set accuracy of 62 percent as a baseline here so the details are I would say not that important so that's a very basic Vision Transformer that I trained on a very simple data set to keep it short we want to dive too much into the details of this Vision Transformer but if you are interested the code is here on my GitHub repository here if you want to check it out and all the examples I'm showing you in the next couple of slides they are also in this GitHub repository so if you are curious about more details yeah you're welcome to check out this repository and if you have any questions yeah please feel free to ask a question in the discussion section or here in the comments and I will be happy to chat so one of the main points I wanted to get at here also is that it's 2023 and training from scratch often does not make sense anymore so with that I mean usually when we train models modern deep learning models we usually start with a pre-trained model so for example when we work with large language models or Vision Transformers we usually use a pre-trained model and fine-tune that pre-trained model so when I do that when I use a pre-trained vision Transformer here on the right hand side and I fine tune it for three epochs I can achieve 95 accuracy in just about 20 minutes which is a big improvement from the left hand side where I try to train this Vision Transformer on the same data set from scratch now so here's just a summary because Yeah the code view was a little bit small just in bigger pictures we can achieve a way better test accuracy using a pre-trained Transformer that we fine-tune and of course it often also requires less epochs so it usually takes much less time than training a model from scratch so one way of doing that is using fabric which is a fast and lightweight wrapper library or API around pytorch which allows us to let's say use Advanced features in pytorch very conveniently so fabric essentially is a way to bring these Advanced functionalities without having to restructure our pie torch code here's an asterisk because there is one catch when we are using Fabric and that is that is 100 free and an open source Library so if you want to use that you have to install the lightning library and open source library and then from there you can import fabric if you are familiar with pytorch lightning so it's essentially by the same creators of pytouch lightning except it's a bit I would say yeah leaner so that means a fabric is essentially a simpler or more lightweight way of bringing this Advanced functionality of pytorch lightning to pytorch models so I will show you first of all before we see the speed advantages I wanted to show you first the small modifications we have to make so suppose we have the simple python code here shown on the right hand side so I kept it very simple and very abstract um you're only focusing on the main parts so what we see here is that we are importing a data set we are defining a pytorch model for example our vision Transformer we have a custom pytorch data set and then when we have a forward path that we usually have is the a loss function then we initialize the pytotch model initialize our Optimizer initialize the data loaders and then we usually call yeah the model on the inputs compute the loss back propagate using dot backward and doing a step with the optimizer so if you have coded in pytouch before this is something you have done probably a thousand times so there's nothing really fancy or new in this code here so the main part I wanted to focus on is the little changes that we have to make when we want to use fabric which then in turn allows us to use a few more of these Advanced features in pytouch more conveniently without having to restructure any of this so here highlighted are the small little changes that I was just referring to so on the left hand side is the code the pytouch code that I just showed you on the right hand side are these little modifications if we want to use Fabric and you can see it's only one two three four five six little modifications here so the first one is importing fabric the second one is here instantiating a fabric object then using fabric.setup on the model and the optimizer and then also calling fabric dot setout data loaders on the data loaders and then yeah essentially in the forward pass the only little difference here is that we call fabric dot backward on the loss instead of calling loss.backword but yeah these are the only little small changes that we have to make here if we want to use fabric but yeah since this was a bit fast and there was a lot of information on this one slide let me do this again step by step one at a time so that you can really see one change at a time before you see the whole big picture and before we apply it to our vision Transformer so the first step like I mentioned is importing Fabric and then instantiating this fabric object so when we instantiate a fabric object there are a lot of options here I'm showing you a few simple options using a Cuda so a GPU accelerator optionally using multiple gpus if we have them and then defining a multi-gpu training strategy but we will get to that more in detail later so the Second Step here then is yeah setting up the model the optimizer in the data loader so here only two lines of code to set up you have the optimizer the model and the data loader then the third step here is calling fabric dot backward on the loss instead of the usual loss.backward so if we do now a performance Benchmark here with plain pie torch and pythagen fabric we see that we get exactly the same performance if you see a little yeah a fluctuation that is yet due to Randomness but that's what we expect because fabric is just yeah a wrapper around pytorch so there's nothing really I would say fancy happening yet so that's just like the same pytouch coat it's just using fabric now so as a baseline we can see we don't get any faster or we don't get any slower the magic happens when we for example use Advanced Techniques like a mixed Precision training so we can see here with mixed Precision training we cut down the training time from about 18 minutes to 6 minutes while maintaining the same predictive performance here on the right hand side and yeah this is great um and so how do we achieve that when we use fabric it's really just one line of code that we have to change or not even a whole line of code just maybe uh 30 or half of a line of Code by just adding this argument here position B F16 mixed so bf16 mixed here stands for a brain float 16 and then mixed position training so what does that mean so usually when we train models on a GPU we use 32-bit Precision which is shown here at the top so it's essentially just a visualization of how much um Precision we have like the fraction here how many numbers behind the decimal point are let's say accurate and then when we use half Precision offload 16 position you can see here at the bottom there is less position so the representation is less accurate for the numbers after after the decimal point now brain float 16 is a I would say special version of float 16 not all gpus support that but if your GPU supports that I would usually recommend that one it usually can represent larger numbers it has a even a bit less precision but you can yeah have a larger dynamic range exactly actually exactly the same dynamic range that float32 has and when I was training a large language models I noticed that usually B float 16 works better than float 16 which means it's a bit more stable so sometimes when you use float 16 especially if you use pure float 16 training you may get nands in the loss and B float 16 is usually a bit more robust but to be honest both work fine so yeah feel free to choose either or especially if you're using mixed Precision not pure precision so mixed position is essentially a trick that uses both um 16 and 32-bit Precision so that there is no loss in accuracy so how it works is essentially it takes the 32-bit um weights converts it to 16 bit weights then computes the gradients in the 16-bit representation to save GPU memory and then convert it back to a 32-bit gradients then use the 32-bit gradients to update the parameters and then yeah continuing on so in this way we usually can cut down the training Time by a lot because Computing the gradients in 16-bit representations is they are much faster than 32-bit and also it saves often a lot of memory so if you are memory constrained or compute constrained mixed Precision training is actually a great way to accelerate the training without having to let's say take impacts on the prediction accuracy now another technique is of course multi-gpu training but only of course if you have multiple gpus but if you have multiple gpus it's great you have to take advantage of that to train the models even faster so one of the more Advanced Techniques behind our multi-gb training is fully sharded data parallelism so there are actually a lot of different methods for multi-gpu training or multiple different paradigms there's model parallelism data parallelism pipeline tensor and sequence parallelism and I would say model parallelism and data parallelism are the classic variants pipeline is a bit I would say more novel sequence parallelism is something we use um for llms sometimes but tensor parallelism I think it's one of the really really um cool things because it enables us to train models that usually don't fit into a single GPU anymore because it's essentially dividing the tensor so fully sharded data parallel training uses essentially both a data parallelism and tensor parallelism so in data parallelism what we do is we usually divide the mini batch and then have a copy of the model on each of the gpus so here I'm showing two gpus where each of the gpus gets let's say a sub set of the mini batch so usually we refer to them as micro batches and um yeah this helps essentially to train the model faster because we can iterate through the data set faster by having multiple gpus work in parallel and yeah it will depending on the method it will then average the gradients or average the loss to get to the data set or get through the data set faster so tensor parallelism is a bit different it divides the model across gpus so if I go back data parallelism divides the mini batch across devices but it doesn't save you memory because one GPU has to fit still the whole model here whereas tensor parallelism splits the model so that if we have a large let's say large language model or a large Vision Transformer we can have it on multiple devices we can break it up and that allows us to train models that we usually can't train so why or how does it work so if you think back of how matrix multiplication works there are essentially two ways we can distribute that we can distribute it by row or by column here I'm showing just by column for brevity so for example if you have this matrix multiplication here shown on the left hand side what we can do is we can break that up into two operations so you can think of this operation here at the top and at the bottom is two separate computations and we can carry out each of the computations on a different GPU and then we can concatenate the results back to get the original results and so in this way we can distribute this large computation if these are two large tensors across multiple devices and how we can enable that in our pytouch code is Again by adding just a small little change here to our fabric initialization so here let's say we have four gpus we set devices to four and we select the strategy fsdp which stands for fully sharded data parallelism and of course it's just one of the million strategies you are also free to experiment with other ones for example there is deep speed and you have many other techniques or colossal AI there are many Advanced multi-gpu paradigms I like fully Charlotte data parallel because it's um yeah it's already part of pytorch so you don't need any additional libraries for that but yeah this is just one of the techniques and if I then use a fully sharded data parallelism and train my model on 4 gpus I'm cutting down here the time that it takes to only two minutes so I'm going or from the original model which trained in 18 minutes down to six minutes when I use mixed position training and then down to two minutes when I use fully sharded data parallelism with mixed Precision training so yeah this was a very short talk um to just show you some of the advanced um I would say data or training acceleration ways we have available in pytouch which are usually really tricky to use so if you ever coded on distributed training in pytouch it requires quite a few code changes so with fabric we can make this a bit easier and of course it's an open source a free library so there is no downside of using that so yeah if you have any questions as you know you can find me on social media or yeah you can also check out my website where I have more information about certain things and um yeah the code like I mentioned the code in the slides are here on GitHub if you want to check them out and um yeah what I've been up to lately is I've been really focused also on writing writing a lot of Articles so you can also check out my head of AI or newsletter or rather magazine where I post about yeah things related to large language model training a lot of technical stuff or Recent research highlights and explanations and so forth so yeah if you want to check that out um I would appreciate it and with that yeah I hope you liked that video and if you have any questions please feel free to ask anything in the comments or here in the discussions on GitHub so yeah see you then