已合并
[UT] adapt 4 official testcase files for dynamo module（test_optimizers.py test_hooks.py test_debug_utils.py test_bytecode_hook.py） #8820
AtomGit-Bot创建于 2024年1月8日
[UT] adapt 4 official testcase files for dynamo module（test_optimizers.py test_hooks.py test_debug_utils.py test_bytecode_hook.py） #8820
已合并
共 4 个文件变更+898-0
| @@ -0,0 +1,29 @@ | |||
| 1 | +# Owner(s): ["module: dynamo"] | ||
| 2 | + | ||
| 3 | +import torch | ||
| 4 | +import torch_npu | ||
| 5 | +import torch._dynamo.test_case | ||
M | |||
| 6 | + | ||
| 7 | + | ||
| 8 | +class BytecodeHookTests(torch._dynamo.test_case.TestCase): | ||
| 9 | + def test_bytecode_hook(self): | ||
| 10 | + def fn(a, b): | ||
| 11 | + return a - b * 10 | ||
| 12 | + | ||
| 13 | + def hook(code, out_code): | ||
| 14 | + return code | ||
M 这个print是为什么存在？   | |||
| 15 | + | ||
| 16 | + torch._dynamo.reset() | ||
| 17 | + handle = torch._dynamo.convert_frame.register_bytecode_hook(hook) | ||
| 18 | + try: | ||
| 19 | + opt_fn = torch.compile(fn) | ||
| 20 | + for i in range(2, 12): | ||
| 21 | + opt_fn(torch.randn(i), torch.randn(i)) | ||
| 22 | + finally: | ||
| 23 | + handle.remove() | ||
| 24 | + | ||
| 25 | + | ||
| 26 | +if __name__ == "__main__": | ||
| 27 | + from torch._dynamo.test_case import run_tests | ||
| 28 | + | ||
| 29 | + run_tests() | ||
| @@ -0,0 +1,57 @@ | |||
| 1 | +# Owner(s): ["module: dynamo"] | ||
| 2 | + | ||
| 3 | +import torch | ||
| 4 | +import torch_npu | ||
| 5 | +from functorch import make_fx | ||
| 6 | +from torch._dynamo import debug_utils | ||
| 7 | +from torch._dynamo.test_case import TestCase | ||
| 8 | + | ||
| 9 | + | ||
| 10 | +class TestDebugUtils(TestCase): | ||
| 11 | + def test_cast_model_to_fp64_dtype_args(self): | ||
| 12 | + # Test that dtype arguments are converted to fp64 | ||
| 13 | + | ||
| 14 | + def fn(x): | ||
| 15 | + return ( | ||
| 16 | + torch.ops.prims.convert_element_type(x, torch.float16), | ||
| 17 | + x.to(torch.float16), | ||
| 18 | + torch.full(x.shape, 2, dtype=torch.float32, device=x.device), | ||
| 19 | + x.new_empty(x.shape), | ||
| 20 | + ) | ||
| 21 | + | ||
| 22 | + x = torch.randn(32, device="cpu") | ||
| 23 | + decomps = torch._decomp.core_aten_decompositions() | ||
| 24 | + fx = make_fx(fn, decomposition_table=decomps)(x) | ||
| 25 | + | ||
| 26 | + self.assertExpectedInline( | ||
| 27 | + fx.code.lstrip(), | ||
| 28 | + """\ | ||
| 29 | +def forward(self, x_1): | ||
| 30 | + convert_element_type = torch.ops.prims.convert_element_type.default(x_1, torch.float16) | ||
| 31 | + _to_copy = torch.ops.aten._to_copy.default(x_1, dtype = torch.float16); x_1 = None | ||
| 32 | + full = torch.ops.aten.full.default([32], 2, dtype = torch.float32, device = device(type='cpu'), pin_memory = False) | ||
| 33 | + empty = torch.ops.aten.empty.memory_format([32], dtype = torch.float32, layout = torch.strided, device = device(type='cpu'), pin_memory = False) | ||
| 34 | + return (convert_element_type, _to_copy, full, empty) | ||
| 35 | + """, # NOQA: B950 | ||
| 36 | + ) | ||
| 37 | + | ||
| 38 | + fp64_model, fp64_examples = debug_utils.cast_to_fp64(fx, (x,)) | ||
| 39 | + self.assertEqual(fp64_examples, (x.to(torch.float64),)) | ||
| 40 | + | ||
| 41 | + self.assertExpectedInline( | ||
| 42 | + fx.code.lstrip(), | ||
| 43 | + """\ | ||
| 44 | +def forward(self, x_1): | ||
| 45 | + convert_element_type = torch.ops.prims.convert_element_type.default(x_1, torch.float64) | ||
| 46 | + _to_copy = torch.ops.aten._to_copy.default(x_1, dtype = torch.float64); x_1 = None | ||
| 47 | + full = torch.ops.aten.full.default([32], 2, dtype = torch.float64, device = device(type='cpu'), pin_memory = False) | ||
| 48 | + empty = torch.ops.aten.empty.memory_format([32], dtype = torch.float64, layout = torch.strided, device = device(type='cpu'), pin_memory = False) | ||
| 49 | + return (convert_element_type, _to_copy, full, empty) | ||
| 50 | + """, # NOQA: B950 | ||
| 51 | + ) | ||
| 52 | + | ||
| 53 | + | ||
| 54 | +if __name__ == "__main__": | ||
| 55 | + from torch._dynamo.test_case import run_tests | ||
| 56 | + | ||
| 57 | + run_tests() | ||
| @@ -0,0 +1,612 @@ | |||
| 1 | +# Owner(s): ["module: dynamo"] | ||
| 2 | + | ||
| 3 | +import contextlib | ||
| 4 | +import functools | ||
| 5 | + | ||
| 6 | +import torch | ||
| 7 | +import torch_npu | ||
| 8 | +import torch._dynamo | ||
| 9 | +import torch._dynamo.test_case | ||
| 10 | +import torch._dynamo.testing | ||
| 11 | +from functorch.compile import nop | ||
| 12 | +from torch._dynamo import compiled_autograd | ||
| 13 | +from torch._functorch.aot_autograd import aot_module_simplified | ||
| 14 | + | ||
| 15 | + | ||
| 16 | +def compiler_fn(gm): | ||
| 17 | + return torch._dynamo.optimize("inductor", nopython=True, dynamic=True)(gm) | ||
| 18 | + | ||
| 19 | + | ||
| 20 | +def global_hook_0(grad): | ||
| 21 | + return grad * 4 | ||
| 22 | + | ||
| 23 | + | ||
| 24 | +def global_hook_1(grad): | ||
| 25 | + return grad / 2 | ||
| 26 | + | ||
| 27 | + | ||
| 28 | +def global_hook_2(grad): | ||
| 29 | + return grad * 3 | ||
| 30 | + | ||
| 31 | + | ||
| 32 | +h0 = None | ||
| 33 | + | ||
| 34 | + | ||
| 35 | +class HooksTests(torch._dynamo.test_case.TestCase): | ||
| 36 | + def test_tensor_only_register_hook_in_graph_lambda(self): | ||
| 37 | + def fn(x): | ||
| 38 | + x.register_hook(lambda grad: grad * 2) | ||
| 39 | + return x | ||
| 40 | + | ||
| 41 | + cnts = torch._dynamo.testing.CompileCounter() | ||
| 42 | + fn = torch._dynamo.optimize(cnts)(fn) | ||
| 43 | + v = torch.tensor([0.0, 0.0, 0.0], requires_grad=True) | ||
| 44 | + v = fn(v) | ||
| 45 | + v.backward(torch.tensor([1.0, 2.0, 3.0])) | ||
| 46 | + self.assertEqual(v.grad, torch.tensor([2.0, 4.0, 6.0])) | ||
| 47 | + self.assertEqual(cnts.frame_count, 0) | ||
| 48 | + | ||
| 49 | + def test_tensor_register_hook_in_graph_lambda(self): | ||
| 50 | + def fn(x, y, z): | ||
| 51 | + x.register_hook(lambda grad: grad * 2) | ||
| 52 | + return x, y * y, z * z | ||
| 53 | + | ||
| 54 | + cnts = torch._dynamo.testing.CompileCounter() | ||
| 55 | + fn = torch._dynamo.optimize(cnts)(fn) | ||
| 56 | + v = torch.tensor([0.0, 0.0, 0.0], requires_grad=True) | ||
| 57 | + v = fn(v, torch.randn([2, 2]), torch.randn([2, 2]))[0] | ||
| 58 | + v.backward(torch.tensor([1.0, 2.0, 3.0])) | ||
| 59 | + self.assertEqual(v.grad, torch.tensor([2.0, 4.0, 6.0])) | ||
| 60 | + self.assertEqual(cnts.frame_count, 1) | ||
| 61 | + | ||
| 62 | + def test_tensor_register_hook_in_graph_break_handle_lambda(self): | ||
| 63 | + def fn(x, y, z): | ||
| 64 | + handle = x.register_hook(lambda grad: grad * 2) | ||
| 65 | + z = z * z | ||
| 66 | + handle.remove() | ||
| 67 | + x.register_hook(lambda grad: grad * 3) | ||
| 68 | + return x, y * y, z | ||
| 69 | + | ||
| 70 | + cnts = torch._dynamo.testing.CompileCounter() | ||
| 71 | + fn = torch._dynamo.optimize(cnts)(fn) | ||
| 72 | + v = torch.tensor([0.0, 0.0, 0.0], requires_grad=True) | ||
| 73 | + v = fn(v, torch.randn([2, 2]), torch.randn([2, 2]))[0] | ||
| 74 | + v.backward(torch.tensor([1.0, 2.0, 3.0])) | ||
| 75 | + self.assertEqual(v.grad, torch.tensor([3.0, 6.0, 9.0])) | ||
| 76 | + self.assertEqual(cnts.frame_count, 2) | ||
| 77 | + | ||
| 78 | + def test_tensor_register_hook_multi_handle_return(self): | ||
| 79 | + def fn(x, y, z): | ||
| 80 | + handle = x.register_hook(lambda grad: grad * 2) | ||
| 81 | + h2 = handle | ||
| 82 | + z = z * z | ||
| 83 | + return x, y * y, z, handle, h2 | ||
| 84 | + | ||
| 85 | + cnts = torch._dynamo.testing.CompileCounter() | ||
| 86 | + fn = torch._dynamo.optimize(cnts)(fn) | ||
| 87 | + v = torch.tensor([0.0, 0.0, 0.0], requires_grad=True) | ||
| 88 | + v, y, z, h, h2 = fn(v, torch.randn([2, 2]), torch.randn([2, 2])) | ||
| 89 | + v.backward(torch.tensor([1.0, 2.0, 3.0])) | ||
| 90 | + self.assertEqual(v.grad, torch.tensor([2.0, 4.0, 6.0])) | ||
| 91 | + self.assertEqual(cnts.frame_count, 1) | ||
| 92 | + self.assertNotEqual(h, None) | ||
| 93 | + self.assertNotEqual(h2, None) | ||
| 94 | + self.assertEqual(h2, h) | ||
| 95 | + | ||
| 96 | + def test_tensor_register_hook_repeated_handle_return(self): | ||
| 97 | + def fn(x, y, z): | ||
| 98 | + handle = x.register_hook(lambda grad: grad * 2) | ||
| 99 | + h2 = handle | ||
| 100 | + z = z * z | ||
| 101 | + return x, y * y, z, handle, handle | ||
| 102 | + | ||
| 103 | + cnts = torch._dynamo.testing.CompileCounter() | ||
| 104 | + fn = torch._dynamo.optimize(cnts)(fn) | ||
| 105 | + v = torch.tensor([0.0, 0.0, 0.0], requires_grad=True) | ||
| 106 | + v, y, z, h, h2 = fn(v, torch.randn([2, 2]), torch.randn([2, 2])) | ||
| 107 | + v.backward(torch.tensor([1.0, 2.0, 3.0])) | ||
| 108 | + self.assertEqual(v.grad, torch.tensor([2.0, 4.0, 6.0])) | ||
| 109 | + self.assertEqual(cnts.frame_count, 1) | ||
| 110 | + self.assertNotEqual(h, None) | ||
| 111 | + self.assertNotEqual(h2, None) | ||
| 112 | + self.assertEqual(h2, h) | ||
| 113 | + | ||
| 114 | + def test_tensor_register_hook_repeated_handle_not_local(self): | ||
| 115 | + def fn(x, y, z, mod): | ||
| 116 | + mod.handle = x.register_hook(lambda grad: grad * 2) | ||
| 117 | + z = z * z | ||
| 118 | + return x, y * y, z | ||
| 119 | + | ||
| 120 | + cnts = torch._dynamo.testing.CompileCounter() | ||
| 121 | + fn = torch._dynamo.optimize(cnts)(fn) | ||
| 122 | + v = torch.tensor([0.0, 0.0, 0.0], requires_grad=True) | ||
| 123 | + | ||
| 124 | + mod = torch.nn.Module() | ||
| 125 | + mod.handle = None | ||
| 126 | + | ||
| 127 | + v, y, z = fn(v, torch.randn([2, 2]), torch.randn([2, 2]), mod) | ||
| 128 | + v.backward(torch.tensor([1.0, 2.0, 3.0])) | ||
| 129 | + | ||
| 130 | + self.assertEqual(v.grad, torch.tensor([2.0, 4.0, 6.0])) | ||
| 131 | + self.assertEqual(cnts.frame_count, 1) | ||
| 132 | + | ||
| 133 | + self.assertNotEqual(mod.handle, None) | ||
| 134 | + | ||
| 135 | + def test_tensor_only_register_hook_in_graph_local(self): | ||
| 136 | + def local_hook(grad): | ||
| 137 | + return grad * 2 | ||
| 138 | + | ||
| 139 | + def fn(x): | ||
| 140 | + x.register_hook(local_hook) | ||
| 141 | + return x | ||
| 142 | + | ||
| 143 | + cnts = torch._dynamo.testing.CompileCounter() | ||
| 144 | + fn = torch._dynamo.optimize(cnts)(fn) | ||
| 145 | + v = torch.tensor([0.0, 0.0, 0.0], requires_grad=True) | ||
| 146 | + v = fn(v) | ||
| 147 | + v.backward(torch.tensor([1.0, 2.0, 3.0])) | ||
| 148 | + self.assertEqual(v.grad, torch.tensor([2.0, 4.0, 6.0])) | ||
| 149 | + self.assertEqual(cnts.frame_count, 0) | ||
| 150 | + | ||
| 151 | + def test_tensor_only_register_hook_in_graph_local_inner(self): | ||
| 152 | + def fn(x): | ||
| 153 | + def local_hook(grad): | ||
| 154 | + return grad * 2 | ||
| 155 | + | ||
| 156 | + z = x * x | ||
| 157 | + x.register_hook(local_hook) | ||
| 158 | + z.register_hook(local_hook) | ||
| 159 | + return x, z | ||
| 160 | + | ||
| 161 | + cnts = torch._dynamo.testing.CompileCounter() | ||
| 162 | + fn = torch._dynamo.optimize(cnts)(fn) | ||
| 163 | + v = torch.tensor([0.0, 0.0, 0.0], requires_grad=True) | ||
| 164 | + v = fn(v) | ||
| 165 | + v[0].backward(torch.tensor([1.0, 2.0, 3.0])) | ||
| 166 | + self.assertEqual(v[0].grad, torch.tensor([2.0, 4.0, 6.0])) | ||
| 167 | + self.assertEqual(cnts.frame_count, 1) | ||
| 168 | + | ||
| 169 | + def test_tensor_register_hook_in_graph_local(self): | ||
| 170 | + def local_hook(grad): | ||
| 171 | + return grad * 2 | ||
| 172 | + | ||
| 173 | + def fn(x, y, z): | ||
| 174 | + x.register_hook(local_hook) | ||
| 175 | + return x, y * y, z * z | ||
| 176 | + | ||
| 177 | + cnts = torch._dynamo.testing.CompileCounter() | ||
| 178 | + fn = torch._dynamo.optimize(cnts)(fn) | ||
| 179 | + v = torch.tensor([0.0, 0.0, 0.0], requires_grad=True) | ||
| 180 | + v = fn(v, torch.randn([2, 2]), torch.randn([2, 2]))[0] | ||
| 181 | + v.backward(torch.tensor([1.0, 2.0, 3.0])) | ||
| 182 | + self.assertEqual(v.grad, torch.tensor([2.0, 4.0, 6.0])) | ||
| 183 | + self.assertEqual(cnts.frame_count, 1) | ||
| 184 | + | ||
| 185 | + def test_tensor_register_hook_in_graph_break_handle_local(self): | ||
| 186 | + def local_hook(grad): | ||
| 187 | + return grad * 2 | ||
| 188 | + | ||
| 189 | + def local_hook2(grad): | ||
| 190 | + return grad * 3 | ||
| 191 | + | ||
| 192 | + def fn(x, y, z): | ||
| 193 | + handle = x.register_hook(local_hook) | ||
| 194 | + z = z * z | ||
| 195 | + handle.remove() | ||
| 196 | + x.register_hook(local_hook2) | ||
| 197 | + return x, y * y, z | ||
| 198 | + | ||
| 199 | + cnts = torch._dynamo.testing.CompileCounter() | ||
| 200 | + fn = torch._dynamo.optimize(cnts)(fn) | ||
| 201 | + v = torch.tensor([0.0, 0.0, 0.0], requires_grad=True) | ||
| 202 | + v = fn(v, torch.randn([2, 2]), torch.randn([2, 2]))[0] | ||
| 203 | + v.backward(torch.tensor([1.0, 2.0, 3.0])) | ||
| 204 | + | ||
| 205 | + self.assertEqual(v.grad, torch.tensor([3.0, 6.0, 9.0])) | ||
| 206 | + | ||
| 207 | + def test_tensor_register_global_hook(self): | ||
| 208 | + def fn(x): | ||
| 209 | + x.register_hook(global_hook_0) | ||
| 210 | + return x, x * x | ||
| 211 | + | ||
| 212 | + cnts = torch._dynamo.testing.CompileCounter() | ||
| 213 | + fn = torch._dynamo.optimize(cnts)(fn) | ||
| 214 | + v = torch.tensor([0.0, 0.0, 0.0], requires_grad=True) | ||
| 215 | + v = fn(v)[0] | ||
| 216 | + v.backward(torch.tensor([1.0, 2.0, 3.0])) | ||
| 217 | + self.assertEqual(v.grad, torch.tensor([4.0, 8.0, 12.0])) | ||
| 218 | + self.assertEqual(cnts.frame_count, 1) | ||
| 219 | + | ||
| 220 | + def test_tensor_register_multiple_hooks(self): | ||
| 221 | + def fn(x): | ||
| 222 | + x.register_hook(global_hook_0) # * 4 | ||
| 223 | + x.register_hook(global_hook_1) # / 2 | ||
| 224 | + x.register_hook(global_hook_2) # * 3 | ||
| 225 | + return x, x * x | ||
| 226 | + | ||
| 227 | + cnts = torch._dynamo.testing.CompileCounter() | ||
| 228 | + fn = torch._dynamo.optimize(cnts)(fn) | ||
| 229 | + v = torch.tensor([0.0, 0.0, 0.0], requires_grad=True) | ||
| 230 | + v = fn(v)[0] | ||
| 231 | + v.backward(torch.tensor([1.0, 2.0, 3.0])) | ||
| 232 | + self.assertEqual(v.grad, torch.tensor([6.0, 12.0, 18.0])) | ||
| 233 | + self.assertEqual(cnts.frame_count, 1) | ||
| 234 | + | ||
| 235 | + def test_tensor_register_multiple_hooks_handles_in_list(self): | ||
| 236 | + def fn(x): | ||
| 237 | + h_0 = x.register_hook(global_hook_0) # * 4 | ||
| 238 | + h_1 = x.register_hook(global_hook_1) # / 2 | ||
| 239 | + h_2 = x.register_hook(global_hook_2) # * 3 | ||
| 240 | + return x, x * x, h_0, h_1, h_2 | ||
| 241 | + | ||
| 242 | + cnts = torch._dynamo.testing.CompileCounter() | ||
| 243 | + fn = torch._dynamo.optimize(cnts)(fn) | ||
| 244 | + v = torch.tensor([0.0, 0.0, 0.0], requires_grad=True) | ||
| 245 | + v, r, handle_0, handle_1, handle_2 = fn(v) | ||
| 246 | + v.backward(torch.tensor([1.0, 2.0, 3.0])) | ||
| 247 | + self.assertEqual(v.grad, torch.tensor([6.0, 12.0, 18.0])) | ||
| 248 | + handle_0.remove() | ||
| 249 | + handle_1.remove() | ||
| 250 | + handle_2.remove() | ||
| 251 | + | ||
| 252 | + v.backward(torch.tensor([1.0, 2.0, 3.0])) | ||
| 253 | + # Handles gone, grad is just applied as is | ||
| 254 | + self.assertEqual(v.grad, torch.tensor([7.0, 14.0, 21.0])) | ||
| 255 | + | ||
| 256 | + self.assertEqual(cnts.frame_count, 1) | ||
| 257 | + | ||
| 258 | + def test_tensor_register_global_hooks_handles_in_list(self): | ||
| 259 | + def fn(x): | ||
| 260 | + global h0 | ||
| 261 | + h0 = x.register_hook(global_hook_0) # * 4 | ||
| 262 | + return x, x * x | ||
| 263 | + | ||
| 264 | + cnts = torch._dynamo.testing.CompileCounter() | ||
| 265 | + fn = torch._dynamo.optimize(cnts)(fn) | ||
| 266 | + v = torch.tensor([0.0, 0.0, 0.0], requires_grad=True) | ||
| 267 | + v, r = fn(v) | ||
| 268 | + | ||
| 269 | + self.assertIsNotNone(h0) | ||
| 270 | + v.backward(torch.tensor([1.0, 2.0, 3.0])) | ||
| 271 | + self.assertEqual(v.grad, torch.tensor([4.0, 8.0, 12.0])) | ||
| 272 | + h0.remove() | ||
| 273 | + | ||
| 274 | + v.backward(torch.tensor([1.0, 2.0, 3.0])) | ||
| 275 | + # Handles gone, grad is just applied as is | ||
| 276 | + self.assertEqual(v.grad, torch.tensor([5.0, 10.0, 15.0])) | ||
| 277 | + | ||
| 278 | + # NYI! | ||
| 279 | + self.assertEqual(cnts.frame_count, 0) | ||
| 280 | + | ||
| 281 | + def test_intermediary_hooks(self): | ||
| 282 | + # Graph breaks because compiled_autograd is not set | ||
| 283 | + def simple_hook(g): | ||
| 284 | + return g * 2 | ||
| 285 | + | ||
| 286 | + def f(x): | ||
| 287 | + y = x + 1 | ||
| 288 | + y.register_hook(simple_hook) | ||
| 289 | + z = y + 1 | ||
| 290 | + return z | ||
| 291 | + | ||
| 292 | + out = torch.randn(1, requires_grad=True) | ||
| 293 | + cnts = torch._dynamo.testing.CompileCounter() | ||
| 294 | + fn = torch._dynamo.optimize(cnts, nopython=False)(f) | ||
| 295 | + res = fn(out) | ||
| 296 | + res.backward() | ||
| 297 | + self.assertEqual(res, f(out)) | ||
| 298 | + self.assertEqual(cnts.frame_count, 2) | ||
| 299 | + self.assertEqual(out.grad, torch.Tensor([2.0])) | ||
| 300 | + | ||
| 301 | + def test_intermediary_hooks_same_on_aot_eager(self): | ||
| 302 | + def my_hook(grad, *, k=0): | ||
| 303 | + return grad + k | ||
| 304 | + | ||
| 305 | + class MyMod(torch.nn.Module): | ||
| 306 | + def forward(self, x): | ||
| 307 | + y = x.mul(2) | ||
| 308 | + hook1 = functools.partial(my_hook, k=3) | ||
| 309 | + hook2 = functools.partial(my_hook, k=4) | ||
| 310 | + y.register_hook(hook1) | ||
| 311 | + y.register_hook(hook2) | ||
| 312 | + z = y.mul(3) | ||
| 313 | + return (z,) | ||
| 314 | + | ||
| 315 | + mod = MyMod() | ||
| 316 | + x0 = torch.ones(4, requires_grad=True) | ||
| 317 | + eager_out = mod(x0) | ||
| 318 | + eager_out[0].backward(torch.ones(4)) | ||
| 319 | + | ||
| 320 | + x1 = torch.ones(4, requires_grad=True) | ||
| 321 | + mod_compiled = aot_module_simplified(mod, (x1,), nop) | ||
| 322 | + aot_out = mod_compiled(x1) | ||
| 323 | + aot_out[0].backward(torch.ones(4)) | ||
| 324 | + | ||
| 325 | + x2 = torch.ones(4, requires_grad=True) | ||
| 326 | + with compiled_autograd.enable(compiler_fn): | ||
| 327 | + dynamo_out = torch._dynamo.optimize("aot_eager", nopython=True)(mod)(x2) | ||
| 328 | + dynamo_out[0].backward(torch.ones(4)) | ||
| 329 | + | ||
| 330 | + self.assertEqual(dynamo_out, aot_out) | ||
| 331 | + self.assertEqual(dynamo_out, eager_out) | ||
| 332 | + | ||
| 333 | + self.assertEqual(x0.grad, x1.grad) | ||
| 334 | + self.assertEqual(x0.grad, x2.grad) | ||
| 335 | + | ||
| 336 | + def test_input_hooks_same(self): | ||
| 337 | + backends = ["eager", "aot_eager", "inductor"] | ||
| 338 | + for backend in backends: | ||
| 339 | + | ||
| 340 | + def my_hook(grad, *, k=0): | ||
| 341 | + return grad + k | ||
| 342 | + | ||
| 343 | + hook = functools.partial(my_hook, k=3) | ||
| 344 | + | ||
| 345 | + class MyMod(torch.nn.Module): | ||
| 346 | + def forward(self, x): | ||
| 347 | + x.register_hook(hook) | ||
| 348 | + y = x.mul(2) | ||
| 349 | + z = y.mul(3) | ||
| 350 | + return (z,) | ||
| 351 | + | ||
| 352 | + mod = MyMod() | ||
| 353 | + x0 = torch.ones(4, requires_grad=True) | ||
| 354 | + eager_out = mod(x0) | ||
| 355 | + eager_out[0].backward(torch.ones(4)) | ||
| 356 | + | ||
| 357 | + x1 = torch.ones(4, requires_grad=True) | ||
| 358 | + mod_compiled = aot_module_simplified(mod, (x1,), nop) | ||
| 359 | + aot_out = mod_compiled(x1) | ||
| 360 | + aot_out[0].backward(torch.ones(4)) | ||
| 361 | + | ||
| 362 | + x2 = torch.ones(4, requires_grad=True) | ||
| 363 | + dynamo_out = torch._dynamo.optimize(backend, nopython=True)(mod)(x2) | ||
| 364 | + with compiled_autograd.enable(compiler_fn): | ||
| 365 | + dynamo_out[0].backward(torch.ones(4)) | ||
| 366 | + | ||
| 367 | + self.assertEqual(dynamo_out, aot_out) | ||
| 368 | + self.assertEqual(dynamo_out, eager_out) | ||
| 369 | + | ||
| 370 | + self.assertEqual(x0.grad, x1.grad) | ||
| 371 | + self.assertEqual(x0.grad, x2.grad) | ||
| 372 | + | ||
| 373 | + def test_intermediary_hooks_same_on_inductor(self): | ||
| 374 | + def my_hook(grad, *, k=0): | ||
| 375 | + return grad + k | ||
| 376 | + | ||
| 377 | + class MyMod(torch.nn.Module): | ||
| 378 | + def forward(self, x): | ||
| 379 | + y = x.mul(2) | ||
| 380 | + hook1 = functools.partial(my_hook, k=3) | ||
| 381 | + hook2 = functools.partial(my_hook, k=4) | ||
| 382 | + y.register_hook(hook1) | ||
| 383 | + y.register_hook(hook2) | ||
| 384 | + z = y.mul(3) | ||
| 385 | + return (z,) | ||
| 386 | + | ||
| 387 | + mod = MyMod() | ||
| 388 | + x0 = torch.ones(4, requires_grad=True) | ||
| 389 | + eager_out = mod(x0) | ||
| 390 | + eager_out[0].backward(torch.ones(4)) | ||
| 391 | + | ||
| 392 | + x1 = torch.ones(4, requires_grad=True) | ||
| 393 | + mod_compiled = aot_module_simplified(mod, (x1,), nop) | ||
| 394 | + aot_out = mod_compiled(x1) | ||
| 395 | + aot_out[0].backward(torch.ones(4)) | ||
| 396 | + | ||
| 397 | + x2 = torch.ones(4, requires_grad=True) | ||
| 398 | + with compiled_autograd.enable(compiler_fn): | ||
| 399 | + dynamo_out = torch._dynamo.optimize("inductor", nopython=True)(mod)(x2) | ||
| 400 | + dynamo_out[0].backward(torch.ones(4)) | ||
| 401 | + | ||
| 402 | + self.assertEqual(dynamo_out, aot_out) | ||
| 403 | + self.assertEqual(dynamo_out, eager_out) | ||
| 404 | + | ||
| 405 | + self.assertEqual(x0.grad, x1.grad) | ||
| 406 | + self.assertEqual(x0.grad, x2.grad) | ||
| 407 | + | ||
| 408 | + def test_complex_state_mutation_in_intermediary_hooks_same_on_inductor(self): | ||
| 409 | + class SomePyClass: | ||
| 410 | + count = 0 | ||
| 411 | + | ||
| 412 | + def do_stuff(self, grad): | ||
| 413 | + if self.count % 2 == 0: | ||
| 414 | + r = grad * grad | ||
| 415 | + else: | ||
| 416 | + r = grad + grad | ||
| 417 | + self.count += 1 | ||
| 418 | + return r | ||
| 419 | + | ||
| 420 | + def complex_state_touching_hook(grad, *, obj): | ||
| 421 | + return obj.do_stuff(grad) | ||
| 422 | + | ||
| 423 | + class MyMod(torch.nn.Module): | ||
| 424 | + def forward(self, x, obj): | ||
| 425 | + y = x.mul(2) | ||
| 426 | + hook1 = functools.partial(complex_state_touching_hook, obj=obj) | ||
| 427 | + hook2 = functools.partial(complex_state_touching_hook, obj=obj) | ||
| 428 | + y.register_hook(hook1) | ||
| 429 | + y.register_hook(hook2) | ||
| 430 | + z = y.mul(3) | ||
| 431 | + return (z,) | ||
| 432 | + | ||
| 433 | + mod = MyMod() | ||
| 434 | + obj = SomePyClass() | ||
| 435 | + x0 = torch.ones(4, requires_grad=True) | ||
| 436 | + eager_out = mod(x0, obj) | ||
| 437 | + eager_out[0].backward(torch.ones(4)) | ||
| 438 | + | ||
| 439 | + # Eager 2 | ||
| 440 | + self.assertEqual(obj.count, 2) | ||
| 441 | + x2 = torch.ones(4, requires_grad=True) | ||
| 442 | + with compiled_autograd.enable(compiler_fn): | ||
| 443 | + dynamo_out = torch._dynamo.optimize("inductor", nopython=True)(mod)(x2, obj) | ||
| 444 | + dynamo_out[0].backward(torch.ones(4)) | ||
| 445 | + | ||
| 446 | + self.assertEqual(dynamo_out, eager_out) | ||
| 447 | + | ||
| 448 | + # Eager 2 + compiled 2 | ||
| 449 | + self.assertEqual(obj.count, 4) | ||
| 450 | + self.assertEqual(x0.grad, x2.grad) | ||
| 451 | + | ||
| 452 | + def test_complex_state_mutation_in_intermediary_hooks_same_on_inductor_with_graph_break( | ||
| 453 | + self, | ||
| 454 | + ): | ||
| 455 | + class SomePyClass: | ||
| 456 | + grad_as_str = "None" | ||
| 457 | + count = 0 | ||
| 458 | + | ||
| 459 | + def write_grad_as_str_and_do_stuff(self, grad): | ||
| 460 | + self.grad_as_str = str(grad) | ||
| 461 | + if self.count % 2 == 0: | ||
| 462 | + r = grad * grad | ||
| 463 | + else: | ||
| 464 | + r = grad + grad | ||
| 465 | + print("Break!") | ||
| 466 | + self.count += 1 | ||
| 467 | + return r | ||
| 468 | + | ||
| 469 | + def complex_state_touching_hook(grad, *, obj): | ||
| 470 | + return obj.write_grad_as_str_and_do_stuff(grad) | ||
| 471 | + | ||
| 472 | + class MyMod(torch.nn.Module): | ||
| 473 | + def forward(self, x, obj): | ||
| 474 | + y = x.mul(2) | ||
| 475 | + hook1 = functools.partial(complex_state_touching_hook, obj=obj) | ||
| 476 | + hook2 = functools.partial(complex_state_touching_hook, obj=obj) | ||
| 477 | + y.register_hook(hook1) | ||
| 478 | + y.register_hook(hook2) | ||
| 479 | + z = y.mul(3) | ||
| 480 | + return (z,) | ||
| 481 | + | ||
| 482 | + mod = MyMod() | ||
| 483 | + obj = SomePyClass() | ||
| 484 | + x0 = torch.ones(4, requires_grad=True) | ||
| 485 | + eager_out = mod(x0, obj) | ||
| 486 | + eager_out[0].backward(torch.ones(4)) | ||
| 487 | + | ||
| 488 | + x2 = torch.ones(4, requires_grad=True) | ||
| 489 | + with compiled_autograd.enable(compiler_fn): | ||
| 490 | + dynamo_out = torch._dynamo.optimize("inductor", nopython=True)(mod)(x2, obj) | ||
| 491 | + with self.assertRaisesRegex( | ||
| 492 | + torch._dynamo.exc.Unsupported, ".*BuiltinVariable\\(str\\).*" | ||
| 493 | + ): | ||
| 494 | + dynamo_out[0].backward(torch.ones(4)) | ||
| 495 | + | ||
| 496 | + self.assertEqual(obj.count, 2) | ||
| 497 | + | ||
| 498 | + def test_no_recompile_on_hook_identity_change(self): | ||
| 499 | + def my_hook(grad, k=0): | ||
| 500 | + return grad + k | ||
| 501 | + | ||
| 502 | + def my_hook2(grad): | ||
| 503 | + return grad * 2 | ||
| 504 | + | ||
| 505 | + class MyMod(torch.nn.Module): | ||
| 506 | + def forward(self, x): | ||
| 507 | + y = x.mul(2) | ||
| 508 | + y.register_hook(my_hook) | ||
| 509 | + y.register_hook(my_hook) | ||
| 510 | + z = y.mul(3) | ||
| 511 | + return (z,) | ||
| 512 | + | ||
| 513 | + mod = MyMod() | ||
| 514 | + x0 = torch.ones(4, requires_grad=True) | ||
| 515 | + eager_out = mod(x0) | ||
| 516 | + eager_out[0].backward(torch.ones(4)) | ||
| 517 | + | ||
| 518 | + x1 = torch.ones(4, requires_grad=True) | ||
| 519 | + with compiled_autograd.enable(compiler_fn): | ||
| 520 | + cnts = torch._dynamo.testing.CompileCounterWithBackend("aot_eager") | ||
| 521 | + comp_mod = torch._dynamo.optimize(cnts, nopython=True)(mod) | ||
| 522 | + comp_out = comp_mod(x1) | ||
| 523 | + comp_out[0].backward(torch.ones(4)) | ||
| 524 | + | ||
| 525 | + self.assertEqual(cnts.frame_count, 1) | ||
| 526 | + my_hook = my_hook2 # noqa: F811 | ||
| 527 | + self.assertEqual(x0.grad, x1.grad) | ||
| 528 | + | ||
| 529 | + eager_out = mod(x0) | ||
| 530 | + eager_out[0].backward(torch.ones(4)) | ||
| 531 | + | ||
| 532 | + comp_out = comp_mod(x1) | ||
| 533 | + | ||
| 534 | + self.assertEqual(cnts.frame_count, 2) | ||
| 535 | + comp_out[0].backward(torch.ones(4)) | ||
| 536 | + self.assertEqual(x0.grad, x1.grad) | ||
| 537 | + | ||
| 538 | + def test_functools_arg_vary(self): | ||
| 539 | + def pre_hook(grad, *, k): | ||
| 540 | + return grad * k | ||
| 541 | + | ||
| 542 | + hook = functools.partial(pre_hook, k=1) | ||
| 543 | + | ||
| 544 | + | ||
| 545 | + def h(x): | ||
| 546 | + y = x.mul(2) | ||
| 547 | + y.register_hook(hook) | ||
| 548 | + return y.mul(3) | ||
| 549 | + | ||
| 550 | + with compiled_autograd.enable(torch.compile(backend="eager", fullgraph=True)): | ||
| 551 | + x = torch.randn(2, requires_grad=True) | ||
| 552 | + h(x).sum().backward() | ||
| 553 | + orig_grad = x.grad | ||
| 554 | + x.grad = None | ||
| 555 | + | ||
| 556 | + hook = functools.partial(pre_hook, k=2) | ||
| 557 | + h(x).sum().backward() | ||
| 558 | + self.assertEqual(orig_grad * 2, x.grad) | ||
| 559 | + | ||
| 560 | + def test_post_acc_grad_hook(self): | ||
| 561 | + def hook(input_t): | ||
| 562 | + input_t.mul_(input_t.grad) | ||
| 563 | + input_t.grad.mul_(5) | ||
| 564 | + | ||
| 565 | + def reg_and_mul(x, y): | ||
| 566 | + x.register_post_accumulate_grad_hook(hook) | ||
| 567 | + return x * y | ||
| 568 | + | ||
| 569 | + cnts = None | ||
| 570 | + | ||
| 571 | + def test_fn(fn): | ||
| 572 | + fn(x, y) | ||
| 573 | + b = torch.tensor([2.0, 2.0, 2.0], requires_grad=True) | ||
| 574 | + x.backward(b) | ||
| 575 | + if cnts: | ||
| 576 | + self.assertEqual(cnts.frame_count, 1) | ||
| 577 | + # These same exact assertions run on both eager and compiled | ||
| 578 | + # X goes to x*2 becaue of mul_ | ||
| 579 | + self.assertEqual(x, torch.tensor([0.5, 0.5, 0.5]) * 2) | ||
| 580 | + # This test proves grad aliasing works - | ||
| 581 | + self.assertEqual(x.grad, b * 5) | ||
| 582 | + | ||
| 583 | + # Eager values | ||
| 584 | + x = torch.tensor([0.5, 0.5, 0.5], requires_grad=True) | ||
| 585 | + y = torch.tensor([1.0, 2.0, 3.0], requires_grad=True) | ||
| 586 | + test_fn(reg_and_mul) | ||
| 587 | + | ||
| 588 | + # Compiled | ||
| 589 | + for backend in ["eager", "aot_eager", "inductor"]: | ||
| 590 | + for compiled_bwd in [False, True]: | ||
| 591 | + torch._dynamo.reset() | ||
| 592 | + x = torch.tensor([0.5, 0.5, 0.5], requires_grad=True) | ||
| 593 | + y = torch.tensor([1.0, 2.0, 3.0], requires_grad=True) | ||
| 594 | + | ||
| 595 | + cnts = torch._dynamo.testing.CompileCounterWithBackend(backend) | ||
| 596 | + compiled_fn = torch._dynamo.optimize(cnts, nopython=True)(reg_and_mul) | ||
| 597 | + | ||
| 598 | + compiled_bwd_ctx = ( | ||
| 599 | + compiled_autograd.enable( | ||
| 600 | + torch.compile(backend=backend, fullgraph=True) | ||
| 601 | + ) | ||
| 602 | + if compiled_bwd | ||
| 603 | + else contextlib.nullcontext() | ||
| 604 | + ) | ||
| 605 | + with compiled_bwd_ctx: | ||
| 606 | + test_fn(compiled_fn) | ||
| 607 | + | ||
| 608 | + | ||
| 609 | +if __name__ == "__main__": | ||
| 610 | + from torch._dynamo.test_case import run_tests | ||
| 611 | + | ||
| 612 | + run_tests() | ||
| @@ -0,0 +1,200 @@ | |||
| 1 | +""" | ||
| 2 | +PYTEST_DONT_REWRITE (prevents pytest from rewriting assertions, which interferes | ||
| 3 | +with test_adam in OptimizerTests) | ||
| 4 | +""" | ||
| 5 | +import functools | ||
| 6 | + | ||
| 7 | +# Owner(s): ["module: dynamo"] | ||
| 8 | + | ||
| 9 | +import inspect | ||
| 10 | + | ||
| 11 | +import torch | ||
| 12 | +import torch_npu | ||
| 13 | +import torch._dynamo | ||
| 14 | +import torch._dynamo.test_case | ||
| 15 | +import torch._dynamo.testing | ||
| 16 | +from torch.nn import Parameter | ||
| 17 | + | ||
| 18 | +input1 = torch.ones([10, 10]) | ||
| 19 | +model = torch.nn.Sequential(*[torch.nn.Linear(10, 10) for _ in range(2)]) | ||
| 20 | +model(input1).sum().backward() | ||
| 21 | + | ||
| 22 | + | ||
| 23 | +def get_optimizer_step(opt_arg, closure=None): | ||
| 24 | + # run the patcher so that step has the expected structure | ||
| 25 | + torch._dynamo.eval_frame.TorchPatcher.patch() | ||
| 26 | + | ||
| 27 | + # unwrap step to avoid a deliberate graph break due to | ||
| 28 | + # a limitation of functionalization/no_grad detection | ||
| 29 | + # see the [Note on graph break] in optimizer.py | ||
| 30 | + # This ignores the outer _use_grad_if_differentiable wrapper, which is fine for now | ||
| 31 | + # as dynamo does not support differentiable optimizers anyway | ||
| 32 | + step_fn = opt_arg.step.__wrapped__ | ||
| 33 | + if closure is not None: | ||
| 34 | + | ||
| 35 | + def fn(): | ||
| 36 | + step_fn(opt_arg, closure) | ||
| 37 | + | ||
| 38 | + else: | ||
| 39 | + | ||
| 40 | + def fn(): | ||
| 41 | + step_fn(opt_arg) | ||
| 42 | + | ||
| 43 | + return fn | ||
| 44 | + | ||
| 45 | + | ||
| 46 | +def make_test(optim_cls, closure=None, **kwargs): | ||
| 47 | + opt = optim_cls(model.parameters(), **kwargs) | ||
| 48 | + | ||
| 49 | + def test_fn(self): | ||
| 50 | + nonlocal opt | ||
| 51 | + | ||
| 52 | + fn = get_optimizer_step(opt, closure=closure) | ||
| 53 | + | ||
| 54 | + with torch.set_grad_enabled(False): | ||
| 55 | + torch.compile(fn, backend="eager", fullgraph=True)() | ||
| 56 | + | ||
| 57 | + return test_fn | ||
| 58 | + | ||
| 59 | + | ||
| 60 | +class OptimizerTests(torch._dynamo.test_case.TestCase): | ||
| 61 | + test_sgd = make_test(torch.optim.SGD, lr=0.01) | ||
| 62 | + # lgbfs has data-dependent control and internally iterates | ||
| 63 | + # calling the closure | ||
| 64 | + # do for later mlazos: re-enable once we have latest pytorch with FakeTensor fix #497 | ||
| 65 | + # test_lbfgs = make_test( | ||
| 66 | + # torch.optim.LBFGS, exp_frame_cnt=3, closure=lambda: model(input).sum() | ||
| 67 | + # ) | ||
| 68 | + | ||
| 69 | + # Has data dependent control for rectification (needs symint) | ||
| 70 | + # RAdam has data-dependent control which breaks the graph; | ||
| 71 | + # furthermore, the break is inside a for loop, so we bail on the frame | ||
| 72 | + # entirely. This is basically an xfail; if the frame count goes up | ||
| 73 | + # you done good | ||
| 74 | + # test_radam = unittest.skipIf(IS_FBCODE, "TypeError: _use_grad() missing")( | ||
| 75 | + # make_test(torch.optim.RAdam, exp_graph_count=0) | ||
| 76 | + # ) | ||
| 77 | + | ||
| 78 | + | ||
| 79 | +# exclude SparseAdam because other areas of the stack don't support it yet | ||
| 80 | +# the others are handled specially above | ||
| 81 | +exclude = { | ||
| 82 | + "SGD", # Handled above | ||
| 83 | + "Optimizer", | ||
| 84 | + "SparseAdam", # Unsupported | ||
| 85 | + "LBFGS", # Unsupported | ||
| 86 | + "RAdam", # Has data dependent control for rectification (needs symint) | ||
| 87 | +} | ||
| 88 | + | ||
| 89 | + | ||
| 90 | +def check_opt(opt_ipt): | ||
| 91 | + if inspect.isclass(opt_ipt) and issubclass(opt_ipt, torch.optim.Optimizer) and opt_ipt.__name__ not in exclude: | ||
| 92 | + return True | ||
| 93 | + return False | ||
| 94 | + | ||
| 95 | + | ||
| 96 | +optimizers = [ | ||
| 97 | + opt | ||
| 98 | + for opt in torch.optim.__dict__.values() | ||
| 99 | + if check_opt(opt) | ||
| 100 | +] | ||
| 101 | + | ||
| 102 | + | ||
| 103 | +for opt in optimizers: | ||
| 104 | + setattr(OptimizerTests, "test_" + opt.__name__.lower(), make_test(opt)) | ||
| 105 | + | ||
| 106 | + | ||
| 107 | +class MyOptimizer(torch.optim.Optimizer): | ||
| 108 | + def __init__(self, params): | ||
| 109 | + super().__init__(params, {}) | ||
| 110 | + | ||
| 111 | + def _init_group(self, params, group): | ||
| 112 | + any_complex = False | ||
| 113 | + for p in group["params"]: | ||
| 114 | + params.append(p) | ||
| 115 | + any_complex |= p.is_complex() | ||
| 116 | + return any_complex | ||
| 117 | + | ||
| 118 | + def step(self): | ||
| 119 | + for group in self.param_groups: | ||
| 120 | + params = [] | ||
| 121 | + any_complex = self._init_group(params, group) | ||
| 122 | + if any_complex: | ||
| 123 | + params[0] -= 1 | ||
| 124 | + else: | ||
| 125 | + params[0] += 1 | ||
| 126 | + | ||
| 127 | + | ||
| 128 | +class End2EndTests(torch._dynamo.test_case.TestCase): | ||
| 129 | + # see torchdynamo issues 1604 | ||
| 130 | + def test_optimizing_over_tensor_with_requires_grad(self): | ||
| 131 | + class Net(torch.nn.Module): | ||
| 132 | + def forward(self, x, y): | ||
| 133 | + z = torch.bmm(x, y) | ||
| 134 | + z = torch.flatten(z, 1) | ||
| 135 | + return z | ||
| 136 | + | ||
| 137 | + def training_iter_fn(batch, model, optimizer): | ||
| 138 | + optimizer.zero_grad() | ||
| 139 | + out = model(**batch) | ||
| 140 | + target = torch.tensor([0, 7]) | ||
| 141 | + loss = torch.nn.CrossEntropyLoss()(out, target) | ||
| 142 | + loss.backward() | ||
| 143 | + optimizer.step() | ||
| 144 | + return loss | ||
| 145 | + | ||
| 146 | + net = Net() | ||
| 147 | + input_1 = torch.randn(2, 1, 4) | ||
| 148 | + input_2 = torch.randn(2, 4, 8, requires_grad=True) | ||
| 149 | + optimizer = torch.optim.Adam([input_2], lr=0.1) | ||
| 150 | + | ||
| 151 | + cnts = torch._dynamo.testing.CompileCounter() | ||
| 152 | + opt_training_iter_fn = torch._dynamo.optimize(cnts)(training_iter_fn) | ||
| 153 | + batch = {"x": input_1, "y": input_2} | ||
| 154 | + for _ in range(2): | ||
| 155 | + opt_training_iter_fn(batch, net, optimizer) | ||
| 156 | + self.assertEqual(cnts.frame_count, 2) | ||
| 157 | + | ||
| 158 | + def test_state_dict(self): | ||
| 159 | + | ||
| 160 | + def _test_state_dict(weight, bias, ipt): | ||
| 161 | + def fn_base(optimizer, weight, bias): | ||
| 162 | + optimizer.zero_grad() | ||
| 163 | + i = ipt | ||
| 164 | + loss = (weight.mv(i) + bias).pow(2).sum() | ||
| 165 | + loss.backward() | ||
| 166 | + return loss | ||
| 167 | + | ||
| 168 | + optimizer = torch.optim.Adagrad([weight, bias]) | ||
| 169 | + fn = functools.partial(fn_base, optimizer, weight, bias) | ||
| 170 | + return optimizer, fn | ||
| 171 | + | ||
| 172 | + optimizer, fn = _test_state_dict( | ||
| 173 | + Parameter(torch.randn(10, 5)), | ||
| 174 | + Parameter(torch.randn(10)), | ||
| 175 | + torch.randn(5, requires_grad=True), | ||
| 176 | + ) | ||
| 177 | + optimizer.step(fn) | ||
| 178 | + | ||
| 179 | + def test_init_group(self): | ||
| 180 | + for dtype in [torch.float32, torch.cfloat]: | ||
| 181 | + tensor = torch.randn(5, 5, dtype=dtype) | ||
| 182 | + params = Parameter(tensor.detach().clone(), requires_grad=False) | ||
| 183 | + opt_params = Parameter(tensor.detach().clone(), requires_grad=False) | ||
| 184 | + print(params, opt_params) | ||
| 185 | + | ||
| 186 | + optim = MyOptimizer([params]) | ||
| 187 | + optim.step() | ||
| 188 | + | ||
| 189 | + opt_optim = MyOptimizer([opt_params]) | ||
| 190 | + opt_step = torch.compile(backend="eager", fullgraph=True)(opt_optim.step) | ||
| 191 | + opt_step() | ||
| 192 | + print(params, opt_params) | ||
| 193 | + | ||
| 194 | + self.assertEqual(params, opt_params) | ||
| 195 | + | ||
| 196 | + | ||
| 197 | +if __name__ == "__main__": | ||
| 198 | + from torch._dynamo.test_case import run_tests | ||
| 199 | + | ||
| 200 | + run_tests() | ||
import库的顺序换一下