Repository: pytorch/torchdynamo

Issue:
TorchBench training error: CUDA: Error- invalid value : speech_transformer
Repro:
```
benchmarks/torchbench.py -d cuda --inductor --training --float32 --use-eval-mode --quiet --no-skip -k  speech_transformer
```
Error:
```
Traceback (most recent call last):
  File "/fsx/users/binbao/torchdynamo-tip/benchmarks/common.py", line 1234, in run_one_model
    new_result = optimized_model_iter_fn(model, example_inputs)
  File "/fsx/users/binbao/torchdynamo-tip/torchdynamo/eval_frame.py", line 166, in _fn
    return fn(*args, **kwargs)
  File "benchmarks/torchbench.py", line 335, in forward_and_backward_pass
    def forward_and_backward_pass(self, mod, inputs, collect_outputs=True):
  File "benchmarks/torchbench.py", line 336, in <graph break in forward_and_backward_pass>
    cloned_inputs = clone_inputs(inputs)
  File "benchmarks/torchbench.py", line 339, in <graph break in forward_and_backward_pass>
    pred = mod(*cloned_inputs)
  File "/fsx/users/binbao/pytorch-release/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/fsx/users/binbao/torchbenchmark/torchbenchmark/models/speech_transformer/speech_transformer/transformer/transformer.py", line 21, in forward
    def forward(self, padded_input, input_lengths, padded_target):
  File "/fsx/users/binbao/pytorch-release/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/fsx/users/binbao/torchbenchmark/torchbenchmark/models/speech_transformer/speech_transformer/transformer/encoder.py", line 36, in forward
    def forward(self, padded_input, input_lengths, return_attns=False):
  File "/fsx/users/binbao/torchbenchmark/torchbenchmark/models/speech_transformer/speech_transformer/transformer/encoder.py", line 48, in <graph break in forward>
    non_pad_mask = get_non_pad_mask(padded_input, input_lengths=input_lengths)
  File "/fsx/users/binbao/torchbenchmark/torchbenchmark/models/speech_transformer/speech_transformer/transformer/encoder.py", line 50, in <graph break in forward>
    slf_attn_mask = get_attn_pad_mask(padded_input, input_lengths, length)
  File "/fsx/users/binbao/torchdynamo-tip/torchdynamo/eval_frame.py", line 166, in _fn
    return fn(*args, **kwargs)
  File "/fsx/users/binbao/pytorch-release/functorch/functorch/_src/aot_autograd.py", line 889, in forward
    return compiled_f(
  File "/fsx/users/binbao/pytorch-release/functorch/functorch/_src/aot_autograd.py", line 880, in new_func
    return compiled_fn(args)
  File "/fsx/users/binbao/pytorch-release/functorch/functorch/_src/aot_autograd.py", line 228, in g
    return f(*args)
  File "/fsx/users/binbao/torchdynamo-tip/torchdynamo/eval_frame.py", line 166, in _fn
    return fn(*args, **kwargs)
  File "/fsx/users/binbao/pytorch-release/functorch/functorch/_src/aot_autograd.py", line 353, in forward
    fw_outs = call_func_with_args(
  File "/fsx/users/binbao/pytorch-release/functorch/functorch/_src/aot_autograd.py", line 250, in call_func_with_args
    out = normalize_as_list(f(args))
  File "/fsx/users/binbao/pytorch-release/functorch/functorch/_src/aot_autograd.py", line 228, in g
    return f(*args)
  File "/tmp/torchinductor_binbao/rd/crdsohhhujfn4p3xi2te6vrpge74gtietpjpk2y7lzy3xx62sqo7.py", line 380, in call
    kernel2[grid(16320)](primals_102, buf10, primals_105, buf13, 16320, 204)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 999, in __call__
    return self.kernel(*wargs, **kwargs, grid=self.grid)
  File "/fsx/users/binbao/torchdynamo-tip/torchinductor/triton_ops/autotune.py", line 50, in __call__
    timings = {
  File "/fsx/users/binbao/torchdynamo-tip/torchinductor/triton_ops/autotune.py", line 51, in <dictcomp>
    config: self._bench(*args, config=config, **kwargs)
  File "/fsx/users/binbao/torchdynamo-tip/torchinductor/triton_ops/autotune.py", line 42, in _bench
    return super()._bench(*args, config=config, **kwargs)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 1054, in _bench
    return triton.testing.do_bench(kernel_call)
  File "/fsx/users/binbao/triton/python/triton/testing.py", line 137, in do_bench
    fn()
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 1053, in kernel_call
    self.kernel(*args, num_warps=config.num_warps, num_stages=config.num_stages, **current)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 988, in __call__
    return _triton.runtime.launch(wargs, self.fn.do_not_specialize, cache_key, self.fn.arg_names,
RuntimeError: CUDA: Error- invalid value
```


Discussion:
The failure is caused by having CPU tensors as graph inputs, and they are coming from https://github.com/pytorch/benchmark/blob/e02e7ff0c15329e3590c2203698d892aed107054/torchbenchmark/models/speech_transformer/speech_transformer/transformer/attention.py#L23, so the scalars are turned into tensors at some step with device set as `CPU`.


I think this may be relevant to unspec variable, as it's a numpy float, we wrap it as tensor before passing to the FX graph. I can take a closer look tomorrow to confirm.  

Resolving patch:
diff --git a/torchinductor/compile_fx.py b/torchinductor/compile_fx.py
--- a/torchinductor/compile_fx.py
+++ b/torchinductor/compile_fx.py
@@ -61,6 +61,10 @@ def complex_memory_overlap(t):
     return torch._debug_has_internal_overlap(indexed_tensor) != 0
 
 
+def is_unspec_input(t):
+    return t.device.type == "cpu" and t.dim() == 0
+
+
 @DebugContext.wrap
 @no_dispatch()
 def compile_fx_inner(
@@ -134,6 +138,7 @@ def run(*new_inputs):
                 if isinstance(new_inputs, tuple):
                     new_inputs = list(new_inputs)
                 new_inputs[i] = clone_preserve_strides(new_inputs[i])
+        new_inputs = [x.to("cuda") if is_unspec_input(x) else x for x in new_inputs]
         return model(*new_inputs)
 
     return run
@@ -190,6 +195,9 @@ def static_input(x):
         return torch.as_strided(buffer, x.size(), x.stride())
 
     assert isinstance(inputs, (list, tuple))
+    # dynamo wraps unspec variable as 0 dim tensor on CPU, need to move to GPU explicitly
+    inputs = [x.to("cuda") if is_unspec_input(x) else x for x in inputs]
+
     static_inputs = [
         static_input(x) if idx not in static_input_idxs else x
         for idx, x in enumerate(inputs)
diff --git a/torchinductor/graph.py b/torchinductor/graph.py
--- a/torchinductor/graph.py
+++ b/torchinductor/graph.py
@@ -225,7 +225,8 @@ def placeholder(self, target, args, kwargs):
         )
         self.graph_inputs[target] = tensor
         self.graph_inputs_original[target] = tensor.data.data
-        self.device_types.add(example.device.type)
+        if example.dim() != 0:
+            self.device_types.add(example.device.type)
         return tensor
 
     def call_function(self, target, args, kwargs):


Test patch:
diff --git a/test/test_torchinductor.py b/test/test_torchinductor.py
--- a/test/test_torchinductor.py
+++ b/test/test_torchinductor.py
@@ -3427,6 +3427,30 @@ def f(x):
         out = compiled(torch.randn(12, 4, device=self.device))
         self.assertEqual(out[0].shape, (24, 2))
 
+    @requires_cuda()
+    @patch.object(config.triton, "cudagraphs", False)
+    def test_unspec_inputs(self):
+        def fn(x, y):
+            return x + y
+
+        inputs = (
+            rand_strided((2, 3), (3, 1), device="cuda"),
+            rand_strided((), (), device="cpu"),
+        )
+        self.assertTrue(same(fn(*inputs), inputs[0] + inputs[1]))
+
+    @requires_cuda()
+    @patch.object(config.triton, "cudagraphs", True)
+    def test_unspec_inputs_cudagraphs(self):
+        def fn(x, y):
+            return x + y
+
+        inputs = (
+            rand_strided((2, 3), (3, 1), device="cuda"),
+            rand_strided((), (), device="cpu"),
+        )
+        self.assertTrue(same(fn(*inputs), inputs[0] + inputs[1]))
+
 
 if HAS_CPU: