Fix FX graph stack traceback issues

[williamwen42]

Fixes #1167 and related bugs where the current translator object is not properly given to FX graph nodes.

Example: https://gist.github.com/williamwen42/1efa7ce86b7ac797e90a13356668d90d
(Note: running the above script before this change requires this modification to the code:

--- a/torchdynamo/output_graph.py
+++ b/torchdynamo/output_graph.py
@@ -373,7 +373,7 @@ class OutputGraph(fx.Tracer):
             # the call to tabulate can cause a lot of memory to be allocated
             if config.log_level <= logging.INFO:
                 log.info(
-                    f"TRACED GRAPH\n {name} {gm.forward.__code__.co_filename} {format_graph_tabular(gm.graph)}\n"
+                    f"TRACED GRAPH\n {name} {gm.forward.__code__.co_filename} {gm.print_readable()}\n"
                 )

Running after this change requires setting torchdynamo.config.output_graph_code = True).

Running python benchmarks/torchbench.py --only BERT_pytorch --performance --verbose now gives:

class GraphModule(torch.nn.Module):
    def forward(self, inputs_0_ : torch.Tensor, inputs_1_ : torch.Tensor):
        
        # Module stack: {'mod': 'BERT'}, File: /scratch/williamwen/work/torchbenchmark/torchbenchmark/models/BERT_pytorch/bert_pytorch/model/bert.py:40, code: mask = (x > 0).unsqueeze(1).repeat(1, x.size(1), 1).unsqueeze(1)
        gt = inputs_0_ > 0
        unsqueeze = gt.unsqueeze(1);  gt = None
        repeat = unsqueeze.repeat(1, 128, 1);  unsqueeze = None
        unsqueeze_1 = repeat.unsqueeze(1);  repeat = None
        
        # Module stack: {'mod': 'BERT', 'mod_embedding': 'BERTEmbedding', 'mod_embedding_token': 'TokenEmbedding'}, File: /scratch/williamwen/work/torchdynamo-env/lib/python3.8/site-packages/torch/nn/modules/sparse.py:161, code: input, self.weight, self.padding_idx, self.max_norm,
        mod_embedding_token_weight = self.mod_embedding_token_weight
        
        # Module stack: {'mod': 'BERT', 'mod_embedding': 'BERTEmbedding', 'mod_embedding_token': 'TokenEmbedding'}, File: /scratch/williamwen/work/torchdynamo-env/lib/python3.8/site-packages/torch/nn/modules/sparse.py:160, code: return F.embedding(
        embedding = torch.nn.functional.embedding(inputs_0_, mod_embedding_token_weight, 0, None, 2.0, False, False);  inputs_0_ = mod_embedding_token_weight = None

[mlazos]

what's the difference between this new stack of tx's that you're adding to OutputGraph and the existing stack (see tx.parent) Could that stack be used instead? Otherwise, LGTM.

[voznesenskym]

Lgtm

[voznesenskym]

Ah, you went with the stack after all :) glad it worked

[voznesenskym]

Nit: idiomatic for pop to return what it popped

[voznesenskym]

Hmm I wonder if we eventually need to do this for all unpacks? As in, make them all take range or provide a default?

[jansel]

I'm about to switchover to core, lets hold off on landing

[jansel]

We have migrated torchdynamo to torch._dynamo and will use the pytorch/pytorch repo for future development. Please resubmit this PR to https://github.com/pytorch/pytorch/

More details and instructions to port this PR over can be found in #1588

[williamwen42]

what's the difference between this new stack of tx's that you're adding to OutputGraph and the existing stack (see tx.parent) Could that stack be used instead? Otherwise, LGTM.

The output graph only keeps a reference to a root (first) tx, not the most current one. The tx stack replaces the current_tx kwarg, since passing around a translator object as a kwarg is quite prone to errors.