PyTorch 2.0 Live Q&A Series: PT2 Profiling and Debugging

Key Takeaways

all right looks like we're live hello everybody um welcome uh this is the uh ask the engineers to point uh follow-up stream to uh the announcement of pytorch 2.0 that took place at pytorch conference in New Orleans um on December 2nd my name is Justin Jeffress and I am a developer Advocate at meta and I have a very special guest with me today his name is Brett Maher um right would you like to introduce yourself yeah sure thanks uh so I am a software engineer here I work on compilers specifically the pytorch compiler so I'm looking forward to taking you all through a bit of the internals today great and uh oops I got the wrong screen up here we go before we get into Bird's talk just wanted to give a little bit of context so uh yes we're going to be talking about profiling and debugging uh this is um going to be really helpful for those of you who choose to start experimenting and playing with uh pytorch 2.0 it's available in the nightly builds and um uh I don't know if you've tuned in earlier this week siraj did a post or a stream about uh how to download and install the nightly so you can get started um but you know one of the things that Pi torch 2.0 introduced was torch.compile and so Bert's going to be talking about um how to use the different protiling profiling tools to help you understand the output of the torch that compile which will help you debug and also help you get help if you go to GitHub and submit an issue and if you venture into the land of contributing to Pride torch knowing how to use these profiling and debugging tools will help you do that as well and speaking of that this afternoon we're also having another presentation um on the same channel um Regarding why you should contribute tutorials and code to Pi torch so I just wanted to call that out um so please you know tune in that's a 2 p.m Pacific time today um so this is part of a series that's going to take place throughout uh December to and January um and there is a schedule of um um the series so you can see if you go I'll share the link here into the chat um but this page here shows you the list of all of the streams that are going to be taking place over the next few weeks or months so um go ahead and take a look at that and I will go ahead and transition over to you Bert so you can start talking about uh um you know the profiler and I'm really excited about this because you've actually built and worked and done a lot of debugging on this so here we go all right so yes uh as uh Justin mentioned uh this is going to be a talk about uh understanding the output of the pi torch 2 compiler and we're going to focus on the profiling and debugging tools that will help you really get to know what your code is doing uh why it's working well why it's working poorly and hopefully tune your applications to get even better performance over time so as you probably know the main feature of pytorch 2 is a compiler this is great right compilers make things fast they make you productive but anybody who's ever tried to debug optimized code let's say you've tried to debug a binary that you built with gcc-o3 you know it can make things tricky right it's tougher to debug optimized code than just pretty straightforward naive say uncompiled code I'm going to help you figure this out right so the goals for this talk let's say there are four the first one just learn the internals as a software engineer I just like digging into this stuff it's really cool I assume that you've probably seen some of the background materials about pt2 so far so you're a bit familiar with the components but just to give background to the rest of the talk I'll go step by step through the major components that we've added to build this compiler stack the second part of the talk we're going to really go under the hood uh we're going to look at a lot of the logs and information that you can get about what the pytorch 2 compiler is doing so that you can really understand what's happening when the pt2 compiler compiles your code the third section yes right just one thing I always wanted to mention that I just pasted a link into the chat um which is just has an overview of Pi torch 2.0 so for those of you who haven't had a chance to to dig into or learn about what's new uh in pi toys 2.0 uh take a look at that link bookmark it and then after this talk go take a read and um you'll get up to speed and learn about uh what's exciting and new in 2.0 so um awesome thank you yeah there's a lot of good info out there on the web I hope people take a look at it it's pretty exciting stuff so the last two sections we're going to figure out some of the tools that you can actually use to improve your code so the first is the troubleshooting part uh what happens when the pt2 stack goes awry how can you uh figure out what's going on how can you produce better bug reports that will help us help you and then the last section is on profiling and tuning performance which I think is some of the really fun stuff that you can do once you have an understanding of the system the format of the talk there are these slides hopefully they'll be somewhat useful as a reference there's lots of material on the web as well this kind of puts a bunch of stuff that's very useful all in one place uh the other thing I'm going to do is some live coding which I think will be uh will be fun yes uh it's always a fun for me you know a danger right like it could be something that goes wrong but uh hopefully you know exactly I'll add some drama right right is it gonna run um so we'll do some live coding I'll show you the actual logs uh you'll see a gigantic amount of logging output hopefully some of it will even make sense uh and then as as you probably observed we'll kind of have a running q a at any point uh if you guys want to ask questions about the material that I'm going through um you know drop a comment in in the chat and uh and Justin can vectorate to me I'd love to help answer questions that you all have about this yes and and feel free to ask questions during the middle of the talk like like Bert said I'll be watching through the comments and I'll Bubble Up the questions to him um while he's speaking in so um so don't be shy yes okay so content really briefly we'll go over the major front end components those being torch Dynamo aot autograd and torch inductor which is the code generator we'll go through a lot of inspecting the code like how you configure the system to give you the debug output how do you actually look at the debug output how do you understand the trace artifacts and then we'll get into actually how to debug the system based on these logs and then the final sections I mentioned uh profiling how you do good benchmarking how do you do good profiling and then I'll actually uh put up some Chrome traces for simplified examples so you can see what what I look at whenever I'm trying to understand a model performance so without further Ado the pytorch 2 component overview so as I mentioned uh pytorch 2's big feature is a compiler stack the API for this is um you know the aptly named torch.compile you compile your model you get an optimized model that you can run and it's hopefully faster than the one you passed in the compiler stack has I I'd say four main components that I've listed here in the boxes so the the first one Dynamo is kind of the front end I'll go into more detail on these in the following slides dynamos the front end it does autograd automatic differentiation for you and then hands off the the graph to torch inductor which is actually our code generator um sort of an intermediate level code generator and then it hands off to the device specific back end so the one we'll focus on in this talk is Triton which is the GPU code generator but we also have a CPU code generator that uses C plus something that's interesting in this stack I think like there's there's the old adage in computer science that you can solve every problem with another layer of indirection in compilers I think I'd modify that you can solve everything with an additional intermediate representation so obviously you write your pytorch model in Python the compiler stack really lowers it down all the way to machine code goes through the torch FX graph representation to the Dr IR which is sort of a tensor Ira in the middle and then to try in your C plus plus so that's kind of the stack at a high level let's go through the components in a bit more detail so Dynamo Dynamo there's a very interesting problem when you're trying to optimize a high level language like python especially something like pytorch we know how to optimize an ml graph this is actually not a new problem there have been ml compilers for years but we're not actually given a nice clean ml graph that looks like you might see uh in your you know textbooks we're given some python code and it turns out that it's really hard to extract the optimizable bits of an arbitrary Python program so Dynamo is the solution to that we're actually going to do bytecode analysis it uses um you know if you're familiar with uh with python proposals pep 523 we extract all the torchops from a program by analyzing the bytecode we put guards on the region we're going to optimize um so that basically means we're going to check what are the things that need to be true before we can execute the optimized graph and honestly this looks a lot like every single tracing jit you've ever used if you use V8 Java they all take this approach they guard on certain facts uh they check those guards and if they all pass you can jump into the optimized code and then the last thing that ensures soundness of the system is we allow graph breaks if the system hits something that it doesn't know how to trace through like arbitrary control flow or a call to uh you know another python Library it breaks the graph it can fall back to eager mode execution with the python interpreter and that way you get soundness without uh without giving up optimization next step aot autograd now the cool thing about this is that pytorch has always had this very flexible eager mode automatic differentiation the autograd engine kind of implicitly builds a graph as your pytorch programming is executing and then when you call backward uh it does the differentiation for you it updates the gradients this is really flexible but it's also really hard to optimize this is all very implicit a lot of it's dispatched through C plus plus behind the scenes your API is just backward and if that's slow it's hard to figure out why and it's hard for a compiler to take these disparate calls being launched from C plus plus and actually turn them into something fast so enter aot autograd which actually uses the existing autograd engine plus some clever tracing magic to build the backward graph and then we can optimize it just like we do the forwards graph the next component which is the actual optimizing compiler heart of this system is torch inductor and this is what generates code its main optimization is fusion um a lot of operators in ml graphs are memory bandwidth bound so if you can fuse them together you reduce memory bandwidth and improve performance but it's also got a lot of other optimizations it'll do memory planning it'll do tiling it also has an auto tuning engine that's really cool we've been experimenting a lot with that lately and seeing nice gains from Auto tuning various operators last component I want to talk about code generation I think one of the keys to the success of torch inductor is the low level compiler which is Triton this is a compiler actually developed at open AI we've been working closely with that team um and it provides a really great clean GPU programming model it lets you write fast uh fast kernels without a whole lot of the low-level details that you might kind of slog through if you're writing just Cuda so this is kind of the the secret sauce for how we get great performance on gpus on CPUs our back end is actually C plus plus uh the great thing about C plus plus is it's easy to read we can all read it we can see what it's doing and you actually get surprisingly good performance modern vectorizing compilers are quite good so we actually have good performance on the CPU backend as well I'll focus on the GPU in this talk since I think that's uh where a lot of the excitement is right now all right so with that kind of uh in in the back of our minds what the stack looks like how do we actually understand what this stack is doing it's great that it's there uh but let's peek under the hood um I wanted a really good quote about the importance of logging whenever you're uh digging into software um and so I figured I would ask chat GPT because I'm not too clever coming up with epigrams I kind of like uh logging Is Like Oxygen for debugging uh or logging is the microscope of software development um I thought those were pretty good I think I could maybe improve on them with some thought uh you know uh um but uh but chat GPT is doing me a favor here so yeah I have to say I like the logging is the heartbeat of your code there you go the heartbeat the soul yeah yeah it's very poetic thank you chat GPT all right so yeah the way I like to dig into a system is to just print out all sorts of stuff and see what it's doing uh and and we have a lot of really good logging options in the pt2 stack so we can I'll show you how to turn them all on uh and really get a sense for what's happening in your program so to do this walkthrough we're going to use a really really simple example this is just your like basic MLP a couple of linear layers a couple of relus uh I I've added a little Quirk to the activation function I I want this to show a speed up and as you might imagine a linear plus a relu there's not too much room to optimize so I've uh created this totally fake activation function where it's like relu squared that gives us something a fuse so we can actually see the compiler really in action and really changing something we'll write a simple Benchmark uh the details of this not super important we run it a bunch of times gather the time and we're getting a 14 speed up on the a100 GPU I'm using for benchmarking so that's not bad um how what's going on let's dig in so first thing I want to put this up front I'm going to show you all of the detailed knobs and kind of walk through them step by step but if you take nothing else away from this talk the best environment variable to use is torch compile debug set that in your environment it prints a nice Trace it's not too detailed you know it won't flood your logs but it gives you the information you need to debug your code I use this one all the time um let's go through the individual knobs though because I think there's a lot of good information here that we can dig into and now I can even start uh showing you some actual code so Dynamo logging uh there are several knobs we might want to see uh first of all we can just turn on a general logging we want to see the byte code and the guards that it's compiling and then we want to see graph breaks so these are kind of the three knobs we want let me pop over to my editor here um on the left you can see this simple MLP that I warned you was coming so linear layers relu squared which I'm sure is going to turn out to be a fantastic activation function just you wait um and and the benchmark so why don't we just first run this and I'll I'll show you uh whether or not I'm uh lying about the speed UPS dramatic pause yes hoping the the code runs and we have no problems exactly what I said okay all right so let's let's dig in with the logging knobs now uh to kind of help me along I've saved some of the knobs over here so you don't have to watch me fumble around typing them in I like to import all these logging modules uh with short names so that they're more convenient we'll just go ahead and get the Dynamo config Funk torch config which is where ao2 autograd lives and the inductor config module and let's turn on Dynamo logging so that's this block here uh logging print graph breaks output code and let's just see what this produces for us okay simple dot Pi here we go dramatic pause again I probably remove the Benchmark so that this can run a little bit faster maybe I'll do that all right wow okay we got a lot of information so let me back up to the very top I'll just maximize this so that we can see now okay I'm well aware that as you're watching this stream you're not going to be pouring over every single line of this very very detailed logging output but I do want to show some of the high level things that are happening okay so step one torch Dynamo tracing forward this is it walking through the byte code and figuring out what it's doing so that it can extract the pi torch operators so you see this long block of Trace some python byte code and there arguments that by code so that's it doing the the trace kind of coming down here it calls the compiler function enters torch inductor and now down here at the bottom we can actually see the graph that it traced so let me uh let me move that up to the top of the screen so it's a little easier to see um you can see okay uh it does FC layer one so your first linear layer calls a relu squares it second linear layer relu Square output okay pretty clean graph yes is this would you call this the stack Trace that we're looking through here of um the compile yeah to to some extent this is kind of like the high level like clean ml graph representation of the program this is like what is extracted from the bytecode you can see what it's done to the byte code actually by looking at this part there's the original bytecode um python bytecode is actually like pretty readable uh we don't really need to go through it it's turned it into a fairly simple function though that just loads up the compiled function calls it and then Returns the value so good job compiler you got the whole graph okay so I mentioned guards earlier this is this is one of the coolest sections to me um the secret sauce with Dynamo is that you really have to know that certain conditions hold before you can execute an optimized graph if various things have changed if you've made some weird edits to the uh to the NN module you're running if you've passed in tensors of different sizes or strides such that your generated code isn't going to work or isn't going to be fast we want to know that up front so that we can recompile a faster program if uh if we want or a more correct program even so in this case we don't need a ton of guards there are there are four there's the input you can see this tensor match function that's the thing that checks some of the fundamental properties of the tensor so you want to know its rank how many dimensions it has probably like the order of those Dimensions you know whether they're transposed or something like that um so that you can jump into a really good kernel you also need to know some things about the MN modules themselves uh it's I think guarding on the actual like object ID so that if you do something weird to the object it's going to say hey wait I just need to recompile this whole thing similarly the parameters uh self fc1 and self fc2 you've got to know whether there are changes to those um uh before you can jump into the optimized version but in the common case you're going to be executing this MLP and all of these guards are going to pass nothing weird is going on and so you can just jump straight to the optimized graph and not repeat any of this analysis so you just have these four checks and the optimize graph and then you generate some code and you can actually even see the output at this extremely extremely long string we'll take a we'll take a minute to look through code in a minute I won't jump into that just yet so okay that's that's Dynamo that gives you some idea of what's happening with Dynamo um if you have a much more complex model which you know I'm sure you will uh the output is going to be correspondingly complicated uh you can have lots and lots of guards uh some of the models I've looked at as we're bringing this up have just absolute pages of guards um it comes with the territory the more complicated your program is the more invariants it has to check if any of the guards fail you get some nice logging output that'll tell you okay your model is recompiling too much um you can dig in and actually figure out why so that's Dynamo the next thing I want to talk about is the aot auto grad pass so this is like digging into the backwards pass um the cool thing about digging into the backwards passes well I don't know about you all but it's been what 25 years I think since I took calculus I don't always remember how to take the derivatives of complicated functions ao2 autograd just you know it does it for me and I can actually look at its output and be like oh yeah that's how you take the derivative of layer Norm for example um the uh sorry I didn't actually go through the uh options um there are some config options that you can turn on here obviously you can just turn on the funk torch lugging to to the max with logging.debug you can print the debug graphs with true but actually the way I usually like to use this debugging info is through the environment if you set aot FX graphs that will dump out all of the information well not all of the information but a lot of the information you care about which is what does the forward pass look like what does the backward pass look like so why don't we try that first I'm going to just remove the Dynamo logging so it was that was that our second uh debugging environment variable yeah that's our second debugging environment variable okay there's I think there's only two in this talk uh but they're okay if you dig around in these config files you can find more debugging environment variables okay because uh someone was asking uh is there a cohesive list of all the debugging environment variables uh there there is a pretty decent one uh later in the talk I'll have a list I'll link to the uh Dynamo troubleshooting document which goes through a lot of these interesting ones you know what I actually forgot too that are in this talk there are two more okay so we have four of hanger yeah yeah they're doubling every few minutes in a minute I love it I definitely recommend uh the Dynamo troubleshooting dock uh Justin can put a link in the chat a little bit later um it's also in the talk the other thing is these um let me flip back to the talk these config files torch Dynamo torch Funk torch and uh torch inductor the configs are just like really valuable reading if you look through them you can see all the knobs they're available available there are a bunch it's pretty obvious usually which are the ones that you can turn on for uh good info things like log level or verbose um are kind of the interesting ones cool okay so flipping back to the simple MLP we can run with uh aot FX graphs did I remember to turn off the Dynamo logging yes I did I don't want to overwhelm everybody's all right aot FX graphs now does this work drama again in fact while this works I'm gonna reduce the drama by just having this call uh opt mod x dot sum whoops thought backward which is uh not a fantastic training Loop if you're uh if you're wondering all right so that way we'll we'll see the output without waiting for the Benchmark okay so if we look at the output let's go back up to the top here we can look at it dumps out the forward graph and the Ford graph like okay this is this is relatively verbose but it's also basically what we typed in for this uh simple MLP um you can see it's got line numbers here so for example like the first line where we do self FC dot prelo squared boils down to to this for the matrix multiplication um for whatever reason linear layers always start with permute of the weights um but then it's just a matrix multiplication this a10. mm thing is Matrix multiply it plus addition of bias um then does the value squared which is a10.relu a10.pal so far so good uh next matrix multiplication at an m relu.pal uh permute it's actually interesting this this permute here is actually uh I think again it come from the backwards graph it partitions things differently depending on how uh what's optimal for performance so that's a little uh little internal Quirk right there and then we can see the backward graph I won't go line by line through this because again you'll see my knowledge of calculus not being as up to date as it should be um glad I'm not the only one so you can kind of see the basic structure of the code though um you know going backwards you have uh actually this comes from the forward pass this is another partitioning Quirk but then this comes from the backward pass so you have um the derivative of uh of x squared is 2 times x so this is kind of funny you can see it uh it's it's clearly applying an unoptimized uh form of x to the first uh we could optimize that away but you know it's in the graph for right now I think inductor actually gets rid of it um and then multiply by two so we can actually see that derivative and you can kind of keep going on from there like this uh less than or equal to uh followed by where is actually the gradient for relu so it's figuring out which things are in the the positive regime and applying the updates only to those same thing matrix multiplication the sum is for bias uh we've got the uh the power and the relu backwards here and then the um matrix multiplication nation and uh and bias sum right here so so there you go backwards pass um it's you know verbose but it gives you all the detail uh that you want and you actually can finally kind of understand what the backwards pass looks like which is pretty cool so now okay we've got forward backwards pass these are high level graphs we're pretty comfortable with those and actually go into the code generation so this is fun uh inductor is the code generator there are some options here for dumping the generated code let's turn those on and start looking at what we get so if I flip over here let's grab my secret logging commands commands yeah exactly very very secret so secret that we published them all right I should probably turn off the aot FX graph so we can uh get a cleaner output so there we go still takes a little while to compile um that's definitely like one of the things to be mindful of when you're running a compiler is sometimes it can be slower but it's not too bad we actually had compile time as a first order goal whenever we were building this um so we paid a lot of attention uh to making sure that models aren't going to take forever to compile uh we kind of scaled back some of the auto tuning and so forth so that you get a result pretty quickly are there any are there any cases where like like where you can get tips on what to avoid to make it so it won't go like like I'm just trying to think of like ideas where it's like okay yeah you've got like thrashing or something that that prevents something from like that makes it look more less optimized so are there things to avoid um that could help like prevent unnecessary slowdown yes like in terms of compile time slow down yes yes there are um the the biggest thing to watch out for um there are kind of two categories of things to watch out for the biggest thing is probably uh just read compilations if you're failing guards and recompiling uh don't do that um it's sometimes not obvious and I'll get to later in the talk uh some tools for diagnosing uh recopolations but that's kind of the main thing um okay one thing you can do to avoid them let's say uh let's say you have um a model that has Dynamic shapes so you're passing in images of different sizes and so forth um and torch inductor can sometimes deal with ease and sometimes can't um if you can narrow down where what you compile to the part that doesn't have Dynamic shapes like usually things will kind of stabilize at some point then you can optimize that part without paying the cost of recompiling every time um the the other thing is auto tuning we've kind of set the defaults to be pretty reasonable um but there is a knob which I don't have in the talk but you can find online to this is actually going the other direction if you want better performance but you will pay for it uh the knob is called Max Auto Tune on Torch compile if you set max autotune on it can take a really long time it basically takes every matrix multiplication tries out sometimes somewhere between 12 and 100 different configurations to figure out which one is fastest so you might sit there for an hour but your your model will be really fast once it's compiled cool so yeah it's one of those things where you get the it takes a long time to build and compile but once it's compiled you've got a really highly tuned model that's going to be really well optimized exactly if you've got a model that you're training for you know weeks or months or whatever you know making it 50 faster could be well worth spending an hour at the beginning right yeah makes a lot of sense cool so I didn't mean to derail you no worries those are great questions if we jump back to look at the output of inductor here um basically when you turn on those knobs that I mentioned it it dumps the code that it generated and the funny thing about the generated code okay as a compiler writer uh I'm used to you know the input of the compiler being you know like C plus plus or whatever and the output being like Assembly Language something low level no no we take in Python and we produce python it's just faster python this is where the power of Triton comes in right like I'm a big fan of Triton I kind of can't overstate it I'm a bit of a fanboy when it comes to Triton Triton lets you write low-level Cuda kernels in Python so um I don't know if this uh example without a bunch of details is going to make sense but I'll just kind of briefly point out what a Triton kernel looks like so this thing here it's a it's a point wise operator that we're going to Jet with Triton um there's kind of a bunch of technical stuff here that I don't want to go through uh in a lot of detail but this is basically doing your forward pass it's relu squared you can kind of see it here maybe if even if you're not super familiar with Triton this where statement uh where is like a select so if uh if temp is less than zero it picks zero otherwise it picks Step Zero right um so that's that's doing your relu and then you know I think we can all read uh multiplication so that's that's temp one squared right um so you can actually kind of see this it's written in Python the magic of the Triton compilers it turns that python into into PTX it actually has a pipeline that goes through Triton Zone ir lvmir and then produces PTX on the back end and it's got a bunch of great optimization passes that do kind of all of the like mechanical but kind of tricky to get right optimizations that you need in in Cuda if we keep scrolling down here hopefully we'll see yeah see there's lots of code generated here I'm looking for a sum because in the backwards pass right we have to sum up the gradients to get the the bias gradient here we go so this is a reduction um and the cool thing about a reduction in Triton uh which you know a okay with with pt2 you don't actually have to care what a reduction looks like in Triton you just write your python code uh and and you get performance for free but even if you're using Triton directly it's pretty nice a reduction looks like you know this for Loop doing accumulations um it's uh it's it's very clean easy uh if if you actually want to do a reduction in Cuda um it's it's non-trivial there's a really good talk by Mark Harris from um gosh years and years ago at this point um they walked through it but it's you know it's a good 20 slides about how to do a fast reduction it's not something you just like type out uh and Triton will just do it for you it's very nice very nice is a uh compiler Target so okay we've gotten through all the Triton kernels now we actually get to the compiled model again this looks very mechanical uh you know nobody's ever going to write uh python variable names like Triton fused Le mobile one value one scale zero but you can kind of read it right right locating some buffers with empty striated you're launching a Triton kernel that writes into it and so forth and so this this call here is implementing your model I think that one is actually the backwards pass since I went looking for the reduction so that's the call for the backwards pass if I kind of scroll up we can see the forwards pass oops yeah forwards passes uh is this guy it's a little bit more concise see Adam M this fused permute power relu is uh is the first activation function there's the second activation function so that's what the actual generated code looks like now I mentioned whoa I mentioned that uh this uh that there were a lot of good trace artifacts that would be printed out by the magic compiler environment variable torch compiled debug so that environment variable uh also can is controlled by some of these inductor logging options if we go back uh to the to the code we can look at where it says debug Trace so if you just search for the the code output here there's model zero backward debug Trace and we can zoom in on what's in that directory so I have a few explanations over on the slide so I'm going to flip back to the slide to explain what these Trace outputs are so there's a debug log that's just the output of logging.debug showing what the you know compiler kind of thinks it's doing there's a readable form of the of the torch FX graph we kind of looked at that from the aot autograd logging so that should be familiar there's the inductor IR both before and after Fusion I'll show you what that looks like I find the IR a little bit hard to read it is like a detailed like mid-level compiler IR but if you want to understand in detail what inductor is doing that is a good place to look and then there's the output python code which we're just looking at a little bit um so let me just dig into these real quick so the readable FX graph this is kind of what we saw a few minutes ago with aot autograd it's called readable uh because it has some nice annotations like you see um on this this tensor relu one there's kind of a fake type annotation that says this is a float32 tensor that's uh 1024 by 1024. it's not a real type annotation this won't it's not valid python but it's useful to look at that so you can know what the the types and shapes of your tensor door there's also a runnable FX graph if you want to just like reproduce a result you can just grab this and and run it it's a normal python script um it's basically the same stuff so I won't go through it in detail it's a little lower level has a little bit less interesting or a little bit less high level information so Bert I had a question from the uh the audience yeah absolutely so uh day I was asking uh do the custom Ops with their own kernels work well in this flow yeah so that is a great question because we definitely have custom Ops inside meta they do um one thing you'll see if you why don't we jump to the generated code one thing you can see in the generated code is their direct calls to a10.mm which is a matrix multiply uh that's in kind of the the parlance of inductor it's an extern kernel so it's one that inductor doesn't really know how to generate code for but it can just call that a normal kernel and exactly if you've got some custom Ops that inductor doesn't understand it's fine with that it will just write them into the output stream it'll optimize everything around that uh and just and just call your kernel directly so that's a very good question awesome I promised I'd show some inductor IR this is definitely like getting into the low-level weeds it looks like a bit of a fusion of uh of the Triton code and the FX code I won't go through it in great detail but you know one thing I'll note um it's kind of unit of work are these uh are these buffers so uh you know computed buffer is a buffer that it's going to compute so it's like the memory location that the tensor lives in so there's this computed buffer um I think this is doing kind of the backwards pass of the relu of the relu squared sorry um this is the extern kernel I was talking about the extern kernel scheduler node as Matrix multiply I'm also lying a little bit when I say that in doctor who doesn't know how to generate code for Matrix multiply it does if you turn on that Max Auto Tune knob I mentioned before it's not like a a matrix multiply template they can drop in um and it will sure auto-tune those Matrix multiplies so okay lots of scheduler nodes lots of code in them I won't go through them all post Fusion you can kind of see you know bigger uh bigger computed buffers with more stuff in them um you can dig into the details if you're really interested okay so there's inductor and it's Trace artifacts I find this stuff really really useful to dig through oh in fact here's here's something really cool that I like about this output code so okay I've got I've got this um I've got this this file called output code let me pull it back up it is the compiled form of My Graph you'll notice there's a main function here um you can actually run this whoops sorry you can actually run this program if you have python output code it'll run the model it will measure performance it will print out performance so this is what's great about this is especially if you've extracted like a small region from a larger model that has a lot of setup or whatever if you dump this output code you can iterate really fast understand like okay are there are there bugs are there optimization opportunities so it's really useful to be able to just run these Snippets we've definitely use that all the time internally to optimize models so okay this graph ran in what is that six milliseconds cool good job graph all right so we went through the trace artifacts um remember I'll say one more time use torch compile debug here I'll leave I'll even show you I'll even show you what it is so torch compile debug and I'll get rid of the logging options in simple dot Pi so that I'm not polluting the output okay torch compile debug you get uh you get some of the information that I've been showing you um you know everything I just showed you really lets you see things in great detail but you get um you get just kind of paging through here you get the FX graphs your forward and backward graphs so you can look through them you get logging about like where compilation is happening you know is it compiling forwards is it compiling backwards and uh I think the really valuable thing you get these debug traces so that you can go dig into the generated code and you know the the naming convention I think should be like relatively clear the model forward model backward uh that's the forward and backward pass so go find those cool all right so the next section is kind of what if something goes wrong how do I debug this stack compilers are obviously more complicated than not having a compiler so we want to know how to uh how to find out what went wrong and importantly how to generate good bug reports we would love to help you all solve your problems and the better bug reports that we can get uh the easier it is so okay first of all there's a really really good troubleshooting guide online I've got a link here we can put in the live stream this troubleshooting guide goes through a lot of the stuff actually in more detail than I'm going to go through in the talk it shows you how to narrow down errors to one of the sub component shows you how to like Minify the example so your model you know might be able might be very large but you want to cut it down something smaller shows you how to like debug compile times and graph breaks excessive recopilations and and even numerical accuracy which I think is a really key one I'll go I'll go through a few of these uh in in this talk so the first one that I think is pretty useful is uh narrowing down an error to a particular sub-component so uh the the way to do this is that you can use this lower level API um you know we've uh We've hidden this all uh behind torch.compile but you can actually tell it to use different compiler back ends that omit different layers of the stack so maybe you just want to like apply Dynamo but then just run the graph eagerly that will ensure that the bug isn't in or or help you narrow it down to see if the bug is in the aot autograd or the um or inductor Cogen similarly aot eager does aot autograd but doesn't apply inductor inductor is equivalent to torch.compile there's also this kind of strange intermediate one called aot inductor debug where it decomposes all the Ops into their lower levels but doesn't actually generate code I haven't used that one as much but I dropped in here because I thought it could be potentially interesting um I'll show I'll show a slightly dangerous example here I'm going to insert a bug into our compiler uh instead of generating code for relu I'm gonna have it raise a runtime error just you know sometimes it happens you think you're typing ops.maxim and you type raise runtime error you know yes I did that yesterday right everybody uh so now now when I run simple.pi hopefully it fails we'll see drama is it gonna fail the opposite of what you raise the runtime error l-o-l okay so I mean let's be honest this stack Trace makes it okay maybe pretty obvious although I'm looking at this and I'm like if I didn't know Dynamo fairly well I'm like there's a lot of stuff in the stack Trace that doesn't necessarily mean anything yeah very overwhelming like what does all this stuff mean there's file names and light numbers and yeah we're in the stack that this happened and a lot of it is the in torch course so if you know like I happen to be a pie torch core developer but if you're not it's like I don't know I just downloaded this from the internet like what is going on here um you know you can kind of see the error okay raise runtime error they usually won't be this obvious um but if we want to actually kind of narrow that down uh like I mentioned instead of using torch.compile we can generate opt mod by calling well first let's import uh torch Dynamo directly as Dynamo and then we'll call Dynamo dot optimize we'll just we'll run the aot auto grad back end so aot eager run that on the module and now let's see what happens so that error I inserted was actually at the inductor level it's you know throwing an error while it's outputting Triton code so by cutting inductor out of the mix we should actually see this succeed okay yeah so we're eliminating different variables from the problem exactly um and there we go it ran correctly so you know obviously like you know a simple example like this it's you know not doing much for you in some sense but I have very much used this to narrow down bugs in complicated uh production models which is really really helpful so okay that lets you know everything's down the next thing that you might want to do is find minimal Repros so okay you have this gigantic model with as many as two linear layers and you're not sure like which part of the model induces the bug um you want a a minified Repro so we have a uh we have a minifier this is one of those environment variables I forgot torch Dynamo Repro after aot oops if we run this torch Dynamo Repro after equals aot run our example sorry this is going to spit out some information that tells us how we can cut down that model to like the smallest thing that actually fails so that's the error we'll have to scroll back up to where it says writing minified Repro 2 and gives you this path so let's let's look at that this program has your model and it runs it using this thing called the minifier and what the minifier does is it just kind of scrapes away layers of the graph until it finds the smallest thing that it possibly can that still reproduces the failure um you know you can run that script directly like this you know python blah blah blah minifier launcher and that will kick it off I'm not going to actually run it right now because it takes a little while it has to run the graph over and over again compile multiple times but I'll show you the output whenever the process is done it creates a little program called repro.pi that I generated last night and if you look at this NN module it's a lot smaller it is a single relu the error reproduces with a single relu which you know you would expect because that's where we inserted the error so it's it's actually working it's cut down the the graph to the point where you can and and then this makes for a great bug report you can say hey my model has a problem the smallest Repro you have for it is this and just paste this file to us on GitHub and that makes it really easy for us solve solve your problems I have a quick question is there any sort of like you know to help with narrowing down to the to the most minimal Repro like how like is there a um like a document that shows like what the flow is like you know first it does you know at autograd first and then torch Dynamo or like yeah is there yeah like to help like sort of Orient like okay where in the pipeline could this be failing yeah so the uh the troubleshooting dock Dimension which here I'll go ahead and just uh just kind of like show you all it looks like this Dynamo troubleshooting dock uh is it's really good it kind of walks through like the whole process that you might go through there's the section on diagnosing runtime errors where it's like okay this is how you can uh zoom in just a little bit oh yeah big monitor yeah all right it's like here here's how you can walk through the various backends here's how you can uh uh let's see where's the minifier in here so I'm just going to scroll down a bit and find it yeah there's a section of how you run the minifier like really everything that I'm telling you in this talk I learned from this document cool thank you yeah absolutely and if you um if you actually run the the minifier which I said I wasn't going to do I won't let it completely uh you can you can kind of see what it's doing the the output is a little bit opaque because I think generally you you don't necessarily care how it's slicing the graph but you can see It'll be like uh you know removing uh you know removing eight nodes or removing you know one output so see it gives you some info about the graph it's like okay it has 15 nodes five inputs it tries eliminating dead code that doesn't work so it tries some other approaches uh cutting off the suffix oh that one worked so it's able to cut it down from 15 to 10 nodes this is very automatic you don't really have to like worry about how it's doing this but that's kind of the process it's going through it's just taking that graph and just iteratively shrinking it until it can't shrink it anymore there's a bunch of other stuff that I think I don't have time to dig into and uh but you know you can learn about it from the troubleshooting page um graph breaks free compilations and accuracy uh accuracy debugging all really important if you're uh experiencing kind of bad performance it's usually one of these uh reasons okay accuracy isn't bad performance but it's a serious problem um the accuracy humidifier is really cool it does basically the same thing uh as the correctness minifier is just showing that actually compares the numerical outputs of the optimized graph to the unoptimized graph and within some tolerance it'll let you know what's the smallest graph where you don't get the same result that is pretty useful for for finding places where the compiler actually like could introduce subtle numerical bugs um you know I I have generally not been a ml modeling engineer by trade but I have done just enough debugging of it that I am well aware of the debugging accuracy is very very tough so having good tools for doing that is is a real life saver last section I want to get to is performance I think we'll actually have enough time to go through some of uh some of this material the first thing I want to say is benchmarking this is the first way people start to uh uh to analyze performance there are so many ways that benchmarking can go awry uh whenever I started here let me let me stop this um whenever I started working on on this program I was actually uh I was actually using a much simpler example that didn't have the relu squared and I was seeing you know a 25 speed up from inductor I was like wow this is pretty cool this model shouldn't have anything to compile but it's making it faster it turned out it was a problem with the clock rate so I wrote this little script based on nvidia's guidance to lock the clocks um if you'd like this is like a safe thing for an a100 it is 1095 megahertz um if you don't do that you can get weird results where it looks like one model is just way faster than the other because the clock happens to run uh run quicker during its benchmarking um there's a whole Nvidia presentation on how to do this well uh I won't go into that but there are other things too uh if you forget to use tensor course I think we print a warning in this case but not using tensor cores is just like a basic mistake you have like an order of magnitude more flops if you use tensor cores um Cuda fan not clearing the GPU cache you know there are really just so many ways my recommendation is Triton has a really nice function in its testing Library called do bench that handles a bunch of this stuff for you it won't lock the clocks but it does like use Cuda event properly it clears the GPU cache um it does all the stuff as you wan