A general question regarding exported graphs using Dynamo

username: abdulazizab2
cooked: <p>Hello all,</p>
<p>I have been learning about Torch Compile and its components, I do remember that when Torch 2.0 was announced for the first time it addressed that many models can’t get its graph exported due the nature of graph breaks … etc.</p>
<p>The question is, when I optimize a function/nn.Module is there a way to get the graph/s exported into ONNX like using <code>torch.onnx.export</code>. Or is the intention of using dynamo to have the “optimized” code run using PyTorch frontend only ?</p>
<p>As it seems Dynamo will help if a <code>nn.Module</code> can’t be exported easily using <code>torch.onnx.export</code> and outputs the “exportable” subgraphs and points out which parts of the function/nn.Module caused graph break/s (I think with python frame evaluation thing?)</p>

username: marksaroufim
cooked: <p>You’re correct in that <code>torch.compile()</code> is not an export workflow it’s a JIT workflow. If you want the whole graph then you can run <code>torch._dynamo.export()</code> but that’s still relatively early days and not ready for large scale customer adoption</p>
<p>That said you can see exactly how the ONNX backend works here <a href="https://github.com/pytorch/pytorch/blob/main/torch/_dynamo/backends/onnxrt.py" class="inline-onebox" rel="noopener nofollow ugc">pytorch/onnxrt.py at main · pytorch/pytorch · GitHub</a> but a high level dynamo sends over 1-n subgraphs which are then exported using the existing ONNX exporter</p>

username: abdulazizab2
cooked: <p>Hi, thanks for your reply!</p>
<p>I understand now that <code>torch.compile</code> is not an export workflow.</p>
<p>I tried to use torch._dynamo.export() to check how it will work when having a graph break and didn’t work (there were calls to “unimplemented” in the error stack) so I guess it is still not stable as you mentioned its in its early days.</p>
<p>Error stack just to double check:</p>
<blockquote>
<p>Blockquote<br>
Traceback (most recent call last):<br>
File “dynamo_playground.py”, line 32, in <br>
torch._dynamo.export(model, example_inputs[0])<br>
File “/home/abdulaziz/miniconda3/envs/torch/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py”, line 777, in export<br>
result_traced = opt_f(*args, **kwargs)<br>
File “/home/abdulaziz/miniconda3/envs/torch/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py”, line 117, in <strong>call</strong><br>
return self.dynamo_ctx(self._orig_mod.<strong>call</strong>)(*args, **kwargs)<br>
File “/home/abdulaziz/miniconda3/envs/torch/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py”, line 253, in _fn<br>
return fn(*args, **kwargs)<br>
File “/home/abdulaziz/miniconda3/envs/torch/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 1501, in _call_impl<br>
return forward_call(*args, **kwargs)<br>
File “/home/abdulaziz/miniconda3/envs/torch/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py”, line 402, in catch_errors<br>
return callback(frame, cache_size, hooks, frame_state)<br>
File “/home/abdulaziz/miniconda3/envs/torch/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py”, line 117, in _fn<br>
return fn(*args, **kwargs)<br>
File “/home/abdulaziz/miniconda3/envs/torch/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py”, line 318, in _convert_frame_assert<br>
return _compile(<br>
File “/home/abdulaziz/miniconda3/envs/torch/lib/python3.8/site-packages/torch/_dynamo/utils.py”, line 169, in time_wrapper<br>
r = func(*args, **kwargs)<br>
File “/home/abdulaziz/miniconda3/envs/torch/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py”, line 386, in _compile<br>
out_code = transform_code_object(code, transform)<br>
File “/home/abdulaziz/miniconda3/envs/torch/lib/python3.8/site-packages/torch/_dynamo/bytecode_transformation.py”, line 683, in transform_code_object<br>
transformations(instructions, code_options)<br>
File “/home/abdulaziz/miniconda3/envs/torch/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py”, line 373, in transform<br>
tracer.run()<br>
File “/home/abdulaziz/miniconda3/envs/torch/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py”, line 1899, in run<br>
super().run()<br>
File “/home/abdulaziz/miniconda3/envs/torch/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py”, line 611, in run<br>
and self.step()<br>
File “/home/abdulaziz/miniconda3/envs/torch/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py”, line 571, in step<br>
getattr(self, inst.opname)(inst)<br>
File “/home/abdulaziz/miniconda3/envs/torch/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py”, line 351, in wrapper<br>
return inner_fn(self, inst)<br>
File “/home/abdulaziz/miniconda3/envs/torch/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py”, line 1022, in CALL_FUNCTION<br>
self.call_function(fn, args, {})<br>
File “/home/abdulaziz/miniconda3/envs/torch/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py”, line 502, in call_function<br>
self.push(fn.call_function(self, args, kwargs))<br>
File “/home/abdulaziz/miniconda3/envs/torch/lib/python3.8/site-packages/torch/_dynamo/variables/builtin.py”, line 596, in call_function<br>
return super().call_function(tx, args, kwargs)<br>
File “/home/abdulaziz/miniconda3/envs/torch/lib/python3.8/site-packages/torch/_dynamo/variables/base.py”, line 230, in call_function<br>
unimplemented(f"call_function {self} {args} {kwargs}")<br>
File “/home/abdulaziz/miniconda3/envs/torch/lib/python3.8/site-packages/torch/_dynamo/exc.py”, line 107, in unimplemented<br>
raise Unsupported(msg)<br>
torch._dynamo.exc.Unsupported: call_function BuiltinVariable(print) [ConstantVariable(str)] {}</p>
</blockquote>
<blockquote>
<p>from user code:<br>
File “dynamo_playground.py”, line 14, in forward<br>
print(“graph break”)<br>
You can suppress this exception and fall back to eager by setting:<br>
torch._dynamo.config.suppress_errors = True</p>
</blockquote>
<p>However, as a workaround I tried using:<br>
<code>explanation, out_guards, graphs, ops_per_graph, break_reasons, explanation_verbose = torch._dynamo.explain(model, example_input)</code><br>
and it worked well. Extracted two subgraphs and I was able to convert it to onnx using the link you provided for onnxrt backend in dynamo.</p>
<p>Since the two functions should at least return subgraphs, should I open an issue that <code>torch._dynamo.export</code> not working when receiving a graph break or that’s the intended behavior for now.</p>
<p>Thanks again for your answer !</p>

username: marksaroufim
cooked: <p>I believe that’s intended, a print is a graph break so export doesn’t know how to export that so you get an error. If you want a partial graphs then torch.compile is the way to go</p>

username: abdulazizab2
cooked: <p>Thanks for your answer. It is reasonable, but was confusing as why <code>torch._dynamo.explain</code> produces several subgraphs. I will experiment more with it.</p>