Repository: pytorch/torchdynamo

Issue:
`UnboundLocalError: ___random_values_0' from hf_Bart
Repro:
```
./benchmarks/torchbench.py -dcuda -k hf_Bart
```

Output:
```
cuda eval  hf_Bart                            ERROR:root:unhandled error
Traceback (most recent call last):
  File "/home/jansel/torchdynamo/benchmarks/common.py", line 690, in run_one_model
    new_result = model_iter_fn(model, example_inputs)
  File "./benchmarks/torchbench.py", line 306, in forward_pass
    def forward_pass(self, mod, inputs, collect_outputs=True):
  File "/home/jansel/pytorch/torch/nn/modules/module.py", line 1186, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/jansel/torchbenchmark/torchbenchmark/util/framework/huggingface/model_factory.py", line 34, in forward
    def forward(self, input_ids, decoder_input_ids):
  File "/home/jansel/pytorch/torch/nn/modules/module.py", line 1186, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/jansel/conda/envs/torchdynamo/lib/python3.8/site-packages/transformers/models/bart/modeling_bart.py", line 1308, in forward
    outputs = self.model(
  File "/home/jansel/pytorch/torch/nn/modules/module.py", line 1186, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/jansel/conda/envs/torchdynamo/lib/python3.8/site-packages/transformers/models/bart/modeling_bart.py", line 1181, in forward
    encoder_outputs = self.encoder(
  File "/home/jansel/pytorch/torch/nn/modules/module.py", line 1186, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/jansel/conda/envs/torchdynamo/lib/python3.8/site-packages/transformers/models/bart/modeling_bart.py", line 706, in forward
    def forward(
  File "/home/jansel/conda/envs/torchdynamo/lib/python3.8/site-packages/transformers/models/bart/modeling_bart.py", line 706, in forward
    def forward(
UnboundLocalError: local variable '___random_values_0' referenced before assignment
```


Discussion:
@yanboliang this may be caused by #392
Thanks for catching this, I'll assign this to me and dig into this.

Resolving patch:
diff --git a/torchdynamo/codegen.py b/torchdynamo/codegen.py
--- a/torchdynamo/codegen.py
+++ b/torchdynamo/codegen.py
@@ -11,7 +11,6 @@
 
 from .bytecode_transformation import Instruction
 from .bytecode_transformation import create_instruction
-from .bytecode_transformation import unique_id
 from .exc import unimplemented
 from .source import AttrSource
 from .source import Source
@@ -36,13 +35,6 @@ def merge(self, other: VariableTracker):
         self.variable = self.variable.add_options(other)
 
 
-def _get_gen_rand_values_fn(random_calls):
-    def _gen_rand_values():
-        return [fn(*args, **kwargs) for fn, args, kwargs in random_calls]
-
-    return _gen_rand_values
-
-
 class PyCodegen(object):
     """
     Helper class uses for constructing Python bytecode
@@ -328,20 +320,6 @@ def make_call_generated_code(self, fn_name: str) -> List[Instruction]:
         self.extend_output(self.load_function_name(fn_name))
 
         graphargs = self.tx.output.graphargs
-
-        # to handle random calls
-        if len(self.tx.random_calls) > 0:
-            rand_fn_name = unique_id("__gen_rand_values")
-            rand_fn = torchdynamo.disable(_get_gen_rand_values_fn(self.tx.random_calls))
-            self.tx.output.install_global(rand_fn_name, rand_fn)
-            self.extend_output(self.load_function_name(rand_fn_name))
-            self.extend_output(
-                [
-                    create_instruction("CALL_FUNCTION", 0),
-                    self.create_store(self.tx.output.random_values_var),
-                ]
-            )
-
         for arg in graphargs:
             if arg.is_unspecialized:
                 self.extend_output(
diff --git a/torchdynamo/output_graph.py b/torchdynamo/output_graph.py
--- a/torchdynamo/output_graph.py
+++ b/torchdynamo/output_graph.py
@@ -38,6 +38,13 @@
 from .variables.tensor import UnspecializedPythonVariable
 
 
+def _get_gen_rand_values_fn(random_calls):
+    def _gen_rand_values():
+        return [fn(*args, **kwargs) for fn, args, kwargs in random_calls]
+
+    return _gen_rand_values
+
+
 class FakeRootModule(torch.nn.Module):
     """Trick the constructor of fx.GraphModule"""
 
@@ -281,6 +288,18 @@ def compile_subgraph(self, tx, partial_convert=False):
             self.side_effects.codegen_update_mutated(pass2)
 
             output = []
+            # to handle random calls
+            if len(tx.random_calls) > 0:
+                rand_fn_name = unique_id("__gen_rand_values")
+                rand_fn = torchdynamo.disable(_get_gen_rand_values_fn(tx.random_calls))
+                self.install_global(rand_fn_name, rand_fn)
+                output.extend(pass2.load_function_name(rand_fn_name))
+                output.extend(
+                    [
+                        create_instruction("CALL_FUNCTION", 0),
+                        pass2.create_store(tx.output.random_values_var),
+                    ]
+                )
             if count_calls(self.graph) != 0 or len(pass2.graph_outputs) != 0:
                 output.extend(
                     self.compile_and_call_fx_graph(tx, pass2.graph_output_vars(), root)


Test patch:
diff --git a/tests/test_misc.py b/tests/test_misc.py
--- a/tests/test_misc.py
+++ b/tests/test_misc.py
@@ -1677,7 +1677,9 @@ def test_unspecialized_primitive_variable4(self):
         def fn(x):
             r1 = random.random()
             y = x + random.uniform(10, 20)
-            r2 = random.randint(2, 18)
+            y.sum().item()
+            r2 = random.randint(2, 18)  # no graph output in this frame
+            y.sum().item()
             return y + r1, r2
 
         x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])