From 4253ed1d24b94fd4ab17413e41a76ccb174a591b Mon Sep 17 00:00:00 2001 From: aws-alexk Date: Tue, 4 Aug 2026 23:04:26 +0000 Subject: [PATCH] Device-generalize more dynamo tests to run on accelerator backends Follow-up to the earlier device-generalization pass. These test/dynamo tests still hardcode CPU input tensors (no device= argument). torch.compile is device-preserving, but accelerator backends that execute the compiled graph on-device return device tensors, so autograd rejects the device mismatch on backward, e.g.: RuntimeError: Function CompiledFunctionBackward returned an invalid gradient at index 0 - expected device cpu but got :0 or the compiled region raises a plain device mismatch when a CPU input meets an on-device intermediate. Create inputs on the current accelerator via the module-level device_type (torch.accelerator.current_accelerator, falling back to "cpu"), adding device=device_type to input tensor factories and .to(device_type) to nn.Module instances in the affected tests. test_activation_checkpointing.py and test_wrap_inductor_compiled_regions.py gain the module-level device_type definition (the other files already have it); the latter also generalizes its DTensor device mesh. Because device_type resolves to "cpu" when no accelerator is present, these edits are a no-op on CPU and CUDA CI and only take effect on accelerator backends. Test Plan: On CPU (no accelerator; device_type == "cpu", edits are a no-op): ``` python test/dynamo/test_autograd_function.py AutogradFunctionTests.test_apply_kwargs_old_style python test/dynamo/test_hooks.py HooksTests.test_input_hooks_same python test/dynamo/test_repros.py ReproTests.test_intermediate_leaf_requires_grad python test/dynamo/test_fwd_loss_bwd.py TestForwardLossBackward.test_backward_dict_inputs ``` On an accelerator backend, the same nodeids (and their _nested_graph_breaks variants) that previously failed with a device mismatch now pass. Authored with an AI assistant (Claude). --- test/dynamo/test_activation_checkpointing.py | 18 +++++++++++------- test/dynamo/test_autograd_function.py | 12 +++++++----- test/dynamo/test_functions.py | 7 ++++++- test/dynamo/test_fwd_loss_bwd.py | 4 ++-- test/dynamo/test_hooks.py | 16 ++++++++-------- test/dynamo/test_misc.py | 12 +++++++----- test/dynamo/test_repros.py | 8 ++++---- .../test_wrap_inductor_compiled_regions.py | 13 +++++++++++-- 8 files changed, 56 insertions(+), 34 deletions(-) diff --git a/test/dynamo/test_activation_checkpointing.py b/test/dynamo/test_activation_checkpointing.py index fd32bc3131102..7b515cab2109c 100644 --- a/test/dynamo/test_activation_checkpointing.py +++ b/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() diff --git a/test/dynamo/test_autograd_function.py b/test/dynamo/test_autograd_function.py index c631ee07b3bb8..1feadef2f470c 100644 --- a/test/dynamo/test_autograd_function.py +++ b/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) diff --git a/test/dynamo/test_functions.py b/test/dynamo/test_functions.py index 1a34158363ec9..63e924206fee7 100644 --- a/test/dynamo/test_functions.py +++ b/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) diff --git a/test/dynamo/test_fwd_loss_bwd.py b/test/dynamo/test_fwd_loss_bwd.py index c5e5310e16ee8..ef722a0967526 100644 --- a/test/dynamo/test_fwd_loss_bwd.py +++ b/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) diff --git a/test/dynamo/test_hooks.py b/test/dynamo/test_hooks.py index 9cdc537e55411..39323dfe9074a 100644 --- a/test/dynamo/test_hooks.py +++ b/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) diff --git a/test/dynamo/test_misc.py b/test/dynamo/test_misc.py index beb2d27150dab..e6aaa1ea616a5 100644 --- a/test/dynamo/test_misc.py +++ b/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() diff --git a/test/dynamo/test_repros.py b/test/dynamo/test_repros.py index f7779c1367d97..60898393e1c31 100644 --- a/test/dynamo/test_repros.py +++ b/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) diff --git a/test/dynamo/test_wrap_inductor_compiled_regions.py b/test/dynamo/test_wrap_inductor_compiled_regions.py index c169433685a8e..72b7fe7fba3ae 100644 --- a/test/dynamo/test_wrap_inductor_compiled_regions.py +++ b/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(),), )