SBTB 2023: Supriya Rao, PyTorch 2.1 - New Developments
Recording: SBTB 2023: Supriya Rao, PyTorch 2.1 - New Developments
[Music] hey everyone very excited to be here um my name is Supra ra and I'm an engineering manager at meta working on pych in today's talk I'll be going over some background about pytorch highlight some newly added features and then conclude with how you can accelerate generative AI models using native pyos code so here's a high level agenda of what we'll be going over today I'll first talk about some background and metrics about py and then do a bit of a deep dive into some of the newly added features in pyo 2.1 I'll then switch gears and talk about how you can accelerate some significant generative AI models using native pyo code and then share Model results for these models with that let's get started cool so what is pyge py started about 7 years ago with the goal of powering researchers today it's one of the leading AI Frameworks powering a diverse set of use cases from self-driving cars to perhaps the robot vacuum that might be running in your living room it's open source it's a community based deep learning platform and if you haven't then you should definitely check it out today so as I mentioned here are some metrics about by Dodge these are some numbers that speak for the growth in pyo over the past few years the chart on the left goes over the growth in researcher usage over the past few years there are close to 7,000 GitHub repos that are using pych and are associated with AI research papers this year 60% of AI research implementations have chosen py Doge as compared to other deep learning Frameworks the graph on the right talks about the number of Open Source repositories that you pitor today in GitHub there are over 600,000 repositories on GitHub that are using pych and we've observed over 20% year-over-year growth in new repositories using pyo pitos has also made a significant impact on the tech industry data from LinkedIn suggests that there are over 65,000 LinkedIn professionals that have pyos listed on their profile there are also 3,000 jobs associated with p p and we have observed over 50% increase in pyo professionals just this year so let's recap a little bit to the last version of pyo pyo 2.0 was released earlier this year it saw over 20 million downloads across different platforms so yes it was a pretty significant release for us it also added a bunch of new features most notable of which is Storch compile with significant sped up training and inference of a large variety of models while offering the same eager mode development experience so let's flash forward to the next release which is also what this talk is mostly going to be about pyo 2.1 this was released earlier this year on October 4th and we also presented a lot of features that we released in p.1 at the page conference which happened in October this year going back to the number numbers we saw close to 7,000 commits and 784 contributors that contributed to pyo since 2.0 so let's dive into some of the features that pyos 2.1 has to offer note that this is by no means the exhaustive list of everything we've introduced in 2.1 for the full list you should definitely check out our release notes so as I mentioned before torch compile was a pretty sign significant feature for us and we are constantly making improvements to it so first I'd like to go over some improvements that we've made to torch compile since the twoo release first is the automatic Dynamic shape support deep learning compilers commonly only work for static shapes that is to say they produce compile programs which work for a specific configuration of input shapes and if you recompile and you need to recompile if any of the input shape changes with Dynamic shape support it allows TCH compile to generate a single kernel that can work for many sizes at only a modest cost to efficiency Dynamic shape has been greatly stabilized in py 2.1 and is now automatically enabled if torch compile notices a recompilation due to varying input shapes the next we've added support to uh how to compile numpy operations in torch compile itself so torch compile understands how to compile numpy operations by translating them into pych equivalent operations because this integration operates in a device agnostic manner you can now GPU accelerate your numpy programs and get the wins for free torch compile also works well with autograph. function now it has improved composability with autograd do function and we can now trace and optimize the backward function of userdefined autograd functions which unlocks training optimizations for models that make heavier use of extension mechanisms so the next feature I want to touch upon is something we newly introduced in 2.1 it's called torch. exort so torch export provides a sound tracing mechanism to capture a full graph from a py program based on new technologies provided in pyo 2.0 this can be intended to run on different environments which can also be python L how torch export does this is that it extracts a single graph representation from py programs by tracing the target function given some example inputs so you might ask how is this different from the previous version of torch effect symbolic Trace well it captures a much higher percentage of models compared to FX on 14,000 GitHub models we see over 88% pass rate on models that are exported via torch export one of the advantages of exporting a model is that you get an IR where users can run custom graph Transformations on for example quantization which I'll be talking about a little bit in a bit more detail in the in the next few slides the export IR can also be serialized saved and transferred and loaded back for execution with or without python so here's a small quote snippet walking through how you can use the export API in this example I have a simple model in which case in this case it's just a a small module you instant iate the module and you call export on the module with sample inputs this generates the exported program and if you print the type you can verify that it's of type exported program and you can further perform other optimizations or graph transformation on this exported program as I mentioned before some one advantage of exporting the model is that you can perform various graph Transformations on the IR one such transformation is quantis ation what is quantization quantization refers to techniques for performing computations and storing tensors at lower prit wids than the floating Point Precision for example in8 in 16 or in4 in P Toge 2.1 we added support for both post trining quantization and quantization aare training on models that are exported via to. export users can use a backend specific quantizer object to enable lowering a single py torch quantise program to different backends these could be backends on server CPU or different Edge backends like your mobile phones or uh specific dsps backend developers can write their own quantizers and expose methods that allow users to express how they want to quantize their model this is done via programmable API for both backend developers as well as users so it's totally configurable we also offer the option to specify quantized computation using integer arithmetic in the pytorch model Itself by doing this you can achieve higher numerical uh Precision on the server and be able to debug your model numerics more accurately so here I'll just go through a short example of showing how you can use these new apis there are two high level steps the first step is capturing the program or exporting here we call the capture pre autograph graph which is is in this case equivalent to torch export and the next step we perform quantization so the user would first have to call a specific quantizer in this example we are calling an xn and pack quantizer which is the quantizer optimized to run on the X andn pack back end which is meant for CPU once this quantizer is instantiated with a config or like a quantization config you call the prepare function what the prepare function does is it takes the model and the quantizer and it inserts observers into the model so that you can collect statistics that are necessary to quantize the model as a Next Step you can either calibrate the model for post trining quantization or train the model for quantization ofare training and generate those statistics that are required to quantize it the final step you would convert the model and the convert step produced a final quantize model which is optimized to run on the target back end in which in this case the X and impact back end the next feature I'd like to highlight is two4 semi-structured sparsity torch. sparse now supports creating an accelerating compute over semi-structured sparse tensors semi-structured sparcity is a data layout that was first introduced in nvidia's ampere architecture we have added support in 2.1 for creating The two4 semi-structured Spar tensors using a oneline API this support accelerated inference via both cutless kernels and qpar LD by running faster sparse mat Ms during inference qpar CT is a library that was that's provided by Nvidia for running two four Spar mmols by compressing the tensors to this format you can achieve a compression ratio of 56.25 for float 16 and 62.5% for in 8 so here's a Cote snippet on how you can use this API you first first have like your input tensor here which is X and then you instantiate a mask this is just for illustration purposes you can you can use different pruning algorithms to determine how you calculate this mask the mask here is essentially setting two out of four values to zero which is a format expected by the sparse kernel then you call the linear layer and you set the weight of the linear layer to the masked value and you call the two Spar semi-structured API which in this case is what is responsible for compressing it to run accelerated inference when you then call the linear operator with a sample input it will automatically take advantage of the uh fast kernels during inference so I now want to shift gears a little bit and talk about how you can accelerate generative AI models using pure pyos code so the space of generative AI is rapidly evolving as as you can see many Innovative models have released in just the last couple of years each of them pushing the state-of-the-art of what's possible in this domain and as with any new model released the community wants a way to be able to run it efficiently on different Hardware the generative AI landscape is no different from llama CPP to Frameworks like Mosaic ml or X Lama we have many different options of Frameworks that have been built to run these models faster during entrance but can we do this with Native pych without having to rely on the external Frameworks so I'd now like to walk through some techniques that you can use in pyo to accelerate these models the first such technique I want to talk about is sdpa or scale do product retention pyo now includes high performance implementation of the Transformer API using a custom kernel architecture for scale do product retention and Transformer a are at the heart of most generative AI models today we have multiple sdpa custom kernels which are supported with a kernel selection logic that will pick the highest performance kernel for a given model and Hardware type some of the options of kernels that we've integrated include the flash attention kernel and the X forers memory efficient attention kernel we also added the flash attention V2 kernel which was released earlier this year and it's now available in pyos 2.1 and is expected to be two time faster than the previous sdpa implementation of lash attention I will share some results later on how this technique helped with the performance of the segment anything model on GPU so here's how you can use the sdpa kernels you simply have to call the the functional operator scale. product attention with the sample inputs so in your model you can find where the Transformer implementation exists and if you go and make the simple on line change you can take advantage of the acceleration at Kels that P has to offer the next technique I want to talk a little bit about is nested tensors firstly what is a nested tensor in a regular tensor each Dimension has a fixed size nested tensors are similar to regular tensors except for their shape where not all Dimensions have regular sizes some of them are Jagged nested tensors are a natural solution for representing sequential data within various Mains like in NLP where sentences can have variable lengths so a batch of sentences forms a nested tensor or in computer vision where images can have variable shapes so a batch of images forms a nested tensor this is pretty common in a lot of Transformer based models or also common for models where you have varying input output shapes where you run batched inference over images by simply batching the images instead of of running them serially using nested tensors we are able to observe close to 2.8x speed up on the segment anything model so this is how the nested tensor API looks like you have two torch tensors here each of different shape A and B and then you invoke the Nester tensor API with the two tensors as inputs which then generates a new tensor object called nested tensors which includes these two tensors as inputs so in your model if you have inputs with varying shapes you could potentially use a nested tensor and take advantage of the batched processing so I touched upon quantization a little bit we also added support to do quantization on the GPU which is pretty relevant for generative AI models today to recap quantization trades of numerical Precision for reduced model size and increased execution speed some of the benefits of using int 8 math over fp32 means that we can reduce the model size by four times on disk and in memory and the matrix multiplication also requires lesser computation and we've seen theoretical speed UPS of up to 16x compared to fp32 so what are some of the techniques that you can use to quantize models on GPU you have Dynamic intake quantization this is dynamic because the activations are quantized dynamically during entrance while the weights are quantized ahead of time you have static inted quantization where both the activations and the weights are quantized ahead of time and then you have weight only quantisation which as the name suggests just runs the weights in a compressed format but stores the activations in full Precision we implemented both Dynamic and Weighton quation for gpus and we see good results on both segment anything and Lama 2 model as you can see below so on the Sam viit a backbone intake Dynamic Quant leads to about 9% speed up and on the Lama 2 7 billion model which is pretty memory bound inate weight only quantization leads to about 47% speed up so we have an API available that you can try out for doing this on your model it's available in pyo Labs under the AO repository here's a sample code of how you could use it you have any user model in this example it's an sequential layer uh sequential module with a single linear layer you have sample inputs and then you call this oneline API depending on what you achieve what you wish to do if you wish to apply weight only quantization you can call apply weight only inate Quant on the model and if you wish to apply Dynamic quantization you can call the apply Dynamic quantization API on the model you can also torch compile the model to improve performance further and these apis compose well with torch compile and then once you run run the model you can take advantage of the accelerated int8 kernels during inference time I want to talk a little bit more about two4 sparcity and how we added support for that on GPU as we spoke before people want faster matrix multiplication and that's one of the most compute intensive operator for most generative AI models and enabling two4 structured Spar Fates gets us that to introducing two4 sparsity can come with some accuracy impact and in order to elevate that you can fine tune your sparse neural networks to reduce the accuracy impact and to do this you can Leverage The Torch AO pruning apis and use different algorithms to prune the weights this also composes well with int8 quantization so you can have two four sparse int8 kernels which will give you 62.5 compression ratio for weights and up to four times theoretical speed up compared to dens fp16 so let me walk through the code a little bit here it's a bit more involved so bear with me I'll I'll go step by step uh I'll first talk through the helper functions that you can use to sparsify or prune your model in this case we will use the torch. ao. pruning package which contains a weight Norm sparf which is a magnitude based sparf these sparsifiers work by applying mask parameterizations to the weight tensors in the model this Les simulate sparsity by masking out the prune weights we also have to decide what layers of the model to apply sparsity to in which in this case we apply it to all the linear layers and we set that in the sparse config we now go to the sparsify do prepare and the prepare step inserts the parameterizations necessary to mask the weights of the model the step function is used to update the mask based on the logic specified in the pruner in this case it will update the mask based on the magnitude of the weights when we reach a final satisfied stage we can call the squash mask function which is used to fuse the weight and The Mask together so the final result in this example would be that you have a weight tensor which is uh compliant with the two4 SPs format and has been pruned using the weight Norm sparf once we have the model in this format we can accelerate for inference using the oneline API that I mentioned before two Spar semi-structured we update the weight of all the linear layers in the model to use this format and once we run the model during inference it will automatically run the sparse kernels during runtime the next technique I'd like to talk about is pretty relevant these days in the context of running large language models which is 4 bit quantization quantizing the weights down to 4 bits can have an impact on model accuracy and gptq is a tech technique to reduce that so gpdq is a one weight quantization method which is based on approximate second order information of the weights we leverage torch
export to prototype an easyto ous GPT gpdq implementation that you can try out today however to run this efficiently on GPU we do need some custom 4bit quantize matrix multiplication which we have implemented in Cuda so here's the example API that we used to implement this you have a GPD 4bit gpdq Handler this Handler is responsible for converting the weights to four bits and uh running the model as well so as a first step you create the quantize state dict you can run the step ahead of time so that your weights are already in the 4-bit format and you can store the weights separately in The Next Step you can prepare your model for runtime so here you would do two steps where you first ready the model for for runtime here which is done by the convert for runtime function and then you can load the previously saved State dict by calling the load State dict function so as end result you have a model that has been 4bit quantized great so now I spoke about a few different techniques let's talk about some of the results that we observed on models using these techniques so I'll mainly go over two models here the segment anything model and Lama 2 segment anything or Sam is a zero vision model for generating prompt image masks our Sam Baseline is Facebook's Facebook research researches unmodified model using a floating Point 32 data type and a bath size of one and here we show The Benchmark results on the viit H backbone so on the graph on the left you can see the Improvement in the performance or the images per second for the segment anything model we start with the floating Point 32 model and as a first step we convert it to run in 16bit Precision this itself gives us a big boost torch compile further improves the performance and once we introduce sdpa or pytorch's native scale dot product retention kernel we can significantly increase the bath size you can see that the bath size here has gone up from 16 to 32 so you can run the model faster and for the following techniques we continue measuring with a bat size of 32 we tested with a custom Triton implementation that further improves the performance of sdpa we added nested tensors and then we added intake quantization and sparsity overall we have achieved over 8 8X speed up over the fp32 eager Baseline and we also observe no accuracy impact when we run the model with inate we do see some accuracy impact with two4 sparsity mainly because we did not apply Supply fine tuning to reduce uh or to prune the model in this case the graph on the light right shows the memory utilization at bath size 32 so as you can see we start with a bath size of 1 and once we reach s sdpa we can increase the bath size to 32 and all the further techniques reduce the bath size quite significantly to make it to make the memory utilization similar to what we observed of floating Point 16 so I'll now talk about some results on the Lama 2 model so the main techniques we applied on the Lama 27b model was Storch compile inate weight only quantization and forbit gptq by applying these techniques we saw over a 7.5% 7.5x speed up over the F uh fp16 eager Baseline you might wonder what's the accuracy impact on the Llama 2 model by applying some of these quantization techniques well there is minimal impact when we apply 8bit quantization as you can see on some of the metrics in the chart on the right and with four bits as well we don't see significant degradation because gpdq helps us uh recover some of that so if you look at the metrics four bits and 8 Bits perform quite similarly in terms of model accuracy with four bits giving a pretty significant boost in performance we also applied some other techniques on Lama 2 including speculative decoding and tensor parallelism by applying these you can see that the final bar here goes up to close to 244 tokens per second on the Lama 2 7B model we believe that this is the current state-of-the-art performance for llama 2 on gpus we also apply the similar techniques on Lama 70b and by applying them we achieve close to 80 tokens per second on an a100 GPU note that without applying these techniques the Lama 270b model would o on an a100 GPU because of the memory restrictions so if you interested in uh checking out these models or trying some of these techniques you can check out the these following repositories on py Labs uh we have one for segment anything fast and we also have the AO repository which includes the GPU quantization techniques I talked about earlier and we are very soon going to release the Lama to repository as well call GPD fast with that that's all I had uh thank you for listening to the talk please check out the other resources on pych and I have a few more minutes I think for questions in case people uh would like to ask some right now [Applause] thank you so I have a quick question about like this Facebook PL uh meta is planning in the future like turn P into a hardware agnostic framework or at this point you still needed to do all the quantization and all this kind of stuff for Hardware specific um so by Hardware agnostic do you mean just pry much if I train my model or I build my model using pter at this point I should be able to run the model in uh different Hardware targets or uh so what kind of Hardwares are you talking talking about about the accelerators you know like about sorry accelerators the GPU like xlaa and all of this kind of stuff is right so with the um the quantization the export based ization that I talked about earlier once you run the convert step you technically get a backend agnostic representation but you can further run like custom graph optimizations to map those um operators into your custom like Hardware accelerator operators so in terms of quantisation you would maybe get a pattern like d quantise floating Point operator followed by quantise and you would map that pattern to your custom Hardware um quantized operator so you can further run lowering passes on the reference quantise model generated to U make it work for your Hardware um the reason why we have backend specific quantizer is to help with accuracy because different Hardwares may have different requirements or specifications on what numerics they can run or what operators they support so we have the ability to specify that ahead of time so that the model uh not just like is able to lower but also performs better but if you do not really care care that much about accuracy you want something that runs across different platforms then you can still use the same flow and lower it on different Hardwares using your own custom lowering stock cool and finally the other question that I have is uh p is doing a really good job with the building blocks to for the modeling part but so what about the deployment part so do you guys uh you know like you have a p is thinking about having an API to deploy models in a mobile device or iot syst or this kind of stuff um so it's actually a little hard for me to hear the question it's U my question is about that fighter is uh doing a really good job with the building blocks to for modeling but so what about the deployment so at this point is there any API to deploy your modeling a mobile device or right uh so one major release we had recently which I did not cover here was exeggutor which pyo released in the conference earlier this year which is the ability to take any pyo program and deploy it on multiple different um mobile backends uh ranging from like uh arm to Apple to qualcom and we are also actively working engaging with these different vendors um and executor is also open source so you can uh check check that out as well all right thanks for your time everybody it was great [Applause] presenting