Warning

This feature is not meant to be used in production and is designed to beused as a tool for debugging torch.export tracing errors.

What sort of errors does it catch?#

Draft-export helps to catch and debug the following errors:

• Guard on data-dependent errors

• Constraint violation errors

• Missing fake kernels

• Incorrectly written fake kernels

How does it work?#

In normal export, we will convert the sample inputs into FakeTensors and usethem to record operations and trace the program into a graph. Input tensorshapes that can change (which are marked throughdynamic_shapes), or valueswithin tensors (typically from an.item() call) will be represented as a symbolicshape (SymInt) instead of a concrete integer. However some issues may occurwhile tracing - we may run into guards that we cannot evaluate, like if we wantto check if some item in a tensor is greater than 0 (u0>=0). Since the tracerdoesn’t know anything about the value ofu0, it will throw a data-dependenterror. If the model uses a custom operator but a fake kernel hasn’t beendefined for it, then we will error withfake_tensor.UnsupportedOperatorExceptionbecause export doesn’t know how to apply this onFakeTensors. If a customoperator has a fake kernel implemented incorrectly, export will silently producean incorrect graph that doesn’t match the eager behavior.

To fix the above errors, draft-export usesreal tensor tracing to guide us onhow to proceed when tracing. As we trace the model with fake tensors, for everyoperation that happens on a fake tensor, draft-export will also run the operatoron stored real tensors which come from the example inputs passed to export. Thisallows us to address the above errors: When we reach a guard that we cannotevaluate, likeu0>=0, we will use the stored real tensor values toevaluate this guard. Runtime asserts will be added into the graph to ensure thatthe graph asserts the same guard that we assumed while tracing. If we run intoa custom operator without a fake kernel, we will run the operator’s normalkernel with the stored real tensors, and return a fake tensor with the same rankbut unbacked shapes. Since we have the real tensor output for every operation,we will compare this with the fake tensor output from the fake kernel. If thefake kernel is implemented incorrectly, we will then catch this behavior andgenerate a more correct fake kernel.

How can I use draft export?#

Let’s say you’re trying to export this piece of code:

classM(torch.nn.Module):defforward(self,x,y,z):res=torch.ops.mylib.foo2(x,y)a=res.item()a=-aa=a//3a=a+5z=torch.cat([z,z])torch._check_is_size(a)torch._check(a<z.shape[0])returnz[:a]inp=(torch.tensor(3),torch.tensor(4),torch.ones(3,3))ep=torch.export.export(M(),inp)

This runs into a “missing fake kernel” error formylib.foo2 and then aGuardOnDataDependentExpression because of the slicing ofz witha,an unbacked symint.

To calldraft-export, we can replace thetorch.export line with the following:

ep=torch.export.draft_export(M(),inp)

ep is a valid ExportedProgram which can now be passed through further environments!

Debugging with draft-export#

In the terminal output from draft-export, you should see the following message:

#########################################################################################WARNING: 2 issue(s) found during export, and it was not able to soundly produce a graph.To view the report of failures in an html page, please run the command:    `tlparse /tmp/export_angelayi/dedicated_log_torch_trace_axpofwe2.log --export`Or, you can view the errors in python by inspecting `print(ep._report)`.########################################################################################

Draft-export automatically dumps logs fortlparse. You can view the tracingerrors by usingprint(ep._report), or you can pass the logs intotlparseto generate an html report.

Running thetlparse command in the terminal will generate atlparseHTML report. Here is an example of thetlparse report:

Clicking into the Data Dependent Error, we will see the following page whichcontains information to help debug this error. Specifically, it contains:

• The stacktrace at which this error occurs

• A list of local variables and their shapes

• Information for how this guard was created

The returned Exported Program#

Because draft-export specializes on code paths based on the example inputs, theexported program resulting from draft-export is guaranteed to be runnable andreturn correct results forat least the given example inputs. Other inputs canwork, as long as they match the same guards that were taken when we weredraft-exporting.

For example, if we have a graph branching on if a value is greater than 5, if indraft-export our example inputs were greater than 5, then the returnedExportedProgram will specialize on that branch, and will assert that the valueis greater than 5. This means that the program will succeed if you pass inanother value greater than 5, but will fail if you pass in a value less than 5.This is more sound thantorch.jit.trace, which will silently specialize on thebranch. The proper way fortorch.export to support both branches would be torewrite the code usingtorch.cond, which will then capture both branches.

Because of the runtime assertions in the graph, the returned exported-program isalso retraceable withtorch.export ortorch.compile, with a minor addition inthe case where a custom operator is missing a fake kernel.

Generating Fake Kernels#

If a custom operator does not contain a fake implementation, currentlydraft-export will use the real-tensor propagation to get an output for theoperator and continue tracing. However, if we run the exported program with faketensors or retrace the exported model, we will still fail because there is stillno fake kernel implementation.

To address this, after draft-export, we will generate an operator profile foreach custom operator call that we encounter, and store this on the reportattached to the exported program:ep._report.op_profiles. Users can then use thecontext managertorch._library.fake_profile.unsafe_generate_fake_kernels togenerate and register a fake implementation based on these operator profiles.This way future fake tensor retracing will work.

The workflow would look something like:

classM(torch.nn.Module):defforward(self,a,b):res=torch.ops.mylib.foo(a,b)# no fake implreturnresep=draft_export(M(),(torch.ones(3,4),torch.ones(3,4)))withtorch._library.fake_profile.unsafe_generate_fake_kernels(ep._report.op_profiles):decomp=ep.run_decompositions()new_inp=(torch.ones(2,3,4),torch.ones(2,3,4),)# Save the profile to a yaml and check it into a codebasesave_op_profiles(ep._report.op_profiles,"op_profile.yaml")# Load the yamlloaded_op_profile=load_op_profiles("op_profile.yaml")

The operator profile is a dictionary mapping operator name to a set of profileswhich describe the input and outputs of the operator, and could be manuallywritten, saved into a yaml file, and checked into a codebase. Here’s an exampleof a profile formylib.foo.default:

"mylib.foo.default":{OpProfile(args_profile=(TensorMetadata(rank=2,dtype=torch.float32,device=torch.device("cpu"),layout=torch.strided,),TensorMetadata(rank=2,dtype=torch.float32,device=torch.device("cpu"),layout=torch.strided,),),out_profile=TensorMetadata(rank=2,dtype=torch.float32,device=torch.device("cpu"),layout=torch.strided,),)}

mylib.foo.default’s profile contains only one profile, which says that for 2input tensors of rank 2, dtypetorch.float32, devicecpu, we will returnone tensor of rank 2, dtypetorch.float32, and devicecpu. Using thecontext manager, will then generate a fake kernel where given 2 input tensors ofrank 2 (and the other tensor metadata), we will output one tensor of rank 2 (andthe other tensor metadata).

If the operator also supports other input ranks, then we can add the profile tothis list of profiles, either by manually adding it into the existing profile orrerunning draft-export with new inputs to get new profiles, so that thegenerated fake kernel will support more input types. Otherwise it will error.

Where to go from here?#

Now that we have successfully created anExportedProgram using draft-export,we can use further compilers such asAOTInductor to optimize its performanceand produce a runnable artifact. This optimized version can then be used fordeployment. In parallel, we can utilize the report generated by draft-export toidentify and fixtorch.export errors that were encountered so that theoriginal model can be directly traceable withtorch.export.

• torchao
• torchrec
• torchft
• TorchCodec
• torchvision
• ExecuTorch
• PyTorch on XLA Devices