Repository: pytorch/torchdynamo

Issue:
Investigate "skipping cudagraphs due to multple devices" issue
```
./benchmarks/torchbench.py -dcuda --float16 --inductor --training -k hf_Bart
```

```
cuda train hf_Bart                            WARNING:root:using triton random, expect difference from eager
WARNING torchinductor.compile_fx: skipping cudagraphs due to multple devices
WARNING torchinductor.compile_fx: skipping cudagraphs due to multple devices
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
WARNING torchinductor.compile_fx: skipping cudagraphs due to multple devices
WARNING:root:using triton random, expect difference from eager
WARNING torchinductor.compile_fx: skipping cudagraphs due to multple devices
WARNING torchinductor.compile_fx: skipping cudagraphs due to multple devices
WARNING:root:using triton random, expect difference from eager
WARNING torchinductor.compile_fx: skipping cudagraphs due to multple devices
WARNING torchinductor.compile_fx: skipping cudagraphs due to multple devices
WARNING:root:using triton random, expect difference from eager
WARNING torchinductor.compile_fx: skipping cudagraphs due to multple devices
WARNING torchinductor.compile_fx: skipping cudagraphs due to multple devices
WARNING:root:using triton random, expect difference from eager
WARNING torchinductor.compile_fx: skipping cudagraphs due to multple devices
WARNING torchinductor.compile_fx: skipping cudagraphs due to multple devices
WARNING:root:using triton random, expect difference from eager
WARNING torchinductor.compile_fx: skipping cudagraphs due to multple devices
WARNING torchinductor.compile_fx: skipping cudagraphs due to multple devices
WARNING:root:using triton random, expect difference from eager
```

These warnings about multiple devices are a bit odd, and may be slowing things down by preventing use of cudagraphs.
1) If these issues aren't real (perhaps due to constants which we can move at compile tile), we should enable cudagraphs in these cases and make the multi-device check less sensitive.
2) If these issues are real and the model truly needs CPU tensors, we should try to bake the host to device copy into the cudagraphs arg copy (which we need to do anyway).


Discussion:
By running one example, `python benchmarks/torchbench.py -d cuda --inductor  --float16 -k hf_Bart`, I found the issue mostly coming from factory ops, like `arange` and `full`. As the following FX graph fragment shows, when `arange` is called without `device` kwarg, it will use the default device which is `cpu` and then does a bunch of computation on `cpu`, and eventually calls `.to` to get a cuda tensor.

We can fix it by forcefully specifying the default device for factory ops, but the root problem is, why is this kind of FX graph is generated in the first place? @Chillee 

BTW, fixing this can bring `hf_Bart` inference perf from 1.07x to 1.33x.


```
__compiled_fn_36 <eval_with_key>.512
opcode         name                                                   target                                                        args                                                                                                                                                                                                                                                                                                                                                                               kwargs
-------------  -----------------------------------------------------  ------------------------------------------------------------  ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------  --------------------------------------------------------------------------------------------------------------------------------------------
placeholder    _stack0_last_hidden_state                              _stack0_last_hidden_state                                     ()                                                                                                                                                                                                                                                                                                                                                                                 {}
placeholder    decoder_input_ids                                      decoder_input_ids                                             ()                                                                                                                                                                                                                                                                                                                                                                                 {}
call_method    view                                                   view                                                          (decoder_input_ids, -1, 512)                                                                                                                                                                                                                                                                                                                                                       {}
get_attr       self_decoder_embed_tokens_weight                       self_decoder_embed_tokens.weight                              ()                                                                                                                                                                                                                                                                                                                                                                                 {}
call_function  embedding_1                                            <function embedding at 0x7ff52c053310>                        (view, self_decoder_embed_tokens_weight)                                                                                                                                                                                                                                                                                                                                           {'padding_idx': 1, 'max_norm': None, 'norm_type': 2.0, 'scale_grad_by_freq': False, 'sparse': False}
call_function  mul                                                    <built-in function mul>                                       (embedding_1, 1.0)                                                                                                                                                                                                                                                                                                                                                                 {}
call_function  tensor                                                 <built-in method tensor of type object at 0x7ff6845b89a0>     (-inf,)                                                                                                                                                                                                                                                                                                                                                                            {}
call_function  full                                                   <built-in method full of type object at 0x7ff6845b89a0>       ((512, 512), tensor)                                                                                                                                                                                                                                                                                                                                                               {}
**call_function  arange                                                 <built-in method arange of type object at 0x7ff6845b89a0>     (512,)                                                                                                                                                                                                                                                                                                                                                                             {}**
call_function  add                                                    <built-in function add>                                       (arange, 1)                                                                                                                                                                                                                                                                                                                                                                        {}
call_method    view_1                                                 view                                                          (add, 512, 1)                                                                                                                                                                                                                                                                                                                                                                      {}
call_function  lt                                                     <built-in function lt>                                        (arange, view_1)                                                                                                                                                                                                                                                                                                                                                                   {}
call_method    masked_fill_                                           masked_fill_                                                  (full, lt, 0)                                                                                                                                                                                                                                                                                                                                                                      {}
call_method    to                                                     to                                                            (full, torch.float16)                                                                                                                                                                                                                                                                                                                                                              {}
call_function  getitem                                                <built-in function getitem>                                   (to, (None, None, slice(None, None, None), slice(None, None, None)))                                                                                                                                                                                                                                                                                                               {}
call_method    expand                                                 expand                                                        (getitem, 8, 1, 512, 512)                                                                                                                                                                                                                                                                                                                                                          {}
call_method    to_1                                                   to                                                            (expand, device(type='cuda', index=0))                                                                                                                                                                                                                                                                                                                                             {}
call_function  arange_1                                               <built-in method arange of type object at 0x7ff6845b89a0>     (0, 512)                                                                                                                                                                                                                                                                                                                                                                           {'dtype': torch.int64, 'device': device(type='cuda', index=0)}
...
```
> We can fix it by forcefully specifying the default device for factory ops, but the root problem is, why is this kind of FX graph is generated in the first place? @Chillee

I guess users are writing bad code lol.
Allocating on CPU and copying to CUDA is a pretty common pattern. Assuming that inductor won't handle it directly, this is the sort of thing that an FX pass would be good for.
>  this is the sort of thing that an FX pass would be good for.

Hrmm... that's interesting. Should we be doing this?
I think there are two ways we could potentially model this. One is an FX pass which looks for tensors that are allocated on CPU, and then moved (sole use) to CUDA; they can be directly allocated to CUDA. The other is in inductor modeling itself; inductor has requirements on the devices of the inputs/outputs, but device allocation for all internal operations is flexible. So it can actually choose to ignore the device annotation on factory functions and make a decision when the inner buffer actually gets used.
As an optimizing compiler, I think inductor should have the freedom to do this kind of optimization. In the worst case, we can control the behavior with a config. 
I guess this makes sense to do at the FX graph level. cc: @SherlockNoMad 
Inductor should already handle this properly.  So no need to do at the FX level.

This optimization should be firing:
https://github.com/jansel/torchdynamo/blob/5123dd0a4832db32855ffbe2b8efe216803d11e1/torchinductor/ir.py#L2306


`constant_to_device()` is able to move a CPU `arange()` to the GPU without any real copies.

I think the issue is our heuristic to disable cudagraphs in this case is too sensitive.

Resolving patch:
diff --git a/torchinductor/ir.py b/torchinductor/ir.py
--- a/torchinductor/ir.py
+++ b/torchinductor/ir.py
@@ -2261,14 +2261,14 @@ def create(cls, a, b):
 class DeviceCopy(ExternKernelOut):
     @classmethod
     def create(cls, x, device):
-        V.graph.device_types.add(device.type)
-        V.graph.device_types.add(x.get_device().type)
-
         if not x.is_extern() and all(
             (r.name in V.graph.constants and hasattr(r, "index")) for r in x.get_reads()
         ):
             return x.constant_to_device(device)
 
+        V.graph.device_types.add(device.type)
+        V.graph.device_types.add(x.get_device().type)
+
         log.warning("DeviceCopy")
         return DeviceCopy(
             FlexibleLayout(


Test patch:
diff --git a/tests/test_torchinductor.py b/tests/test_torchinductor.py
--- a/tests/test_torchinductor.py
+++ b/tests/test_torchinductor.py
@@ -1387,6 +1387,14 @@ def test_layer_norm(self):
         if self.device != "cpu":
             self.assertEqual(torchinductor.metrics.generated_kernel_count, 1)
 
+    def test_move_arange(self):
+        def fn(x):
+            return torch.arange(len(x), device="cpu").to(x.device) + x
+
+        self.common(fn, (torch.randn([32]),), check_lowp=False)
+        # if we have a copy there will be more than 1 kernel
+        self.assertEqual(torchinductor.metrics.generated_kernel_count, 1)
+
     def test_leaky_relu(self):
         def fn(x):
             return aten.leaky_relu(x, 0.2) + 2, aten.leaky_relu(x + 1)