Draft Export - PyTorch Compiler Series Episode 3

Key Takeaways

All right. So, today we have Angela who works on Torch export and she's going to talk about this new tool that we have called draft export. Um, just to give some context um export is a tool that that uh gives you um graphs from from models um using uh using the PyTorch compiler. Um however it's an ahead of time um graph capturing mechanism and you can get errors. So the normal workflow of export works by basically giving you one error at a time. But a lot of users asked us whether they can get all errors at once and somehow managed to get a full graph. So that's where draft export comes in and uh I'll let Angela talk all about it. Yeah, thanks Eic. So yeah, I'm going to be talking about draft export. And so uh example that EIC was talking about. So here I have this um piece of code that it's just a dummy dummy function. But let's say I have this model and then we have this like custom operator that does like pre-processing on it. And um some reasons why it might be a custom operator is because it could either contain some like C++ logic or it contains some like other mysterious logic that like you don't want to include inside of the model itself. You just want to hide it away in a black box. Um and then afterwards we have some like other random piece of code in the model that I just put together so that it will error correctly. Um so if we run this piece of code, it gives us you know some some tensor values. Um but if I try to export, so this is the API for exporting. It takes in the model and then the sample inputs and a dynamic shape specification saying whether or not like we want some shapes to be dynamic. And in this case, I've marked the zeroth dimension of the first input to be dynamic. So in this case we have size six but I want to say that in other cases we could all be also passing in values of size like seven eight or nine and so this is why I marked it as dynamic. Um but the for the second value I don't want it to be dynamic so I just marked it as none. And so maybe you like came across this model and you want to try to apply export on it. And so if we try to run it, uh, oh, I didn't say the first thing that happens is no fake, no fake imple. Um, so what this is, so for every custom operator, we must have this thing called a fake implementation uh registered for. This is because while exporting, we use this thing called fake tensors to help us trace through. And what it does is just records the shapes and dypes that go through the method. Um, and so in this case we to tell us how to get through pre-proc with fake tensors, we must have like a fake tensor implementation. Um so let's say that we add it register fake my preproc and um yeah so our fake tensor implementation can be exactly the same as the original implementation because all it needs is tell to tell us like how do we propagate the shapes and d types but we could also just return like a dummy tensor with like the same shape. Um but in this case like it works if we just return that the same thing. So now that we have the fake Oh sorry. So now that we have the fake implementation, we run it again and we have this new issue which says canal guard on data dependent expression uh u 0 / 3 + 5. And if we make this error a little bit bigger, we can see that it occurs on uh test.py on line 22. Um what else does it tell us? And I guess I guess it doesn't really based on oh based on the sax trace we can see that it's going into this slice function um inside of pietorch and there's this guard here that checks whether or not end val is less than zero. Um and so we can kind of uh think about it and see that this is checking to see that if whether or not this value a is uh greater than or equal to zero. Um because when we slice we can slice with negative numbers or we can slice with positive numbers and it could and the behavior of the uh slice operation could change depending on this. Um and so we just want to tell it that a is actually greater than zero. So here I will add a torch check that a uh so to do this I will actually add a torch check that a is a size value and this will automatically tell it that a is being used as like the value for a size and that this is like greater than greater than equal to zero. So if we run it again um we can see a new one which says U0 it must be whether or not U0 is greater than 12. And so if we look at the stack trace a little more we can see okay there's some line here that says like statically non true something something 12 comes from the cat operation right before. Yes. Yeah. Okay. Yeah. So from just looking at this error, it's a little bit complicated and we can't really tell like where this expression is coming from, where the 12 is coming from. The still on the slice again. Um, and so it's like a little bit annoying, but actually the correct solution here is that we need to add a check to make sure that A is less than the shape of like the shape of Z. Um, and so after we run that, we actually run into this constraints violated error. Uh, which says that we mark seek length as dynamic, but our code specialized it to be a constant stick. Yeah, I was wondering about that because the moment you said you wanted to do dynamic and then I saw the assert in there. That's sneaky. Yeah. So I was trying to Yeah, I was trying to demonstrate that like sometimes inside of a code Yeah. Um Yeah. I mean I mean the the real life scenario is um modeling is that you you haven't really looked at the source code of the model or you in in general the assumptions in the model might be baked in pretty deep and you're just trying out you your mental model says that something should be dynamic but you don't know whether it's going to work or not. Yeah. So if we comment this out cuz maybe like this assumption isn't true or like we don't actually want this assumption then we can export the model and yay it gives us an exported program and so we see the call to pre-proc the mole or sorry the divide and we see the cat and the slice. So, yay, it worked. But that was like a pretty tiring experience, right? Like we had to write a fake kernel. We had to add these we had to understand all of the guards and then add these checks and then like also like delete the assert. Um, so in this case like the assert was like pretty easy to figure out, but in some cases was like a lot harder. And so in some cases users really just want the graph itself. Like what if we're just trying this out? We just want the graph to play around with to see like to bring it to the rest of our workflow. And we don't really care about all of these correctness like whether or not like x of that.shaped is like actually equal to six or something like that. And so what if like we just really wanted a graph and so this is where draft export comes in. So draft export is a version of export which will always successfully export a graph and it will bypass any of like the previous export errors and instead these errors will be compiled into a report for a clear visualization which you can then later address maybe after you're done like bringing the graph through your pipeline. So what are the types of errors it catches? Here is a chart of like an estimate of torch expert errors. Um, this is not like an actual like representation of everything, but sort of like a compilation after like running through a bunch of internal models and looking at a bunch of GitHub issues. This is sort of like the estimate of what most people run into. Um, and so a lot of it is like missing a fake kernel guard on data dependent errors or like constraint violation errors. There are also some other errors such as like missing a PI tree registration, missing like or having an attribute assignment or some other random things. But ex draft export will be able to catch um these set of errors which are actually like 90% of most of the errors that um export users will will run into. And for the other ones we are looking into how to fix them but sometimes like we're not actually able to automatically fix them for you. So they might result in some Yeah. Yeah. So just to be clear so if if you do get some of these other random errors uh what will happen? So draft export will fail to produce a graph in that case right? Yes. And it will error the same way that normal expert would error and it might require you to like actually fix the model or file an issue on GitHub. Yeah. So how does this work? So this is how a normal tracing looks like in in normal export. So we have the eager code here. It calls like pre-prog item all this stuff. And every time we call an operation the tracer will record this operator inside the graph. So when we call preproc this corresponds to like the myib.prepo call um calling item will go into item and then calling cat will return like a a10.cat into the graph and like slicing will have a slice in the graph but at the same time as a tracer runs we run it using fake tensors to help us record all of these notes and the fake tensor just stores um like the shape and d type of what this uh operator returns. So in the preproc case, we're returning something of shape s0. This means um like s0 is just like a symbolic size. And in the case where we have this y do item, we return this u0. Um u0 means this is unbacked meaning we don't actually know what the value behind it is. Um oh sorry I forgot to mention so s0 was the the shape of the input. uh but you but you mentioned it's dynamic and there therefore uh we we create a symbol for it except that we know that the backing value for for that is six but for you zero you don't know that. Yeah. Yeah. So fix tensors are just tensors that do not have any values behind it. So when we do item do item it's trying to inspect like a value inside of the tensor. But because we're running with fake tensors which don't have actual values we don't actually know what the value is there. So we have this user which doesn't have a backing behind it. maps. Um, and when we do like a divide by three in like we just store this entire expression of u0 / 3 and this plus 5 creates like this even longer expression. And then once we get to this slicing call, what happens is we see a check for whether or not this value is greater than zero and whether or not this value is less than the shape of Z. And so this will um have these this will then turn into checks for whether or not u0 is greater than zero or if this expression is like less than this other expression. And in this case since we don't know how to evaluate this expression we will throw a guard on data dependent error. Um, so I guess by the name it says there's this data dependent behavior of this U0 because it's dependent on the data we've provided and we couldn't figure out how to resolve this and so we've errored. Um, and so in the case of draft export what we do is this thing called real tensor tracing where every time if we run the values with fake tensors we also run it with real tensors that we took from the sample inputs. Um so now u0 will be evaluate will have now a fake backing behind it which is the sample impulse we pass in which is two and then we'll be able to evaluate like all the other expressions. And so when we get to this whether or not u0 is greater than equal to zero check we will just plug in the two that we had from our um sample inputs and then evaluate this to be true and then evaluate this other expression to be true. So in this case s0 already has a backing which is six and so this expression can also be evaluated and then we can continue running and actually get the graph. Cool. So so from what I understand here is that the the graph would still be generalized with respect to a zero. Uh but uh now you're like specializing the graph to the like the the case where y dot item was exactly two. Um yeah. So if in our sample inputs we passed in negative one then this expression would evaluate to false and then we would take a different branch down the graph. Right. Right. So um in the case where so in the case where we reach a guard that we cannot evaluate we just use the fake tensors behind it to help us evaluate it and that helps us bypass all the guard on data dependent errors. But for the missing fake tensor errors, what happens is if we have an operator that doesn't have a fake implementation, then we will just run that operator with the real tensors yet the output tensor from that and then return in our fake tensor implementation another tensor that has the same rank but with unback shapes. So we just say like uh we know that this tensor will be of rank two but we don't actually know what the shape is um because it could be dependent on what your inputs are. So we just return unback shapes again. So the graph that gets created by draft export is guaranteed to work on the given sample inputs that you've passed in and it's specialized on that sample input. So depending on which branch it takes based on those sample inputs and so it might not necessarily work on any other inputs. And to show that we also add in asserts into the graph to so that if you do run it with other inputs it will error and tell you that um this graph has not been made to work for these types of input. But but if you were lucky um the the searchs would still pass even if the the data was a little bit different. If you if you were not depending on on the exact values then it would still be fine. Yeah. Yeah. Yeah. Yeah. So then how does what does draft export look like? Well, what happens is we can just replace the export call with draft export and they have the same uh signature. Uh and so in the case where draft expert succeeds, we will return uh an exported program or I guess most of the time draft expert should succeed and then we will return an exported program. And in the case where there are no errors then the behavior will be exactly like export. But in the case where there are errors that draft expert has uh bypassed, we will see this terminal output that says like hey four issues were found and then please you can view the report of all the errors with this TL parse command or by printing it out if you're in like a Python uh setting. So, let's go back to our example here and let me comment out all of these fixes that we did. So, if we run this again, it should error with the the first error that we encountered, which is the no fake imple. And if we replace this with draft export um we now get a we now get a graph. Woo. And we can see in the output like 3 hours uh were found. Um here please run this command or we can directly print it out with this. And there was another issue that we found which was that an operator doesn't have a fake kernel but we'll get to that later. Um so if we print out ep.report this will contain a report of all of the issues that we encountered. Um so the first thing is we were missing the fake kernel which we bypassed and then we had a data dependent error for this uh assert x is equal to shape six. So in this case it's considered a data dependent error because um we replaced this fake implementation with another implementation that just returned uh tensors of unback shapes. And so when it checked the shape it was unable to determine whether or not this value is equal to six. Um so it's labeled as a data dependent error because of the fake implementation that was used to bypass the first fake implementation. Mhm. And then later we have another data dependent error for the slicing operation here. So the alternative is we can use TL parse. So while we're running but this is already like for the for the less than equal to zero or um that that error is already better. Uh the message error message is already better because it tells you yeah the locals and so on and what actual value. So a it tells you that a was that sim uh that that expression, right? Yes. Yeah. Yeah. So the error here also like shows a little bit more information like it clears up all of the errors like logs that we saw before to just extract out this the slicing to show that this was actually within the framework code. This is where it errored. And then here are some of like the locals information. So this expression corresponds to what a is here. Yeah. Yeah. Um so while we're draft exporting, draft export also returns this log file of um everything that has occurred. And so if we use TL parse um by just doing TL parse here like that. Uh actually so if we run TL parse this will generate an HTML report for us to view all of the errors in this nice HTML file. And in our case in internally we can actually get a nice HTML file uh like website for us. But in open source, you'll have to like load it up onto your Chrome page. So what the TL parse report looks like is something like this. Um, which lists out all the errors that occurred. And for errors like this data dependent error, we can then click more into it and it'll show us a little bit more information than what the print of the report did. Specifically, it'll show us this provenence information which tells us where each value came from. So this u0 it shows that it came from the do item column and the u0 divided by 3 came from this divide by three operation and this five came from here and then finally this like less than check came from the slicing code here. So hopefully this should give you like a better sense of how um the guard was created so that you can help add checks for this. Yeah, this is cool because uh you know most of the time when when you when you see guards, internal guards being dumped into the terminal, you have no idea how they were created. So this basically breaks it down. Hey, I I had a quick question. Um so what what happened to the greater than 12 or less than equal to 12 thing? Yeah. So all of that is within this one slice operation. And so there's like multiple guards that are occurring inside of the slice. And to, you know, not clump up the terminal with like too many things, we just merged all of the guards that occurred on the same line. Okay. I see. I see. I see. So like after fixing this error, if you run draft expert again, the 12 should show up. I see. Okay. Okay. Oh. So um what's next? Oh yeah. So another so another uh thing that we generate uh if we look at oh yeah going back to the graph we do see a bunch of asserts checking that this value must be greater than or equal to some value and less than or equal some other value. And those are the run-time asserts that we put into the graph to show uh to to show the assumptions that we made while draft exporting. Um okay. So now going to the second part which is with custom operators. Um while we're exporting draft exporting, we will also record a profile. Uh I'm getting kicked out of my room. Oh, okay. I'll go find another room. Okay. Okay. All right. So um at this while we're draft exporting we also record a bunch of profiles for all of the operators that do not contain meta kernels. Um and so what this looks like is or this is stored in ep.report.up profiles. Um, so what this looks like is it's a mapping of the custom operator that does not contain the prof uh that does not contain a meta kernel and a set of profiles for what are the types of inputs and outputs um that can occur. And so we can actually save this to a file and then load it up later. But we can load it up even. So we can save this profile to a YAML file and and it looks something like this. Um so it says for the op for the inputs to this pre-proc we found yeah we found a tensor of rank one and with d type seven seven I think corresponds to like float 32 and for this type of input which is one tensor of rank one with a d type of uh float 32 we will return it returns an output tensor of rank one and also d type of float 32 too. Um, and so with this profile, we can then use it to generate a fake tensor profile for a fake kernel for us. So that every time we see a tensor input of this profile, we will also return an output of this profile. Um, and this will allow us to generate fake tensor kernels for you so that you don't have to write it yourself. Um what this looks like is we can do load this YAML file and then with this unsafe fake kernel generator um we can then call export again and it should run successfully. Oh wait actually first we should probably fix those errors. Yeah. So let's let's fix these data dependent errors and in the CA if we run this code again. Oh sorry. So if we run expert again we uh without this registration here we see that there's no fake impul registration but if we load up the profile and then add this context this will based on this YAML file that we stored it will load up a bunch of profiles and then with this um decorator it will temporarily register fake kernels with this profile. Um so if we do this then it will register a fake profile. Oh I didn't print the exported program and then we now get a graph with export. So in this case if you are too lazy to write your own fake tensor implementation um with draft export you can just save the profile that it returns and then load it um and continue running and continue like exporting. Um so this profile is specific to what is encountered. So if we pass in a tensor of rank two now this profile will error because it has never seen the kernel output of type two. Um so to fix that we can either like manually append to this YAML file or we can call draft expert again with um a new profile and it will generate you a new profile here. So if if it encounters different shapes then what happens does it does it try to um make a kernel that has dynamic shapes in it or uh it'll error. Okay. Um but draft expert will u make try to make you a new kernel. So it so what what will happen will you will you generate multiple kernels that are all specialized to to their own inputs or it will it will generate one kernel that contains like multiple of these profiles. Okay. And then it'll just like case match based on the profiles. But so this pro uh this way of like generating fake kernels is not entirely accurate because if you have um data dependent behavior like oh if my rank if my input of rank one then returns something of like rank one * 5 like and then or sorry like if you have a value inside of the tensor and then you return a tensor of that shape that specific shape like or that rank like that could not work in this case cuz it like hardcodes It's based on like what rank you've seen and what rank what rank you see in the input and what rank you see in the output. All right, makes sense. Cool. [Music] Um yeah so for some future work um we'll continue to get feedback and export errors and try to improve the report with more information also to support more errors and we will also work on improving the slowness of draft export because it is running um the real tensors at the same time there is some slowness so that is something we'll continue to work on so please try it out and let us know what you think and file issues to GitHub Nice. All right. Uh so thanks Angela and yeah uh please uh I I eco Angela please try it out. This is pretty cool and a bunch of our internal users already getting a lot of joy out of it. All right we'll end here for today. Thanks. Bye-bye. Thank you.