Repository: pytorch/torchdynamo

Issue:
NameError - torch is not defined - speech_transformer
Repro -  `TORCHDYNAMO_DEBUG_FUNCTION=forward python benchmarks/torchbench.py --training --float32 --backend=eager --use-eval-mode --only=speech_transformer`

And a diff

~~~
diff --git a/torchdynamo/convert_frame.py b/torchdynamo/convert_frame.py
index 7fa490d..c3af426 100644
--- a/torchdynamo/convert_frame.py
+++ b/torchdynamo/convert_frame.py
@@ -232,6 +232,9 @@ def convert_frame_assert(compiler_fn: Callable, guard_export_fn=None, one_graph=
             and os.environ.get("TORCHDYNAMO_DEBUG_FUNCTION") != code.co_name
         ):
             return None
+
+        if "encoder" not in code.co_filename:
+            return None
         if code.co_name == "<genexpr>" and code.co_filename.endswith(
             ("transformers/file_utils.py", "transformers/utils/generic.py")
         ):
~~~


This one leads to error like 
~~~
NameError: name 'torch' is not defined
ERROR
name 'torch' is not defined
~~~

I do not fully understand how globals are handled during inlining. The file containing the parent function has import `import torch.nn as nn` (no `import torch`), but the child function file has `import torch`. It seems that children imports are not passed on the parent function while inlining.

cc @voznesenskym @jansel 




Discussion:
Repro confirmed. I have a theory here, looking into this. 
Yeah this needs instruction translation for imports, I think.

Just (hilariously) jamming in 
```
hack_output.append(create_instruction("LOAD_CONST", 2, 0))
hack_output.append(create_instruction("LOAD_CONST", 0, None))
hack_output.append(create_instruction("IMPORT_NAME",  0, "torch"))
hack_output.append(create_instruction("STORE_FAST", 1, "torch"))
instructions = hack_output + instructions
```

Gets me past this import error, onto another one:

```
ModuleNotFoundError: No module named 'torchbenchmark.models.speech_transformer.speech_transformer.transformer.get_non_pad_mask'
```

Obviously not a real fix - but demonstrative of the issue.

There's also another problem here - when we inline code the stack trace is wrong. 
I somehow jammed it in using https://github.com/pytorch/torchdynamo/pull/1128. But, idk.

Resolving patch:
diff --git a/torchdynamo/codegen.py b/torchdynamo/codegen.py
--- a/torchdynamo/codegen.py
+++ b/torchdynamo/codegen.py
@@ -315,6 +315,19 @@ def make_function_with_closure(
         output.extend(self.rot_n(num_on_stack + 1))
         self.clear_tos()
 
+    def create_load_python_module(self, mod):
+        """
+        Generate a LOAD_GLOBAL instruction to fetch a given python module.
+        """
+        root_globals = self.tx.output.root_globals
+        name = re.sub(r"^.*[.]", "", mod.__name__)
+        if root_globals.get(name, None) is mod:
+            return self.create_load_global(name, add=True)
+        mangled_name = f"___module_{name}_{id(mod)}"
+        if mangled_name not in root_globals:
+            self.tx.output.install_global(mangled_name, mod)
+        return self.create_load_global(mangled_name, add=True)
+
     def make_call_generated_code(self, fn_name: str) -> List[Instruction]:
         """Call the generated code function stored in fn_name"""
         self.extend_output(self.load_function_name(fn_name))
@@ -324,7 +337,7 @@ def make_call_generated_code(self, fn_name: str) -> List[Instruction]:
             if arg.is_unspecialized:
                 self.extend_output(
                     [
-                        self.create_load_global("torch", add=True),
+                        self.create_load_python_module(torch),
                         self.create_load_attr("tensor"),
                     ]
                 )


Test patch: