Load Large Language Model (2B Param) on Colab without OOM Memory Issue | Pytorch Tensor GPU

Key Takeaways

hey friends welcome to online recorder recently in one of my popular videos which is about um building your own github co-pilot using salesforce's code gen ai model i mentioned that you know i was using google collab and i was using the lowest size model which was 350 mb because i could not load large size model large language model llb llm sorry large language model with for example to say that two billion model after i mentioned this actually the founder um the creator of code gen um replied to one of my tweet on twitter saying that we could load the large model by setting default tensor allocation to gpu memory so i want to stand corrected and i want to also show you how to use um how to load large language model especially code gen related i've not tested it with any other language model but for code gym if you want to build your own github co-pilot or you want to build your own programming assistant ai assistant then you can actually indeed load large language model especially in my case i'm testing the two billion model and by using by asking by setting the default tensor allocation to deep gpu what does it mean and i'm going to show you in this video that's what this video is going to be very short if you already have seen this code um the only change you would see from my previous code and the current code is this part as usual we're downloading the we are installing the latest transformers from github the next thing is we are saying we are telling torch to set the default tensor type to torch dot cuda dot float tensor by default if you see the torch cuda um let me show you torch set default um so if you see the default tensor type is by default it is pointing to the tensor type t and then it uses the default type what we are asking now on pi dots to do is we are telling pythons that use the gpu use torch dot co dot float tensor so now the point is why cannot i do this for all the models i think this is not um aligning in line with what pi touch is advising us to do there has been a little bit of discussion in the community as well for example why why don't i just set it up by default and then there are there are like people i can show you um yeah is there anything i'll link this discussion in the youtube description so you can see so around 2018 this kind of discussion has happened why this is not a great idea why it breaks some models why it doesn't go well with what pytorch is saying so i'll link all these conversations for you to make a decision for yourself if you want to use this setup in all your projects but at least for this project for this code gen if you set if you force pytorch um the tensor allocation to use gpu memory then you can indeed use um load the large model if you do not do this by the time the two billion model gets downloaded your google collab session will crash and restart but if you do that i can quickly show you my memory so you can see that i've maxed almost maxed out my gpu memory at this point and uh it still works like i can just simply run the code so i've loaded the two billion model now i can say for example if let me let me keep this in the side so that you get to see this as well and i can say what do i want to say i want to say def um default slower to upper i can say input text it's a very silly code but let's see so i'm going to i'm running this code and you can see that this code is running and it is going to potentially produce the result at this point let's let's wait for wait for it to and you also you know you you know that you have loaded a two billion um model a two billion parameter model in the last video i said dollars so just being careful so we have loaded a two billion um parameter model so it's going to take a little bit of time in executing you can see that it took about 18 seconds for it to execute this is just inference where you take the text tokenize it use the model to generate the output and then decode it and then show it show the result but it has given you the result like whatever you wanted if you if you tell um it it works completely fine and then you have a you have a larger model which means it must be more um it must have more capability in producing the result than having a 350 million 350 mb um million parameters 350 million parameters so two billion parameters is always better i just wanted to quickly show you the memory allocation just to show you that this actually works when i run this this is from torch coda memory summary you can see that the gpu uh reserve memory is completely full a couple of um like it's totally allocated but you can see that there was no cuda out of memories so that that is a good thing so even when we have a two billion parameter model loaded um our machine is working fine and we could do inference fine without maxing out so i think this is a very good tip so thanks to thanks to eric for giving us this tip so this something that i did not know that you could do even in the previous case we were actually you know at the tokenization side and at the model site we were converting um everything to coda but we did not set the default intertype and like i said this is um this is a point of conversation and discussion in python's community you can see on the community pages and there was a github issue which talks about that we should get away with this so i would link those things in the youtube description but but for now you know this just works you can load a two billion parameter code gen model on google collab and build your own github cobit if you have not seen that video i would i mean as as as a youtuber i would strongly recommend you to watch my video on how to build your own keyword co-pilot with code gin dot ai sorry code gen ei from salesforce but otherwise um this is a this is a new tip on how to force um the default tensor allocation to gpu memory and then load large language models any questions let me know in the comment section otherwise happy coding see you in the next video peace