Sponsored Session: Lightning Talk: Optimizing Model Inference with PyTorch 2.0 - Devansh Ghatak

Key Takeaways

All right. Hi, how are you doing? Uh, good afternoon. Hope you're having a good time at PyTorch. I am Dimanch from Simply Smart and today we're going to talk about model optimizations using PyTorch on cloud. So, if you guys looked at the app and thought this was edge, I'm sorry to disappoint you. This is not we're talking about cloud today. And Simply Smart, we are an end toend MLOps platform primarily focusing on inference. So as you can imagine uh I spend my days thinking about inference optimization more often than not and honestly this has been frequent uh in the recent months uh especially when you compare it to the past few years right because when we were optimizing models a few years back uh the larger size of the models that we were working with a few hundred million parameters uh there are very few companies who are working at a scale where inference optimization really becomes a problem And also uh there were very limited use cases that we're serving right but today we're serving Kimik K2 a trillion parameter model at with ease uh we're serving a variety of use cases with this large models and inference is essentially democratized right so now is a good time to sort of start thinking about inference optimization and given that we primarily work with inference we think about this day in day out uh we are optimizing new model architectures we're integrating new optimization techniques new inference techniques and When we start optimizing these, we start thinking about or we ask ourselves what are we optimizing for, right? What we realized was there are different use cases that people want to serve using the same models. Uh but they have very different KPIs and SLAs, right? Take an example of an LM that is serving two different use cases. Say a voice bot which of course has a real-time mandate, right? You need to serve your time to first token should be within let's say 300 milliseconds. Whereas if you're serving the same LM for a document parsing use case, you would probably want to maximize the throughput, minimize the cost, right? So depending on what model you're serving, for what scale, for what workflows, you will need to have different sort of optimizations, different inference stacks, different deployment strategies in order to optimize for those workloads, right? And we figured out that there are majorly four KPIs or SLAs that people focus on when they're optimizing their models. And that's cost, latency, quality, and throughput. Right? Uh you can care one or about one or more of these things. Uh and you may have SLAs for one or more of these things, right? So uh we'll of course dive deeper into some production case studies which explore these problems deeply uh using uh some optimization techniques, right? But uh in order to optimize these we of course think about inference optimization in three layers. Uh so one is the infra layer. Uh how fast can you load the model? How fast can you spin up an instance? Uh some application layer optimizations like how are you batching uh your requests and majorly what we're going to talk about today is the model level optimization specifically quantization compilation and custom tensor ops. So let's start with one of the most tabooed optimizations which is quantization of course and uh of course torch gives us very straightforward ways to quantize models right you have torchio which exposes very uh simple APIs to sort of quantize and calibrate models you also have a lot of libraries like lm compressor which are very use case specific they extend these torch capabilities in order to quantize the model but let's be honest none of us really wants to quantize a model right we would like to preserve the accuracy as much as possible And that makes us apprehensive about quantization. But there are nuances to this. Uh not all models are affected by quantization equally. Uh larger model tend to do well with quantization. And it's also very task and calibration data specific. Right? If you have really good calibration data for your tasks, you can actually get away with quantization really really well. Case in point uh we had a user who was deploying a 12 billion SLM for a real time voicebot use case which was translating languages right now of course because it's a voice bot you have a realtime SLA that you have to adhere to so your TTFT becomes your primary SLA but there's also business mandate that the per translation cost of the model should be low right now given that the model is decently sized uh the task of translation also fairly deterministic and you can easily generate calibration data for this right int 4 made a lot of sense for us uh with good calibration we actually get uh the accuracy difference down to 0.01% 01% of the baseline model and we were able to unlock smaller GPUs with lower memory footprint that eventually ended up saving significant amount of cost right so uh by choosing the right task or the right model we were able to get to a state which was suitable for a particular business requirement and of course this is one of the more straightforward optimizations that we're talking about today right let's talk about something that's a little more nuanced when we actually take it into production uh which is compilation right so of course with the introduction of torch compile and torch export in torch 2.0. Uh we have again very simple APIs to do that. We see direct uh comparisons with eager mode and we see really good speed ups there. But you know what may seem like a silver bullet is rarely. So uh so we were working with LLM essentially for production use case where you need the full output in a certain amount of time. Right now, the team uh that we were working with, they spent weeks fine-tuning the VLM config in order to get to the right state, right? But they were still missing their SLAs's in total latency. Uh when we actually ended up optimizing the model to its most optimal state, uh we actually changed primarily one part of the compilation which was essentially how you compile the CUDA graph. Uh instead of using the full mode, which seems like a good default choice, right? You would want to optimize the full model. uh we actually went with a full and peace-wise which actually focuses more on the decode. It makes the decode layer fully compiled whereas it does piece wise for the rest of uh the parts right and it makes models perform better in terms of their output throughput right so eventually for spiky workloads where they were missing the SLAs's at P95 once you change the compilation mode from full to full and peacewise you are easily able to meet those spiky workloads uh by reducing the TPOT by 15% right and By taking these simple trade-offs, what seems like the most optimal state by taking the simple trade-offs by looking at the actual data, we're actually able to get to a state which is much more optimized for your particular workload. Right? And that brings us to our last optimization that we're going to talk about today, which are custom tensor ops. Uh, of course, uh, in an ideal world, every equation, every function that you have in a paper or every model that you're implementing, uh, for inference, uh, the ops for any chipset should be in the most optimal state, right? But that's rarely. So we actually choose a lot of different frameworks, a lot of libraries that offer us a lot of these custom tensor ops out of the box and then uh and that's that's generally the pattern that we've seen as well that and then we write custom tensor ops on top of them. Right? When we're doing this uh the big challenge that comes about is compatibility with versions with different libraries with different frameworks because there's not a single umbrella where all of these optimizations or custom tensor ops are available, right? uh and this happened with us fairly recently when we were optimizing the flux pipeline. So we have a direct integration in our platform where users could bring in their fine-tune diffusion or flux models and we optimize them very frequently with sage attention. Right now when sage attention came out and we patched it with the diffusers library uh it gave us a certain baseline speed and we considered that as our ground truth. Recently uh de uh diffusers actually came out with the newer version which gave sage attention out of the box. Now what this entailed was you could use it directly from diffusers and with some new features. But when we did a head-to-head comparison of sage attention we actually figured out that the original implementation of the paper is actually 20% faster than what uh diffusers exposes natively. Now the choice in front of us is do we let go of the new features? Do we let go of the newer optimizations in the new version or do we keep using the older version, keep getting the core speed that we're getting in the model? Right? Of course, the answer for us was best of both worlds. Patch the older Sage attention into the newer diffusers library in order to get the right speeds. But it's easier said than done, right? We as ML engineers, as inference uh engineers, we are constantly making choices. is we're constantly making trade-offs between different techniques, different ops, different libraries, different optimizations, different quantizations, right? And we talked about just three optimizations today, but when you start including different techniques like speculative decoding, KV cache management, things like that, we face a combinatorial explosion, there are so many choices that it eventually leads us to a grid search hell. What this entails is we generally when we whenever we are creating or assembling our inference stack, we are generally going through the most walked path instead of actually exploring the combinatorial possibilities that exist today. Right? And that's why at simply we sort of reimagined the way we assemble our infant stack, right? Modularity and abstraction is key. So at first what seems like a Jenga tower, you replace those Jenga towers with Lego blocks. You make these techniques pluggable. You abstract them away. You feed com compatibility data as well as benchmark data into these uh different techniques so that you can basically approximate what sort of speeds or what sort of cost or what sort of quality are you going to get for particular optimizations when you swap them out. Right? And that leads us to a lot of non-intuitive solution. Like if it is a real-time use case, do we not go for the fastest GPU that is available out there? Turns out you can actually do that given the quantization study that we just saw. You can actually work with smaller GPUs like 8NG with info quantization and still achieve this SLAs's get higher availability, get lower costs and that is what we are trying to do at Simplysmart as well. uh we're trying to solve one problem at a time, iterate faster with our plug-and-play framework that lets us reach novel solutions for well-known problems while optimizing for four KPIs that we just talked about. So working backwards from these four KPIs we can actually approximate a lot of the solutions find the best choices that are possible for these uh SLAs or KPIs and eventually reach to a solution which is optimal for the business and uh technical use case right uh we are at booth G9 and we are hosting an inference clinic all day uh we are trying to solve this grit cell exactly so if you are working on inference if you're working on ML models if you're optimizing one of these models through these techniques Do drop by at booth G9. We'd love to talk to you. Uh that's my time. I'm Dwanch. We are Simply Smart and we are at booth G9. See you.