Repository: pytorch/torchdynamo

Issue:
hf_Bart `assert dst.data_ptr() == src.data_ptr()` error
@Chillee is this heuristic for static inputs still valid with your latest changes?
```
    def is_not_gradout(x):
        return "tangents" not in x.name
```


I am seeing the error below, which I believe means one of the "tangents" is changing when we didn't expect it to.


```
cuda train hf_Bart                            WARNING:root:using triton random, expect difference from eager
WARNING:root:using triton random, expect difference from eager
WARNING:root:using triton random, expect difference from eager
WARNING:root:using triton random, expect difference from eager
WARNING:root:using triton random, expect difference from eager
WARNING:root:using triton random, expect difference from eager
WARNING:root:using triton random, expect difference from eager
WARNING:root:using triton random, expect difference from eager
WARNING torchinductor.compile_fx: skipping cudagraphs due to multple devices
Similarity score=0.766252875328064
Accuracy failed for key name logits
INCORRECT
Traceback (most recent call last):
  File "./benchmarks/torchbench.py", line 390, in <module>
    main(TorchBenchmarkRunner(), original_dir)
  File "/home/jansel/torchdynamo/benchmarks/common.py", line 1547, in main
    runner.run_one_model(
  File "/home/jansel/torchdynamo/benchmarks/common.py", line 900, in run_one_model
    model_iter_fn(model, example_inputs)
  File "./benchmarks/torchbench.py", line 374, in forward_and_backward_pass
    def forward_and_backward_pass(self, mod, inputs, collect_outputs=True):
  File "./benchmarks/torchbench.py", line 374, in forward_and_backward_pass
    def forward_and_backward_pass(self, mod, inputs, collect_outputs=True):
  File "./benchmarks/torchbench.py", line 380, in forward_and_backward_pass
    self.grad_scaler.scale(loss).backward()
  File "/home/jansel/pytorch/functorch/functorch/_src/monkey_patching.py", line 77, in _backward
    return _old_backward(*args, **kwargs)
  File "/home/jansel/pytorch/torch/_tensor.py", line 484, in backward
    torch.autograd.backward(
  File "/home/jansel/pytorch/torch/autograd/__init__.py", line 191, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/home/jansel/pytorch/torch/autograd/function.py", line 265, in apply
    return user_fn(self, *args)
  File "/home/jansel/torchdynamo/torchdynamo/eval_frame.py", line 94, in _fn
    return fn(*args, **kwargs)
  File "/home/jansel/pytorch/functorch/functorch/_src/aot_autograd.py", line 271, in backward
    out = normalize_as_list(compiled_bw(*ctx.saved_tensors, *contiguous_args))
  File "/home/jansel/torchdynamo/torchdynamo/eval_frame.py", line 94, in _fn
    return fn(*args, **kwargs)
  File "/home/jansel/torchdynamo/torchinductor/compile_fx.py", line 205, in run
    assert dst.data_ptr() == src.data_ptr()
AssertionError

```


Discussion:
It should be, I'll take a look.
Looks like it might be an issue with inductor - if I disable the inductor compilation (but still use cudagraphs), there's no longer an error.
Digging into it a bit more, figured it out.

The issue is that for a pair of forwards + backwards graphs, the current cudagraph implementation assumes that the activations returned from the forwards pass are always in the same location. However, in this case, we fail to cudagraph the forwards pass but still cudagraph the backwards pass. So the backwards pass assumes that the activations are saved in the same place, but it isn't actually true. 

<img width="840" alt="image" src="https://user-images.githubusercontent.com/6355099/183777180-e437ba99-1155-499a-9e04-7a908d042322.png">


Resolving patch:
diff --git a/torchinductor/compile_fx.py b/torchinductor/compile_fx.py
--- a/torchinductor/compile_fx.py
+++ b/torchinductor/compile_fx.py
@@ -1,3 +1,4 @@
+import dataclasses
 import functools
 import itertools
 import logging
@@ -26,6 +27,14 @@
 log = logging.getLogger(__name__)
 
 
+@dataclasses.dataclass
+class BoxedBool:
+    value: bool
+
+    def __bool__(self):
+        return self.value
+
+
 class CheckEachNode(torch.fx.Interpreter):
     def call_function(self, target, args, kwargs):
         expected = target(*args, **kwargs)
@@ -157,10 +166,15 @@ def compile_fx_inner(
             return cudagraphify(
                 compiled_fn, example_inputs, static_input_idxs=range(num_fixed)
             )
-        elif cudagraphs and set(graph.device_types) == {"cuda"}:
-            log.warning("skipping cudagraphs due to input mutation")
-        elif cudagraphs and len(graph.device_types) > 1:
-            log.warning("skipping cudagraphs due to multple devices")
+        elif cudagraphs:
+            if isinstance(cudagraphs, BoxedBool):
+                # Disable cudagraphs in the backwards pass too:
+                cudagraphs.value = False
+
+            if set(graph.device_types) == {"cuda"}:
+                log.warning("skipping cudagraphs due to input mutation")
+            elif set(graph.device_types) != {"cpu"}:
+                log.warning("skipping cudagraphs due to multiple devices")
         return compiled_fn
     except Exception:
         if os.environ.get("TORCHINDUCTOR_DUMP_REPRO") == "1":
@@ -248,20 +262,25 @@ def compile_fx_aot(model_: torch.fx.GraphModule, example_inputs_: List[torch.Ten
         model_ = normalize_ir(model_, example_inputs_)
         model_ = overrides.replace_fx(model_)
     num_example_inputs = len(example_inputs_)
+    cudagraphs = BoxedBool(config.triton.cudagraphs)
 
     def fw_compiler(model: torch.fx.GraphModule, example_inputs):
         if config.debug:
             print("FORWARD GRAPH:")
             model.graph.print_tabular()
         fixed = len(example_inputs) - num_example_inputs
-        return compile_fx_inner(model, example_inputs, num_fixed=fixed)
+        return compile_fx_inner(
+            model, example_inputs, num_fixed=fixed, cudagraphs=cudagraphs
+        )
 
     def bw_compiler(model: torch.fx.GraphModule, example_inputs):
         if config.debug:
             print("BACKWARD GRAPH:")
             model.graph.print_tabular()
         fixed = count_tangents(model)
-        return compile_fx_inner(model, example_inputs, num_fixed=fixed)
+        return compile_fx_inner(
+            model, example_inputs, num_fixed=fixed, cudagraphs=cudagraphs
+        )
 
     with overrides.patch_functions():
         return aot_autograd(


Test patch: