Device-generalize more dynamo tests to run on accelerator backends

test/dynamo/test_activation_checkpointing.py

@@ -56,7 +56,9 @@
 )
 
 
-device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"
+device_type = (
+    acc.type if (acc := torch.accelerator.current_accelerator(True)) else "cpu"
+)
 
 
 if HAS_GPU_AND_TRITON:
@@ -2101,7 +2103,7 @@ def gn(x):
         def fn(x):
             return torch.utils.checkpoint.checkpoint(gn, x, use_reentrant=True)
 
-        x = torch.randn(4, 4, requires_grad=True)
+        x = torch.randn(4, 4, device=device_type, requires_grad=True)
         fn(x).sum().backward()
         # The mutation is reapplied in the backward as well
         self.assertEqual(counter, 2)
@@ -2270,7 +2272,9 @@ def checkpointed_forward(inp):
                     preserve_rng_state=True,
                 )
 
-            input_eager = InputNode(x=torch.randn(2, 4, requires_grad=True))
+            input_eager = InputNode(
+                x=torch.randn(2, 4, device=device_type, requires_grad=True)
+            )
             torch.manual_seed(0)
             output_eager = checkpointed_forward(input_eager)
             output_eager.y.sum().backward()
@@ -2584,8 +2588,8 @@ def fn(x, y):
                 return (torch.mm(a, y) + 1).relu()
 
         cfn = torch.compile(fn, backend="aot_eager", fullgraph=True)
-        x = torch.randn(8, 8, requires_grad=True)
-        y = torch.randn(8, 8, requires_grad=True)
+        x = torch.randn(8, 8, device=device_type, requires_grad=True)
+        y = torch.randn(8, 8, device=device_type, requires_grad=True)
         with self.assertRaisesRegex(RuntimeError, "conflicting budgets"):
             cfn(x, y).sum().backward()
 
@@ -2600,8 +2604,8 @@ def fn(x, y):
             return (a * 2).relu()
 
         cfn = torch.compile(fn, backend="aot_eager", fullgraph=True)
-        x = torch.randn(8, 8, requires_grad=True)
-        y = torch.randn(8, 8, requires_grad=True)
+        x = torch.randn(8, 8, device=device_type, requires_grad=True)
+        y = torch.randn(8, 8, device=device_type, requires_grad=True)
         with self.assertRaisesRegex(RuntimeError, "must cover the entire forward"):
             cfn(x, y).sum().backward()
 

test/dynamo/test_autograd_function.py

@@ -462,13 +462,13 @@ def backward(ctx, grad):
                 return grad * 3
 
         def fn(w):
-            y = torch.ones(2, 2) @ w
+            y = torch.ones(2, 2, device=device_type) @ w
             TimesThreeInplace.apply(y)
             return y.sum()
 
         def grad_for(compiled):
             torch._dynamo.reset()
-            w = torch.eye(2, requires_grad=True)
+            w = torch.eye(2, device=device_type, requires_grad=True)
             f = torch.compile(fn, backend="eager", fullgraph=True) if compiled else fn
             loss = f(w)
             loss.backward()
@@ -1956,7 +1956,9 @@ def fn(x):
             return Foo.apply(x, x).sum()
 
         def check_fallback():
-            x = torch.tensor([0.7, -1.3, 2.1, 0.05, -0.5], requires_grad=True)
+            x = torch.tensor(
+                [0.7, -1.3, 2.1, 0.05, -0.5], device=device_type, requires_grad=True
+            )
 
             x_ref = x.detach().clone().requires_grad_(True)
             ref = fn(x_ref)
@@ -2647,7 +2649,7 @@ def backward(ctx, grad_output):
         def fn(x):
             return MySin.apply(x, factor=6)
 
-        x = torch.tensor([0.812], requires_grad=True)
+        x = torch.tensor([0.812], device=device_type, requires_grad=True)
         ref = fn(x)
         opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
         res = opt_fn(x)
@@ -2678,7 +2680,7 @@ def backward(ctx, gO):
         def fn(x):
             return MyScale.apply(x, factor=3)
 
-        x = torch.tensor([2.0], requires_grad=True)
+        x = torch.tensor([2.0], device=device_type, requires_grad=True)
         ref = fn(x)
         opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
         res = opt_fn(x)

test/dynamo/test_functions.py

@@ -655,7 +655,12 @@ def fn():
     def test_itertools_compress_tensors(self):
         def fn():
             return itertools.compress(
-                [torch.tensor([0]), torch.tensor([1]), torch.tensor([2])], [1, 0, 1]
+                [
+                    torch.tensor([0], device=device_type),
+                    torch.tensor([1], device=device_type),
+                    torch.tensor([2], device=device_type),
+                ],
+                [1, 0, 1],
             )
 
         opt_fn = torch.compile(fn, backend="eager", fullgraph=True)

test/dynamo/test_fwd_loss_bwd.py

@@ -176,8 +176,8 @@ def fn(x):
 
     @skipIfCrossRef
     def test_backward_dict_inputs(self):
-        mod = torch.nn.Linear(4, 4)
-        x = torch.randn(2, 4)
+        mod = torch.nn.Linear(4, 4).to(device_type)
+        x = torch.randn(2, 4, device=device_type)
 
         def fn(x):
             res = mod(x)

test/dynamo/test_hooks.py

@@ -599,20 +599,20 @@ def forward(self, x):
                     z = y.mul(3)
                     return (z,)
 
-            mod = MyMod()
-            x0 = torch.ones(4, requires_grad=True)
+            mod = MyMod().to(device_type)
+            x0 = torch.ones(4, device=device_type, requires_grad=True)
             eager_out = mod(x0)
-            eager_out[0].backward(torch.ones(4))
+            eager_out[0].backward(torch.ones(4, device=device_type))
 
-            x1 = torch.ones(4, requires_grad=True)
+            x1 = torch.ones(4, device=device_type, requires_grad=True)
             mod_compiled = aot_module_simplified(mod, (x1,), nop)
             aot_out = mod_compiled(x1)
-            aot_out[0].backward(torch.ones(4))
+            aot_out[0].backward(torch.ones(4, device=device_type))
 
-            x2 = torch.ones(4, requires_grad=True)
+            x2 = torch.ones(4, device=device_type, requires_grad=True)
             dynamo_out = torch.compile(mod, backend=backend, fullgraph=True)(x2)
             with compiled_autograd._enable(compiler_fn):
-                dynamo_out[0].backward(torch.ones(4))
+                dynamo_out[0].backward(torch.ones(4, device=device_type))
 
             self.assertEqual(dynamo_out, aot_out)
             self.assertEqual(dynamo_out, eager_out)
@@ -967,7 +967,7 @@ def reg_and_mul(x, y):
 
         def test_fn(fn):
             fn(x, y)
-            b = torch.tensor([2.0, 2.0, 2.0], requires_grad=True)
+            b = torch.tensor([2.0, 2.0, 2.0], device=device_type, requires_grad=True)
             x.backward(b)
             if cnts:
                 self.assertEqual(cnts.frame_count, 1)

test/dynamo/test_misc.py

@@ -1982,13 +1982,15 @@ def __init__(self):
             def forward(self, x):
                 batch_size = x.size(0)
                 h = torch.zeros(
-                    self.num_layers, batch_size, self.hidden_size
+                    self.num_layers, batch_size, self.hidden_size, device=x.device
                 ).share_memory_()
-                c = torch.zeros(self.num_layers, batch_size, self.hidden_size)
+                c = torch.zeros(
+                    self.num_layers, batch_size, self.hidden_size, device=x.device
+                )
                 return x + h.sum() + c.sum()
 
         model = Model()
-        x = torch.randn(4, 10)
+        x = torch.randn(4, 10, device=device_type)
         expected = model(x)
         compiled_model = torch.compile(model, fullgraph=False, backend="eager")
         actual = compiled_model(x)
@@ -11984,8 +11986,8 @@ def test_compile_with_userland_fake_tensor_mode(self):
         from torch._subclasses.fake_tensor import FakeTensorMode
 
         with FakeTensorMode():
-            model = torch.nn.Linear(4, 4)
-            inp = torch.rand(4, 4)
+            model = torch.nn.Linear(4, 4).to(device_type)
+            inp = torch.rand(4, 4, device=device_type)
             loss = torch.compile(model, backend="aot_eager")(inp).sum()
             loss.backward()
 

test/dynamo/test_repros.py

@@ -1326,7 +1326,7 @@ def fn(x):
             log_det += torch.zeros(x.size(0), device="meta")
             return log_det
 
-        x = torch.randn(2, 4)
+        x = torch.randn(2, 4, device=device_type)
         eager_out = fn(x)
         compiled_fn = torch.compile(fn, backend="eager", fullgraph=True)
         compiled_out = compiled_fn(x)
@@ -1354,7 +1354,7 @@ def f(x):
     # https://github.com/pytorch/pytorch/issues/90552
     def test_intermediate_leaf_requires_grad(self):
         def f(x):
-            leaf = torch.ones(2, requires_grad=True)
+            leaf = torch.ones(2, device=x.device, requires_grad=True)
             return leaf, leaf * 2
 
         f_compiled = torch.compile(f, backend="aot_eager")
@@ -7886,8 +7886,8 @@ def forward(self, input_seq, input_lengths):
                 )
                 return self.proj(context)
 
-        model = Model().eval()
-        input_seq = torch.arange(10, dtype=torch.long).view(5, 2)
+        model = Model().eval().to(device_type)
+        input_seq = torch.arange(10, dtype=torch.long, device=device_type).view(5, 2)
         input_lengths = torch.tensor([3, 5], dtype=torch.int64)
 
         expected = model(input_seq, input_lengths)

test/dynamo/test_wrap_inductor_compiled_regions.py

@@ -23,6 +23,11 @@
 )
 
 
+device_type = (
+    acc.type if (acc := torch.accelerator.current_accelerator(True)) else "cpu"
+)
+
+
 def count_ops(
     gm, args, freq=None, freq_ge=None, op=None, freqs=None, freqs_ge=None, ops=None
 ):
@@ -1413,7 +1418,9 @@ def test_sac_cached_value_fifo_mismatch(self):
         dist.init_process_group(backend="fake", store=fake_store, rank=0, world_size=1)
 
         try:
-            mesh = init_device_mesh("cpu", mesh_shape=(1,), mesh_dim_names=("dp",))
+            mesh = init_device_mesh(
+                device_type, mesh_shape=(1,), mesh_dim_names=("dp",)
+            )
 
             @torch.compile(dynamic=False, fullgraph=True)  # noqa: UNSPECIFIED_BACKEND
             def compute_int_stuff(n):
@@ -1449,7 +1456,9 @@ def policy_fn(ctx, op, *args, **kwargs):
             )
 
             x = DTensor.from_local(
-                torch.randn(4, 4, dtype=torch.float32, requires_grad=True),
+                torch.randn(
+                    4, 4, dtype=torch.float32, device=device_type, requires_grad=True
+                ),
                 mesh,
                 (Replicate(),),
             )