Repository: pytorch/torchdynamo

Issue:
KeyError: 'example_value' while trying to trace TransformerEncoderLayer
See this gist for a repro: https://gist.github.com/bertmaher/f581734cf80b967592ccdff6d513774b

The model here is more or less a copypaste of nn.TransformerEncoderLayer, but applying `torchdynamo.optimize` yields the following stack trace:
```
ERROR FROM offset=18 filename encoder.py 42 KeyError
ERROR FROM offset=20 filename encoder.py 36 KeyError
========== TorchDynamo Stack Trace ==========
Traceback (most recent call last):
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/convert_frame.py", line 163, in _convert_frame_assert
    code = transform_code_object(frame.f_code, transform)
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/bytecode_transformation.py", line 284, in transform_code_object
    transformations(instructions, code_options)
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/convert_frame.py", line 139, in transform
    tracer.run()
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/symbolic_convert.py", line 278, in run
    and self.step()
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/symbolic_convert.py", line 256, in step
    getattr(self, inst.opname)(inst)
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/symbolic_convert.py", line 142, in wrapper
    return inner_fn(self, inst)
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/symbolic_convert.py", line 556, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/symbolic_convert.py", line 195, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/variables/functions.py", line 182, in call_function
    return super().call_function(tx, args, kwargs)
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/variables/functions.py", line 64, in call_function
    return tx.inline_user_function_return(
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/symbolic_convert.py", line 227, in inline_user_function_return
    result = InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/symbolic_convert.py", line 1150, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/symbolic_convert.py", line 1192, in inline_call_
    tracer.run()
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/symbolic_convert.py", line 278, in run
    and self.step()
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/symbolic_convert.py", line 256, in step
    getattr(self, inst.opname)(inst)
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/symbolic_convert.py", line 142, in wrapper
    return inner_fn(self, inst)
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/symbolic_convert.py", line 605, in CALL_FUNCTION_KW
    self.call_function(fn, args, kwargs)
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/symbolic_convert.py", line 195, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/variables/nn_module.py", line 172, in call_function
    return variables.TensorVariable.create(
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/variables/tensor.py", line 122, in create
    TensorVariable.create(
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/variables/tensor.py", line 75, in create
    args, kwargs = cls.propagate_args_kwargs(proxy.node)
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/variables/tensor.py", line 60, in propagate_args_kwargs
    return torch.fx.node.map_arg((node.args, node.kwargs), visit)
  File "/fsx/users/bertrand/conda/envs/triton3/lib/python3.8/site-packages/torch/fx/node.py", line 598, in map_arg
    return map_aggregate(a, lambda x: fn(x) if isinstance(x, Node) else x)
  File "/fsx/users/bertrand/conda/envs/triton3/lib/python3.8/site-packages/torch/fx/node.py", line 606, in map_aggregate
    t = tuple(map_aggregate(elem, fn) for elem in a)
  File "/fsx/users/bertrand/conda/envs/triton3/lib/python3.8/site-packages/torch/fx/node.py", line 606, in <genexpr>
    t = tuple(map_aggregate(elem, fn) for elem in a)
  File "/fsx/users/bertrand/conda/envs/triton3/lib/python3.8/site-packages/torch/fx/node.py", line 606, in map_aggregate
    t = tuple(map_aggregate(elem, fn) for elem in a)
  File "/fsx/users/bertrand/conda/envs/triton3/lib/python3.8/site-packages/torch/fx/node.py", line 606, in <genexpr>
    t = tuple(map_aggregate(elem, fn) for elem in a)
  File "/fsx/users/bertrand/conda/envs/triton3/lib/python3.8/site-packages/torch/fx/node.py", line 616, in map_aggregate
    return fn(a)
  File "/fsx/users/bertrand/conda/envs/triton3/lib/python3.8/site-packages/torch/fx/node.py", line 598, in <lambda>
    return map_aggregate(a, lambda x: fn(x) if isinstance(x, Node) else x)
  File "/fsx/users/bertrand/torchdynamo/torchdynamo/variables/tensor.py", line 58, in visit
    return n.meta["example_value"]
KeyError: 'example_value'
```

It seems to have something to do with the MultiheadAttention module attr, but I got only far enough into debugging to realize that it was better to ask y'all.

cc @anj-s 


Discussion:
I can take a look at this tomorrow.

Resolving patch:
diff --git a/torchdynamo/variables/tensor.py b/torchdynamo/variables/tensor.py
--- a/torchdynamo/variables/tensor.py
+++ b/torchdynamo/variables/tensor.py
@@ -118,16 +118,22 @@ def create(cls, tx, proxy, example_value=None, nnmodule=None, **options):
         elif isinstance(example_value, (tuple, list)):
             unpacked = []
             for i, val in enumerate(example_value):
-                unpacked.append(
-                    TensorVariable.create(
-                        tx,
-                        proxy.tracer.create_proxy(
-                            "call_function", operator.getitem, (proxy, i), {}
-                        ),
-                        example_value=val,
-                        **options,
+                if val is None:
+                    # nn.MultiheadAttention() can return None, see issue #175
+                    unpacked.append(
+                        variables.ConstantVariable(None, **options),
+                    )
+                else:
+                    unpacked.append(
+                        TensorVariable.create(
+                            tx,
+                            proxy.tracer.create_proxy(
+                                "call_function", operator.getitem, (proxy, i), {}
+                            ),
+                            example_value=val,
+                            **options,
+                        )
                     )
-                )
             if istype(example_value, tuple):
                 return variables.TupleVariable(unpacked, **options)
             elif istype(example_value, (list, immutable_list)):


Test patch:
diff --git a/tests/test_repros.py b/tests/test_repros.py
--- a/tests/test_repros.py
+++ b/tests/test_repros.py
@@ -7,6 +7,7 @@
 from typing import List
 
 import torch
+from torch import nn
 from torch.nn import functional as F
 
 import torchdynamo.testing
@@ -659,6 +660,61 @@ def _get_sorted_bucket_idx_and_undo_sorted_bucket_idx(buckets):
     return sorted_bucket_idx, undo_sorted_bucket_idx
 
 
+class FeedForwardLayer(nn.Module):
+    def __init__(self, d_model, dim_feedforward, activation, dropout) -> None:
+        super(FeedForwardLayer, self).__init__()
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.activation = activation
+        self.dropout1 = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+        self.dropout2 = nn.Dropout(dropout)
+
+    def forward(self, x):
+        return self.dropout2(
+            self.linear2(self.dropout1(self.activation(self.linear1(x))))
+        )
+
+
+class TransformerEncoderLayer(nn.Module):
+    def __init__(
+        self,
+        d_model,
+        nhead,
+        dim_feedforward=2048,
+        dropout=0.1,
+        activation=nn.ReLU(),
+        layer_norm_eps=1e-5,
+    ):
+        super(TransformerEncoderLayer, self).__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        self.norm1 = nn.LayerNorm(d_model, eps=layer_norm_eps)
+        self.norm2 = nn.LayerNorm(d_model, eps=layer_norm_eps)
+        self.dropout = nn.Dropout(dropout)
+        self.ff_block = FeedForwardLayer(d_model, dim_feedforward, activation, dropout)
+
+    def forward(self, src, src_mask=None, src_key_padding_mask=None):
+        x = src
+        x = self.norm1(x + self._sa_block(x, src_mask, src_key_padding_mask))
+        x = self.norm2(x + self._ff_block(x))
+        return x
+
+    # self-attention block
+    def _sa_block(self, x, attn_mask, key_padding_mask):
+        x = self.self_attn(
+            x,
+            x,
+            x,
+            attn_mask=attn_mask,
+            key_padding_mask=key_padding_mask,
+            need_weights=False,
+        )[0]
+        return self.dropout(x)
+
+    # feed forward block
+    def _ff_block(self, x):
+        return self.ff_block(x)
+
+
 class ReproTests(torchdynamo.testing.TestCase):
     def test_do_paste_mask(self):
         torchdynamo.utils.counters.clear()
@@ -951,3 +1007,15 @@ def toy_example(a, b, v):
             )
         self.assertEqual(cnt.frame_count, 1)
         self.assertEqual(cnt.op_count, 4)
+
+    def test_issue175(self):
+        n_heads = 2
+        d_model = 64
+        model = TransformerEncoderLayer(d_model, n_heads)
+        inp = torch.randn(1, d_model)
+        cnt = torchdynamo.testing.CompileCounter()
+        with torchdynamo.optimize(cnt, nopython=True):
+            model(inp)
+            model(inp)
+        self.assertEqual(cnt.frame_count, 1)
+        self.assertEqual(cnt.op_count, 12)