Difference between the graph break reasons: `Dynamic control flow is not supported at the moment.` and `generic_jump TensorVariable()`

username: Abhishek_Ghosh
cooked: <p>I was trying to understand the reason behind graph breaks, where I came across certain graph breaks in models from <a href="https://github.com/pytorch/pytorch/tree/50efbb9f1e7111b4b6d5b8e9a6064ee9783930be/benchmarks/dynamo" rel="noopener nofollow ugc">PyTorch Benchmarks</a>.</p>
<p><code>export TORCH_LOGS="graph_breaks"</code></p>
<p>Case 1: <code>BartForCausalLM</code> (<code>huggingface.py</code>)<br>
Command:</p>
<pre><code class="lang-auto">pytorch$ ./benchmarks/dynamo/huggingface.py --performance --training --amp --backend=inductor --only=BartForCausalLM
</code></pre>
<p>I get the following:</p>
<pre><code class="lang-auto">torch/_dynamo/symbolic_convert.py:527] [2/0] [__graph_breaks] Graph break: from user code at:
torch/_dynamo/symbolic_convert.py:527] [2/0] [__graph_breaks]   File "/media/disk1/abhishek/pytorch/benchmarks/dynamo/huggingface.py", line 564, in torch_dynamo_resume_in_forward_and_backward_pass_at_562
torch/_dynamo/symbolic_convert.py:527] [2/0] [__graph_breaks]     pred = mod(**cloned_inputs)
torch/_dynamo/symbolic_convert.py:527] [2/0] [__graph_breaks]   File "/home/abhishek/pytorch-venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
torch/_dynamo/symbolic_convert.py:527] [2/0] [__graph_breaks]     return forward_call(*args, **kwargs)
torch/_dynamo/symbolic_convert.py:527] [2/0] [__graph_breaks]   File "/home/abhishek/pytorch-venv/lib/python3.10/site-packages/transformers/models/bart/modeling_bart.py", line 2245, in forward
torch/_dynamo/symbolic_convert.py:527] [2/0] [__graph_breaks]     outputs = self.model.decoder(
torch/_dynamo/symbolic_convert.py:527] [2/0] [__graph_breaks]   File "/home/abhishek/pytorch-benchmarks/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
torch/_dynamo/symbolic_convert.py:527] [2/0] [__graph_breaks]     return forward_call(*args, **kwargs)
torch/_dynamo/symbolic_convert.py:527] [2/0] [__graph_breaks]   File "/home/abhishek/pytorch-venv/lib/python3.10/site-packages/transformers/models/bart/modeling_bart.py", line 1451, in forward
torch/_dynamo/symbolic_convert.py:527] [2/0] [__graph_breaks]     if dropout_probability &lt; self.layerdrop:
...
...
torch/_dynamo/symbolic_convert.py:527] [2/0] [__graph_breaks] torch._dynamo.exc.UserError: Dynamic control flow is not supported at the moment. Please use functorch.experimental.control_flow.cond to explicitly capture the control flow. For more information about this error, see: https://pytorch.org/docs/main/generated/exportdb/index.html#cond-operands
</code></pre>
<p>So, it complains about the <code>if</code> condition here:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/huggingface/transformers/blob/c409cd81777fb27aadc043ed3d8339dbc020fb3b/src/transformers/models/bart/modeling_bart.py#L1360C1-L1360C57">
  <header class="source">

      <a href="https://github.com/huggingface/transformers/blob/c409cd81777fb27aadc043ed3d8339dbc020fb3b/src/transformers/models/bart/modeling_bart.py#L1360C1-L1360C57" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/huggingface/transformers/blob/c409cd81777fb27aadc043ed3d8339dbc020fb3b/src/transformers/models/bart/modeling_bart.py#L1360C1-L1360C57" target="_blank" rel="noopener nofollow ugc">huggingface/transformers/blob/c409cd81777fb27aadc043ed3d8339dbc020fb3b/src/transformers/models/bart/modeling_bart.py#L1360C1-L1360C57</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="1360" style="counter-reset: li-counter 1359 ;">
          <li>if dropout_probability &lt; self.layerdrop:</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>I tried to produce it with a minimalistic example:</p>
<pre><code class="lang-auto">@torch.compile(mode="reduce-overhead")
def f(x, y):
    z = torch.rand([])
    if z &gt; 0.5:
        return x + y
    else:
        return x - y
x = torch.tensor([1.0, 2.0], device="cuda")
y = torch.tensor([1.0, 2.0], device="cuda")
f(x,y)
</code></pre>
<p>The simple <code>export TORCH_LOGS="graph_breaks"</code> didn’t print out anything. So, using “+dynamo” I got:</p>
<pre><code class="lang-auto">torch/_dynamo/symbolic_convert.py:774] [5/0] [__trace_source]         z = torch.rand([])
torch/_dynamo/symbolic_convert.py:797] [5/0] [__trace_bytecode] TRACE LOAD_GLOBAL torch []
torch/_dynamo/symbolic_convert.py:797] [5/0] [__trace_bytecode] TRACE LOAD_ATTR rand [PythonModuleVariable(&lt;module 'torch' from '/media/abhishek/pytorch-venv/lib/python3.10/site-packages/torch/__init__.py'&gt;)]
torch/_dynamo/symbolic_convert.py:797] [5/0] [__trace_bytecode] TRACE BUILD_LIST 0 [TorchInGraphFunctionVariable(&lt;built-in method rand of type object at 0x7fe1fda66500&gt;)]
torch/_dynamo/symbolic_convert.py:797] [5/0] [__trace_bytecode] TRACE CALL_FUNCTION 1 [TorchInGraphFunctionVariable(&lt;built-in method rand of type object at 0x7fe1fda66500&gt;), ListVariable(length=0)]
torch/_dynamo/symbolic_convert.py:797] [5/0] [__trace_bytecode] TRACE STORE_FAST z [TensorVariable()]
torch/_dynamo/symbolic_convert.py:774] [5/0] [__trace_source] TRACE starts_line /tmp/ipykernel_1940666/3830298969.py:4 in f
torch/_dynamo/symbolic_convert.py:774] [5/0] [__trace_source]         if z &gt; 0.5:
torch/_dynamo/symbolic_convert.py:797] [5/0] [__trace_bytecode] TRACE LOAD_FAST z []
torch/_dynamo/symbolic_convert.py:797] [5/0] [__trace_bytecode] TRACE LOAD_CONST 0.5 [TensorVariable()]
torch/_dynamo/symbolic_convert.py:797] [5/0] [__trace_bytecode] TRACE COMPARE_OP &gt; [TensorVariable(), ConstantVariable()]
torch/_dynamo/symbolic_convert.py:797] [5/0] [__trace_bytecode] TRACE POP_JUMP_IF_FALSE 26 [TensorVariable()]
torch/_dynamo/symbolic_convert.py:322] [5/0] generic_jump triggered compile
torch/_dynamo/output_graph.py:971] [5/0] COMPILING GRAPH due to GraphCompileReason(reason='generic_jump TensorVariable()', user_stack=[&lt;FrameSummary file /tmp/ipykernel_1940666/3830298969.py, line 4 in f&gt;], graph_break=True)
</code></pre>
<p>Why this difference?</p>
<p>Similarly,</p>
<p>Case 2: <code>detectron2_fasterrcnn_r_101_c4</code> (<code>torchbench.py</code>)</p>
<pre><code class="lang-auto">pytorch$ ./benchmarks/dynamo/torchbench.py --performance --inference --amp --backend=inductor --only=detectron2_fasterrcnn_r_101_c4```
</code></pre>
<p>Gives out the following graph break:</p>
<pre><code class="lang-auto">torch/_dynamo/symbolic_convert.py:527] [8/0] [__graph_breaks] Graph break: from user code at:
torch/_dynamo/symbolic_convert.py:527] [8/0] [__graph_breaks]   File "/home/abhishek/pytorch-venv/lib/python3.10/site-packages/detectron2/modeling/meta_arch/rcnn.py", line 208, in torch_dynamo_resume_in_inference_at_203
torch/_dynamo/symbolic_convert.py:527] [8/0] [__graph_breaks]     proposals, _ = self.proposal_generator(images, features, None)
torch/_dynamo/symbolic_convert.py:527] [8/0] [__graph_breaks]   File "/home/abhishek/pytorch-venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
torch/_dynamo/symbolic_convert.py:527] [8/0] [__graph_breaks]     return forward_call(*args, **kwargs)
torch/_dynamo/symbolic_convert.py:527] [8/0] [__graph_breaks]   File "/home/abhishek/pytorch-venv/lib/python3.10/site-packages/detectron2/modeling/proposal_generator/rpn.py", line 477, in forward
torch/_dynamo/symbolic_convert.py:527] [8/0] [__graph_breaks]     proposals = self.predict_proposals(
torch/_dynamo/symbolic_convert.py:527] [8/0] [__graph_breaks]   File "/home/abhishek/pytorch-venv/lib/python3.10/site-packages/detectron2/modeling/proposal_generator/rpn.py", line 503, in predict_proposals
torch/_dynamo/symbolic_convert.py:527] [8/0] [__graph_breaks]     return find_top_rpn_proposals(
torch/_dynamo/symbolic_convert.py:527] [8/0] [__graph_breaks]   File "/home/abhishek/pytorch-venv/lib/python3.10/site-packages/detectron2/modeling/proposal_generator/proposal_utils.py", line 106, in find_top_rpn_proposals
torch/_dynamo/symbolic_convert.py:527] [8/0] [__graph_breaks]     if not valid_mask.all():
...
...
torch/_dynamo/symbolic_convert.py:527] [8/0] [__graph_breaks]     self.dispatch_table[inst.opcode](self, inst)
torch/_dynamo/symbolic_convert.py:527] [8/0] [__graph_breaks]   File "/home/abhishek/pytorch-venv/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 470, in inner
torch/_dynamo/symbolic_convert.py:527] [8/0] [__graph_breaks]     raise exc.UserError(
torch/_dynamo/symbolic_convert.py:527] [8/0] [__graph_breaks] torch._dynamo.exc.UserError: Dynamic control flow is not supported at the moment. Please use functorch.experimental.control_flow.cond to explicitly capture the control flow. For more information about this error, see: https://pytorch.org/docs/main/generated/exportdb/index.html#cond-operands
</code></pre>
<p>Actually, it complains of the following <code>if</code> condition:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/facebookresearch/detectron2/blob/5b72c27ae39f99db75d43f18fd1312e1ea934e60/detectron2/modeling/proposal_generator/proposal_utils.py#L106C1-L106C33">
  <header class="source">

      <a href="https://github.com/facebookresearch/detectron2/blob/5b72c27ae39f99db75d43f18fd1312e1ea934e60/detectron2/modeling/proposal_generator/proposal_utils.py#L106C1-L106C33" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/facebookresearch/detectron2/blob/5b72c27ae39f99db75d43f18fd1312e1ea934e60/detectron2/modeling/proposal_generator/proposal_utils.py#L106C1-L106C33" target="_blank" rel="noopener nofollow ugc">facebookresearch/detectron2/blob/5b72c27ae39f99db75d43f18fd1312e1ea934e60/detectron2/modeling/proposal_generator/proposal_utils.py#L106C1-L106C33</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="106" style="counter-reset: li-counter 105 ;">
          <li>if not valid_mask.all():</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>But when I tried to create a minimalistic reproduction:</p>
<pre><code class="lang-auto">@torch.compile(mode="reduce-overhead")
def f(x, y):
    if (x == y).all():
        return x + y 
    else:
        return x - y
x = torch.tensor([1.0, 2.0], device="cuda")
y = torch.tensor([1.0, 2.0], device="cuda")
f(x,y)
</code></pre>
<p>The situation is same as Case 1, “graph_breaks” option did not say anything but, “+dynamo” showed the following:</p>
<pre><code class="lang-auto">torch/_dynamo/symbolic_convert.py:774] [3/0] [__trace_source] TRACE starts_line /tmp/ipykernel_1940666/2411286153.py:3 in f
torch/_dynamo/symbolic_convert.py:774] [3/0] [__trace_source]         if (x == y).all():
torch/_dynamo/symbolic_convert.py:797] [3/0] [__trace_bytecode] TRACE LOAD_FAST x []
torch/_dynamo/symbolic_convert.py:797] [3/0] [__trace_bytecode] TRACE LOAD_FAST y [LazyVariableTracker()]
torch/_dynamo/symbolic_convert.py:797] [3/0] [__trace_bytecode] TRACE COMPARE_OP == [LazyVariableTracker(), LazyVariableTracker()]
torch/_dynamo/output_graph.py:2029] [3/0] create_graph_input L_x_ L['x']
torch/_dynamo/variables/builder.py:2268] [3/0] wrap_to_fake L['x'] (2,) StatefulSymbolicContext(dynamic_sizes=[&lt;DimDynamic.STATIC: 2&gt;], constraint_sizes=[None], view_base_context=None, tensor_source=LocalSource(local_name='x', cell_or_freevar=False), shape_env_to_source_to_symbol_cache={}) &lt;class 'torch.Tensor'&gt;
torch/_dynamo/output_graph.py:2029] [3/0] create_graph_input L_y_ L['y']
torch/_dynamo/variables/builder.py:2268] [3/0] wrap_to_fake L['y'] (2,) StatefulSymbolicContext(dynamic_sizes=[&lt;DimDynamic.STATIC: 2&gt;], constraint_sizes=[None], view_base_context=None, tensor_source=LocalSource(local_name='y', cell_or_freevar=False), shape_env_to_source_to_symbol_cache={}) &lt;class 'torch.Tensor'&gt;
torch/_dynamo/symbolic_convert.py:797] [3/0] [__trace_bytecode] TRACE LOAD_ATTR all [TensorVariable()]
torch/_dynamo/symbolic_convert.py:797] [3/0] [__trace_bytecode] TRACE CALL_FUNCTION 0 [GetAttrVariable()]
torch/_dynamo/symbolic_convert.py:797] [3/0] [__trace_bytecode] TRACE POP_JUMP_IF_FALSE 20 [TensorVariable()]
torch/_dynamo/symbolic_convert.py:322] [3/0] generic_jump triggered compile
torch/_dynamo/output_graph.py:971] [3/0] COMPILING GRAPH due to GraphCompileReason(reason='generic_jump TensorVariable()', user_stack=[&lt;FrameSummary file /tmp/ipykernel_1940666/2411286153.py, line 3 in f&gt;], graph_break=True)
</code></pre>
<p>Why this difference?</p>