已合并
test(utils): adapt torch.utils.checkpoint.checkpoint, torch.utils.checkpoint.checkpoint_sequential and torch.overrides.is_tensor_like on NPU #39485
yulin520创建于 6月27日
test(utils): adapt torch.utils.checkpoint.checkpoint, torch.utils.checkpoint.checkpoint_sequential and torch.overrides.is_tensor_like on NPU #39485
已合并
共 2 个文件变更+104-24
| @@ -1,28 +1,82 @@ | |||
| 1 | diff --git a/test/dynamo/test_misc.py b/test/dynamo/test_misc.py | 1 | diff --git a/test/dynamo/test_misc.py b/test/dynamo/test_misc.py |
| 2 | -index 98f5972..a28b01a 100644 | 2 | +index 98f5972..be78a08 100644 |
| 3 | --- a/test/dynamo/test_misc.py | 3 | --- a/test/dynamo/test_misc.py |
| 4 | +++ b/test/dynamo/test_misc.py | 4 | +++ b/test/dynamo/test_misc.py |
| 5 | -@@ -11536 +11536 @@ def ___make_guard_fn(): | 5 | +@@ -35,6 +35,8 @@ from unittest.mock import patch |
| 6 | + import numpy as np | ||
| 7 | + | ||
| 8 | + import torch | ||
| 9 | ++import torch_npu | ||
| 10 | ++from torch_npu.contrib import transfer_to_npu | ||
| 11 | + import torch._dynamo.testing | ||
| 12 | + import torch._inductor.config | ||
| 13 | + import torch._inductor.test_case | ||
| 14 | +@@ -5131,9 +5133,9 @@ not ___dict_contains('cccccccc', G['sys'].modules)""", | ||
| 15 | + def f(x): | ||
| 16 | + if torch.overrides.is_tensor_like(x): | ||
| 17 | + return (x * 2,) | ||
| 18 | +- return (torch.ones(10) + x,) | ||
| 19 | ++ return (torch.ones(10).npu() + x,) | ||
| 20 | + | ||
| 21 | +- x = torch.randn(10) | ||
| 22 | ++ x = torch.randn(10).npu() | ||
| 23 | + ref0 = f(x) | ||
| 24 | + ref1 = f(4) | ||
| 25 | + opt_f = torch.compile(f, backend=cnts, fullgraph=True) | ||
| 26 | +@@ -11533,7 +11535,7 @@ def ___make_guard_fn(): | ||
| 27 | + | ||
| 28 | + # Call create_symbolic_sizes_strides_storage_offset on both of them. | ||
| 29 | + r = main.create_symbolic_sizes_strides_storage_offset( | ||
| 6 | - torch.randn(3, 2), ConstantSource("x") | 30 | - torch.randn(3, 2), ConstantSource("x") |
| 7 | + torch.randn(3, 2).npu(), ConstantSource("x") | 31 | + torch.randn(3, 2).npu(), ConstantSource("x") |
| 8 | -@@ -11590 +11590 @@ ShapeEnv not equal: field values don't match: | 32 | + ) |
| 33 | + | ||
| 34 | + # Create a guard: size[0] == 3 (call evaluate_expr) | ||
| 35 | +@@ -11587,7 +11589,7 @@ ShapeEnv not equal: field values don't match: | ||
| 36 | + def test_shape_env_equal_create_symbolic_sizes_strides_storage_offset(self): | ||
| 37 | + main, other = ShapeEnv(), ShapeEnv() | ||
| 38 | + main.create_symbolic_sizes_strides_storage_offset( | ||
| 9 | - torch.randn(3, 2), ConstantSource("x") | 39 | - torch.randn(3, 2), ConstantSource("x") |
| 10 | + torch.randn(3, 2).npu(), ConstantSource("x") | 40 | + torch.randn(3, 2).npu(), ConstantSource("x") |
| 11 | -@@ -11660 +11660 @@ ShapeEnv not equal: field values don't match: | 41 | + ) |
| 42 | + self.assertExpectedRaisesInline( | ||
| 43 | + NotEqualError, | ||
| 44 | +@@ -11657,10 +11659,10 @@ ShapeEnv not equal: field values don't match: | ||
| 45 | + | ||
| 46 | + # Call create_symbolic_sizes_strides_storage_offset on both of them. | ||
| 47 | + r = main.create_symbolic_sizes_strides_storage_offset( | ||
| 12 | - torch.randn(3, 2), ConstantSource("x") | 48 | - torch.randn(3, 2), ConstantSource("x") |
| 13 | + torch.randn(3, 2).npu(), ConstantSource("x") | 49 | + torch.randn(3, 2).npu(), ConstantSource("x") |
| 14 | -@@ -11663 +11663 @@ ShapeEnv not equal: field values don't match: | 50 | + ) |
| 51 | + other.create_symbolic_sizes_strides_storage_offset( | ||
| 15 | - torch.randn(3, 2), ConstantSource("x") | 52 | - torch.randn(3, 2), ConstantSource("x") |
| 16 | + torch.randn(3, 2).npu(), ConstantSource("x") | 53 | + torch.randn(3, 2).npu(), ConstantSource("x") |
| 17 | -@@ -11700 +11700 @@ ShapeEnv not equal: field values don't match: | 54 | + ) |
| 55 | + | ||
| 56 | + # Create a guard: size[0] % 3 == 0 (only in the main ShapeEnv) | ||
| 57 | +@@ -11697,10 +11699,10 @@ ShapeEnv not equal: field values don't match: | ||
| 58 | + | ||
| 59 | + # Call create_symbolic_sizes_strides_storage_offset on both of them. | ||
| 60 | + r = main.create_symbolic_sizes_strides_storage_offset( | ||
| 18 | - torch.randn(3, 2), ConstantSource("x") | 61 | - torch.randn(3, 2), ConstantSource("x") |
| 19 | + torch.randn(3, 2).npu(), ConstantSource("x") | 62 | + torch.randn(3, 2).npu(), ConstantSource("x") |
| 20 | -@@ -11703 +11703 @@ ShapeEnv not equal: field values don't match: | 63 | + ) |
| 64 | + other.create_symbolic_sizes_strides_storage_offset( | ||
| 21 | - torch.randn(3, 2), ConstantSource("x") | 65 | - torch.randn(3, 2), ConstantSource("x") |
| 22 | + torch.randn(3, 2).npu(), ConstantSource("x") | 66 | + torch.randn(3, 2).npu(), ConstantSource("x") |
| 23 | -@@ -11743 +11743 @@ ShapeEnv not equal: field values don't match: | 67 | + ) |
| 68 | + | ||
| 69 | + # Create a guard: size[0] == 3 (only in the main ShapeEnv) | ||
| 70 | +@@ -11740,10 +11742,10 @@ ShapeEnv not equal: field values don't match: | ||
| 71 | + | ||
| 72 | + # Call create_symbolic_sizes_strides_storage_offset on both of them. | ||
| 73 | + r = main.create_symbolic_sizes_strides_storage_offset( | ||
| 24 | - torch.randn(3, 2), ConstantSource("x") | 74 | - torch.randn(3, 2), ConstantSource("x") |
| 25 | + torch.randn(3, 2).npu(), ConstantSource("x") | 75 | + torch.randn(3, 2).npu(), ConstantSource("x") |
| 26 | -@@ -11746 +11746 @@ ShapeEnv not equal: field values don't match: | 76 | + ) |
| 77 | + other.create_symbolic_sizes_strides_storage_offset( | ||
| 27 | - torch.randn(3, 2), ConstantSource("x") | 78 | - torch.randn(3, 2), ConstantSource("x") |
| 28 | + torch.randn(3, 2).npu(), ConstantSource("x") | 79 | + torch.randn(3, 2).npu(), ConstantSource("x") |
| 80 | + ) | ||
| 81 | + | ||
| 82 | + # Create a guard: size[0] >= 3 (only in the main ShapeEnv) | ||
| @@ -1,16 +1,42 @@ | |||
| 1 | diff --git a/test/test_autograd.py b/test/test_autograd.py | 1 | diff --git a/test/test_autograd.py b/test/test_autograd.py |
| 2 | -index eb2ad5c..9d1e3fb 100644 | 2 | +index eb2ad5c..88d6071 100644 |
| 3 | --- a/test/test_autograd.py | 3 | --- a/test/test_autograd.py |
| 4 | +++ b/test/test_autograd.py | 4 | +++ b/test/test_autograd.py |
| 5 | -@@ -28,6 +28,7 @@ from operator import mul | 5 | +@@ -28,6 +28,8 @@ from operator import mul |
| 6 | from typing import TYPE_CHECKING | 6 | from typing import TYPE_CHECKING |
| 7 | 7 | ||
| 8 | import torch | 8 | import torch |
| 9 | +import torch_npu | 9 | +import torch_npu |
| 10 | ++from torch_npu.contrib import transfer_to_npu | ||
| 10 | import torch.autograd._functions | 11 | import torch.autograd._functions |
| 11 | import torch.autograd.forward_ad as fwAD | 12 | import torch.autograd.forward_ad as fwAD |
| 12 | from torch import inf, nan, nn | 13 | from torch import inf, nan, nn |
| 13 | -@@ -7588,7 +7589,7 @@ for shape in [(1,), ()]: | 14 | +@@ -7490,11 +7492,11 @@ for shape in [(1,), ()]: |
| 15 | + checkpoint(lambda x: x, a) | ||
| 16 | + | ||
| 17 | + def test_checkpoint_sequential_warns_if_use_reentrant_not_passed_explcitly(self): | ||
| 18 | +- a = torch.randn(3, requires_grad=True) | ||
| 19 | ++ a = torch.randn(3, device="npu", requires_grad=True) | ||
| 20 | + modules_list = [ | ||
| 21 | +- torch.nn.Linear(3, 3), | ||
| 22 | +- torch.nn.Linear(3, 3), | ||
| 23 | +- torch.nn.Linear(3, 3), | ||
| 24 | ++ torch.nn.Linear(3, 3).npu(), | ||
| 25 | ++ torch.nn.Linear(3, 3).npu(), | ||
| 26 | ++ torch.nn.Linear(3, 3).npu(), | ||
| 27 | + ] | ||
| 28 | + | ||
| 29 | + # Passing explicitly should not warn | ||
| 30 | +@@ -7574,7 +7576,7 @@ for shape in [(1,), ()]: | ||
| 31 | + return x.sin().exp() | ||
| 32 | + | ||
| 33 | + def save_2_tensors_alt(x): | ||
| 34 | +- return x.sin() * torch.tensor([1.0, 2.0]) | ||
| 35 | ++ return x.sin() * torch.tensor([1., 2.], device=x.device) | ||
| 36 | + | ||
| 37 | + def get_non_det_fn(orig_fn, recompute_fn): | ||
| 38 | + counter = [0] | ||
| 39 | +@@ -7588,7 +7590,7 @@ for shape in [(1,), ()]: | ||
| 14 | 40 | ||
| 15 | return fn | 41 | return fn |
| 16 | 42 | ||
| @@ -19,7 +45,7 @@ index eb2ad5c..9d1e3fb 100644 | |||
| 19 | 45 | ||
| 20 | # Save fewer tensors during recompute | 46 | # Save fewer tensors during recompute |
| 21 | fn = get_non_det_fn(orig_fn=save_3_tensors, recompute_fn=save_2_tensors) | 47 | fn = get_non_det_fn(orig_fn=save_3_tensors, recompute_fn=save_2_tensors) |
| 22 | -@@ -7737,7 +7738,7 @@ for shape in [(1,), ()]: | 48 | +@@ -7737,7 +7739,7 @@ for shape in [(1,), ()]: |
| 23 | self.assertEqual(checkpoint_param.grad, param.grad) | 49 | self.assertEqual(checkpoint_param.grad, param.grad) |
| 24 | 50 | ||
| 25 | def test_checkpoint_valid_reset_on_error(self): | 51 | def test_checkpoint_valid_reset_on_error(self): |
| @@ -28,7 +54,7 @@ index eb2ad5c..9d1e3fb 100644 | |||
| 28 | 54 | ||
| 29 | with self.assertRaisesRegex( | 55 | with self.assertRaisesRegex( |
| 30 | Exception, "torch.utils.checkpoint is incompatible" | 56 | Exception, "torch.utils.checkpoint is incompatible" |
| 31 | -@@ -14946,8 +14947,8 @@ class TestSelectiveActivationCheckpoint(TestCase): | 57 | +@@ -14946,8 +14948,8 @@ class TestSelectiveActivationCheckpoint(TestCase): |
| 32 | def fn(x, y): | 58 | def fn(x, y): |
| 33 | return x.view(-1), y.sin().cos() | 59 | return x.view(-1), y.sin().cos() |
| 34 | 60 | ||
| @@ -39,7 +65,7 @@ index eb2ad5c..9d1e3fb 100644 | |||
| 39 | 65 | ||
| 40 | context_fn = functools.partial( | 66 | context_fn = functools.partial( |
| 41 | create_selective_checkpoint_contexts, | 67 | create_selective_checkpoint_contexts, |
| 42 | -@@ -14963,7 +14964,7 @@ class TestSelectiveActivationCheckpoint(TestCase): | 68 | +@@ -14963,7 +14965,7 @@ class TestSelectiveActivationCheckpoint(TestCase): |
| 43 | # | 69 | # |
| 44 | # 1. size/strides are dispatched upon | 70 | # 1. size/strides are dispatched upon |
| 45 | # 2. our policy saves sizes | 71 | # 2. our policy saves sizes |
| @@ -48,7 +74,7 @@ index eb2ad5c..9d1e3fb 100644 | |||
| 48 | 74 | ||
| 49 | class CustomSizeDynamicShapesTensor(torch.Tensor): | 75 | class CustomSizeDynamicShapesTensor(torch.Tensor): |
| 50 | @staticmethod | 76 | @staticmethod |
| 51 | -@@ -15023,7 +15024,7 @@ class TestSelectiveActivationCheckpoint(TestCase): | 77 | +@@ -15023,7 +15025,7 @@ class TestSelectiveActivationCheckpoint(TestCase): |
| 52 | self.assertEqual(view_out.grad_fn._saved_self_sym_sizes, [12]) | 78 | self.assertEqual(view_out.grad_fn._saved_self_sym_sizes, [12]) |
| 53 | return view_out.exp() | 79 | return view_out.exp() |
| 54 | 80 | ||
| @@ -57,7 +83,7 @@ index eb2ad5c..9d1e3fb 100644 | |||
| 57 | x_wrapper = CustomSizeDynamicShapesTensor(x) | 83 | x_wrapper = CustomSizeDynamicShapesTensor(x) |
| 58 | context_fn = functools.partial(create_selective_checkpoint_contexts, policy_fn) | 84 | context_fn = functools.partial(create_selective_checkpoint_contexts, policy_fn) |
| 59 | out = checkpoint(fn, x_wrapper, use_reentrant=False, context_fn=context_fn) | 85 | out = checkpoint(fn, x_wrapper, use_reentrant=False, context_fn=context_fn) |
| 60 | -@@ -15072,7 +15073,7 @@ class TestSelectiveActivationCheckpoint(TestCase): | 86 | +@@ -15072,7 +15074,7 @@ class TestSelectiveActivationCheckpoint(TestCase): |
| 61 | def fn(x): | 87 | def fn(x): |
| 62 | return x.sin().sin().sin() | 88 | return x.sin().sin().sin() |
| 63 | 89 | ||
| @@ -66,7 +92,7 @@ index eb2ad5c..9d1e3fb 100644 | |||
| 66 | context_fn = functools.partial( | 92 | context_fn = functools.partial( |
| 67 | create_selective_checkpoint_contexts, | 93 | create_selective_checkpoint_contexts, |
| 68 | Policy(), | 94 | Policy(), |
| 69 | -@@ -15124,7 +15125,7 @@ class TestSelectiveActivationCheckpoint(TestCase): | 95 | +@@ -15124,7 +15126,7 @@ class TestSelectiveActivationCheckpoint(TestCase): |
| 70 | with disable_gc(): | 96 | with disable_gc(): |
| 71 | # Case 1: If graph goes away without backward, make sure there's no reference cycle | 97 | # Case 1: If graph goes away without backward, make sure there's no reference cycle |
| 72 | # keeping storage alive. | 98 | # keeping storage alive. |
| @@ -75,7 +101,7 @@ index eb2ad5c..9d1e3fb 100644 | |||
| 75 | context_fn = functools.partial( | 101 | context_fn = functools.partial( |
| 76 | create_selective_checkpoint_contexts, policy_fn | 102 | create_selective_checkpoint_contexts, policy_fn |
| 77 | ) | 103 | ) |
| 78 | -@@ -15134,7 +15135,7 @@ class TestSelectiveActivationCheckpoint(TestCase): | 104 | +@@ -15134,7 +15136,7 @@ class TestSelectiveActivationCheckpoint(TestCase): |
| 79 | self.assertIsNone(ref()) | 105 | self.assertIsNone(ref()) |
| 80 | 106 | ||
| 81 | # Case 2: After backward, even if retain_graph=True, the storage should go away | 107 | # Case 2: After backward, even if retain_graph=True, the storage should go away |
| @@ -84,7 +110,7 @@ index eb2ad5c..9d1e3fb 100644 | |||
| 84 | context_fn = functools.partial( | 110 | context_fn = functools.partial( |
| 85 | create_selective_checkpoint_contexts, policy_fn | 111 | create_selective_checkpoint_contexts, policy_fn |
| 86 | ) | 112 | ) |
| 87 | -@@ -15156,7 +15157,7 @@ class TestSelectiveActivationCheckpoint(TestCase): | 113 | +@@ -15156,7 +15158,7 @@ class TestSelectiveActivationCheckpoint(TestCase): |
| 88 | def fn(x): | 114 | def fn(x): |
| 89 | return x.sin().mul_(2).cos().exp() | 115 | return x.sin().mul_(2).cos().exp() |
| 90 | 116 | ||
| @@ -93,7 +119,7 @@ index eb2ad5c..9d1e3fb 100644 | |||
| 93 | context_fn = functools.partial(create_selective_checkpoint_contexts, policy_fn) | 119 | context_fn = functools.partial(create_selective_checkpoint_contexts, policy_fn) |
| 94 | out = checkpoint(fn, x, use_reentrant=False, context_fn=context_fn) | 120 | out = checkpoint(fn, x, use_reentrant=False, context_fn=context_fn) |
| 95 | 121 | ||
| 96 | -@@ -15164,7 +15165,7 @@ class TestSelectiveActivationCheckpoint(TestCase): | 122 | +@@ -15164,7 +15166,7 @@ class TestSelectiveActivationCheckpoint(TestCase): |
| 97 | with self.assertRaisesRegex(RuntimeError, "has been mutated"): | 123 | with self.assertRaisesRegex(RuntimeError, "has been mutated"): |
| 98 | out.sum().backward() | 124 | out.sum().backward() |
| 99 | 125 | ||
| @@ -102,7 +128,7 @@ index eb2ad5c..9d1e3fb 100644 | |||
| 102 | context_fn = functools.partial( | 128 | context_fn = functools.partial( |
| 103 | create_selective_checkpoint_contexts, | 129 | create_selective_checkpoint_contexts, |
| 104 | policy_fn, | 130 | policy_fn, |
| 105 | -@@ -15192,7 +15193,7 @@ class TestSelectiveActivationCheckpoint(TestCase): | 131 | +@@ -15192,7 +15194,7 @@ class TestSelectiveActivationCheckpoint(TestCase): |
| 106 | a, b = torch.var_mean(x) | 132 | a, b = torch.var_mean(x) |
| 107 | return a * b | 133 | return a * b |
| 108 | 134 | ||
| @@ -111,7 +137,7 @@ index eb2ad5c..9d1e3fb 100644 | |||
| 111 | context_fn = functools.partial(create_selective_checkpoint_contexts, policy_fn) | 137 | context_fn = functools.partial(create_selective_checkpoint_contexts, policy_fn) |
| 112 | out = checkpoint(fn, x, use_reentrant=False, context_fn=context_fn) | 138 | out = checkpoint(fn, x, use_reentrant=False, context_fn=context_fn) |
| 113 | x_grad = torch.autograd.grad(out.sum(), (x,)) | 139 | x_grad = torch.autograd.grad(out.sum(), (x,)) |
| 114 | -@@ -15223,14 +15224,14 @@ class TestSelectiveActivationCheckpoint(TestCase): | 140 | +@@ -15223,14 +15225,14 @@ class TestSelectiveActivationCheckpoint(TestCase): |
| 115 | "mylib::sin_with_extra", backward, setup_context=setup_context | 141 | "mylib::sin_with_extra", backward, setup_context=setup_context |
| 116 | ) | 142 | ) |
| 117 | 143 | ||
| @@ -128,7 +154,7 @@ index eb2ad5c..9d1e3fb 100644 | |||
| 128 | context_fn = functools.partial( | 154 | context_fn = functools.partial( |
| 129 | create_selective_checkpoint_contexts, ops_list | 155 | create_selective_checkpoint_contexts, ops_list |
| 130 | ) | 156 | ) |
| 131 | -@@ -15254,7 +15255,7 @@ class TestSelectiveActivationCheckpoint(TestCase): | 157 | +@@ -15254,7 +15256,7 @@ class TestSelectiveActivationCheckpoint(TestCase): |
| 132 | def fn(x): | 158 | def fn(x): |
| 133 | return x.sin().cos().exp() | 159 | return x.sin().cos().exp() |
| 134 | 160 | ||