已合并
test(utils): adapt checkpoint and dataloader upstream tests on NPU #36417
Jinfan Liu创建于 5月21日
test(utils): adapt checkpoint and dataloader upstream tests on NPU #36417
已合并
共 2 个文件变更+201-0
| @@ -0,0 +1,148 @@ | |||
| 1 | +diff --git a/test/test_autograd.py b/test/test_autograd.py | ||
| 2 | +index cdb4f9e..c2d1578 100644 | ||
| 3 | +--- a/test/test_autograd.py | ||
| 4 | ++++ b/test/test_autograd.py | ||
| 5 | + from operator import mul | ||
| 6 | + from typing import TYPE_CHECKING | ||
| 7 | + | ||
| 8 | + import torch | ||
| 9 | ++import torch_npu | ||
| 10 | + import torch.autograd._functions | ||
| 11 | + import torch.autograd.forward_ad as fwAD | ||
| 12 | + from torch import inf, nan, nn | ||
| 13 | + for shape in [(1,), ()]: | ||
| 14 | + | ||
| 15 | + return fn | ||
| 16 | + | ||
| 17 | +- a = torch.randn(1, requires_grad=True) | ||
| 18 | ++ a = torch.randn(1, device="npu", requires_grad=True) | ||
| 19 | + | ||
| 20 | + # Save fewer tensors during recompute | ||
| 21 | + fn = get_non_det_fn(orig_fn=save_3_tensors, recompute_fn=save_2_tensors) | ||
| 22 | + for shape in [(1,), ()]: | ||
| 23 | + self.assertEqual(checkpoint_param.grad, param.grad) | ||
| 24 | + | ||
| 25 | + def test_checkpoint_valid_reset_on_error(self): | ||
| 26 | +- a = torch.randn(2, 2, requires_grad=True) | ||
| 27 | ++ a = torch.randn(2, 2, device="npu", requires_grad=True) | ||
| 28 | + | ||
| 29 | + with self.assertRaisesRegex( | ||
| 30 | + Exception, "torch.utils.checkpoint is incompatible" | ||
| 31 | + class TestSelectiveActivationCheckpoint(TestCase): | ||
| 32 | + def fn(x, y): | ||
| 33 | + return x.view(-1), y.sin().cos() | ||
| 34 | + | ||
| 35 | +- x = torch.tensor([1, 2, 3], dtype=torch.int64) | ||
| 36 | +- y = torch.randn(3, requires_grad=True) | ||
| 37 | ++ x = torch.tensor([1, 2, 3], dtype=torch.int64, device="npu") | ||
| 38 | ++ y = torch.randn(3, device="npu", requires_grad=True) | ||
| 39 | + | ||
| 40 | + context_fn = functools.partial( | ||
| 41 | + create_selective_checkpoint_contexts, | ||
| 42 | + class TestSelectiveActivationCheckpoint(TestCase): | ||
| 43 | + # | ||
| 44 | + # 1. size/strides are dispatched upon | ||
| 45 | + # 2. our policy saves sizes | ||
| 46 | +- ta = torch.randn(6, 2) | ||
| 47 | ++ ta = torch.randn(6, 2, device="npu") | ||
| 48 | + | ||
| 49 | + class CustomSizeDynamicShapesTensor(torch.Tensor): | ||
| 50 | + @staticmethod | ||
| 51 | + class TestSelectiveActivationCheckpoint(TestCase): | ||
| 52 | + self.assertEqual(view_out.grad_fn._saved_self_sym_sizes, [12]) | ||
| 53 | + return view_out.exp() | ||
| 54 | + | ||
| 55 | +- x = torch.randn(4, 3, requires_grad=True) | ||
| 56 | ++ x = torch.randn(4, 3, device="npu", requires_grad=True) | ||
| 57 | + x_wrapper = CustomSizeDynamicShapesTensor(x) | ||
| 58 | + context_fn = functools.partial(create_selective_checkpoint_contexts, policy_fn) | ||
| 59 | + out = checkpoint(fn, x_wrapper, use_reentrant=False, context_fn=context_fn) | ||
| 60 | + class TestSelectiveActivationCheckpoint(TestCase): | ||
| 61 | + def fn(x): | ||
| 62 | + return x.sin().sin().sin() | ||
| 63 | + | ||
| 64 | +- x = torch.randn(3, requires_grad=True) | ||
| 65 | ++ x = torch.randn(3, device="npu", requires_grad=True) | ||
| 66 | + context_fn = functools.partial( | ||
| 67 | + create_selective_checkpoint_contexts, | ||
| 68 | + Policy(), | ||
| 69 | + class TestSelectiveActivationCheckpoint(TestCase): | ||
| 70 | + with disable_gc(): | ||
| 71 | + # Case 1: If graph goes away without backward, make sure there's no reference cycle | ||
| 72 | + # keeping storage alive. | ||
| 73 | +- x = torch.randn(3, requires_grad=True) | ||
| 74 | ++ x = torch.randn(3, device="npu", requires_grad=True) | ||
| 75 | + context_fn = functools.partial( | ||
| 76 | + create_selective_checkpoint_contexts, policy_fn | ||
| 77 | + ) | ||
| 78 | + class TestSelectiveActivationCheckpoint(TestCase): | ||
| 79 | + self.assertIsNone(ref()) | ||
| 80 | + | ||
| 81 | + # Case 2: After backward, even if retain_graph=True, the storage should go away | ||
| 82 | +- x = torch.randn(3, requires_grad=True) | ||
| 83 | ++ x = torch.randn(3, device="npu", requires_grad=True) | ||
| 84 | + context_fn = functools.partial( | ||
| 85 | + create_selective_checkpoint_contexts, policy_fn | ||
| 86 | + ) | ||
| 87 | + class TestSelectiveActivationCheckpoint(TestCase): | ||
| 88 | + def fn(x): | ||
| 89 | + return x.sin().mul_(2).cos().exp() | ||
| 90 | + | ||
| 91 | +- x = torch.randn(3, requires_grad=True) | ||
| 92 | ++ x = torch.randn(3, device="npu", requires_grad=True) | ||
| 93 | + context_fn = functools.partial(create_selective_checkpoint_contexts, policy_fn) | ||
| 94 | + out = checkpoint(fn, x, use_reentrant=False, context_fn=context_fn) | ||
| 95 | + | ||
| 96 | + class TestSelectiveActivationCheckpoint(TestCase): | ||
| 97 | + with self.assertRaisesRegex(RuntimeError, "has been mutated"): | ||
| 98 | + out.sum().backward() | ||
| 99 | + | ||
| 100 | +- x = torch.randn(3, requires_grad=True) | ||
| 101 | ++ x = torch.randn(3, device="npu", requires_grad=True) | ||
| 102 | + context_fn = functools.partial( | ||
| 103 | + create_selective_checkpoint_contexts, | ||
| 104 | + policy_fn, | ||
| 105 | + class TestSelectiveActivationCheckpoint(TestCase): | ||
| 106 | + a, b = torch.var_mean(x) | ||
| 107 | + return a * b | ||
| 108 | + | ||
| 109 | +- x = torch.randn(3, requires_grad=True) | ||
| 110 | ++ x = torch.randn(3, device="npu", requires_grad=True) | ||
| 111 | + context_fn = functools.partial(create_selective_checkpoint_contexts, policy_fn) | ||
| 112 | + out = checkpoint(fn, x, use_reentrant=False, context_fn=context_fn) | ||
| 113 | + x_grad = torch.autograd.grad(out.sum(), (x,)) | ||
| 114 | + class TestSelectiveActivationCheckpoint(TestCase): | ||
| 115 | + "mylib::sin_with_extra", backward, setup_context=setup_context | ||
| 116 | + ) | ||
| 117 | + | ||
| 118 | +- x = torch.randn(3, requires_grad=True) | ||
| 119 | ++ x = torch.randn(3, device="npu", requires_grad=True) | ||
| 120 | + | ||
| 121 | + def fn(x): | ||
| 122 | + return (torch.ops.mylib.sin_with_extra(x)[0] * x.sin().exp()).sin() | ||
| 123 | + | ||
| 124 | + ops_list = [torch.ops.mylib.sin_with_extra.default] | ||
| 125 | + | ||
| 126 | +- x = torch.randn(3, requires_grad=True) | ||
| 127 | ++ x = torch.randn(3, device="npu", requires_grad=True) | ||
| 128 | + context_fn = functools.partial( | ||
| 129 | + create_selective_checkpoint_contexts, ops_list | ||
| 130 | + ) | ||
| 131 | + class TestSelectiveActivationCheckpoint(TestCase): | ||
| 132 | + def fn(x): | ||
| 133 | + return x.sin().cos().exp() | ||
| 134 | + | ||
| 135 | +- x = torch.randn(3, requires_grad=True) | ||
| 136 | ++ x = torch.randn(3, device="npu", requires_grad=True) | ||
| 137 | + context_fn = functools.partial(create_selective_checkpoint_contexts, policy_fn) | ||
| 138 | + out = checkpoint(fn, x, use_reentrant=False, context_fn=context_fn) | ||
| 139 | + out.sum().backward(retain_graph=True) | ||
| 140 | + class TestSelectiveActivationCheckpoint(TestCase): | ||
| 141 | + fwd_decisions.append(decision) | ||
| 142 | + return decision | ||
| 143 | + | ||
| 144 | +- x = torch.randn(4, requires_grad=True) | ||
| 145 | ++ x = torch.randn(4, device="npu", requires_grad=True) | ||
| 146 | + context_fn = functools.partial( | ||
| 147 | + create_selective_checkpoint_contexts, policy_fn | ||
| 148 | + ) | ||
| @@ -0,0 +1,53 @@ | |||
| 1 | +diff --git a/test/test_dataloader.py b/test/test_dataloader.py | ||
| 2 | +index df5c53b..2dac4f6 100644 | ||
| 3 | +--- a/test/test_dataloader.py | ||
| 4 | ++++ b/test/test_dataloader.py | ||
| 5 | + import unittest | ||
| 6 | + import warnings | ||
| 7 | + | ||
| 8 | + import torch | ||
| 9 | ++import torch_npu | ||
| 10 | + import torch.utils.data.datapipes as dp | ||
| 11 | + from torch import multiprocessing as mp | ||
| 12 | + from torch._utils import ExceptionWrapper | ||
| 13 | + class TestDataLoader(TestCase): | ||
| 14 | + self.dataset = TensorDataset(self.data, self.labels) | ||
| 15 | + self.persistent_workers = False | ||
| 16 | + | ||
| 17 | ++ def _move_dataset_to_npu(self): | ||
| 18 | ++ self.data = self.data.npu() | ||
| 19 | ++ self.labels = self.labels.npu() | ||
| 20 | ++ self.dataset = TensorDataset(self.data, self.labels) | ||
| 21 | ++ | ||
| 22 | + def _get_data_loader(self, dataset, **kwargs): | ||
| 23 | + persistent_workers = kwargs.get("persistent_workers", self.persistent_workers) | ||
| 24 | + if persistent_workers and kwargs.get("num_workers", 0) == 0: | ||
| 25 | + except RuntimeError as e: | ||
| 26 | + self.assertRaises(ValueError, fn) | ||
| 27 | + | ||
| 28 | + def test_sequential_nonbatch(self): | ||
| 29 | ++ self._move_dataset_to_npu() | ||
| 30 | + self._test_sequential(self._get_data_loader(self.dataset, batch_size=None)) | ||
| 31 | + | ||
| 32 | + def test_sequential_batch(self): | ||
| 33 | ++ self._move_dataset_to_npu() | ||
| 34 | + self._test_sequential(self._get_data_loader(self.dataset)) | ||
| 35 | + self._test_sequential(self._get_data_loader(self.dataset, batch_size=2)) | ||
| 36 | + | ||
| 37 | + except RuntimeError as e: | ||
| 38 | + self.assertEqual(collated, torch.tensor(arr)) | ||
| 39 | + self.assertEqual(collated.dtype, torch.bool) | ||
| 40 | + | ||
| 41 | ++ npu_batch = [ | ||
| 42 | ++ torch.tensor([1, 2], device="npu"), | ||
| 43 | ++ torch.tensor([3, 4], device="npu"), | ||
| 44 | ++ ] | ||
| 45 | ++ npu_collated = dataloader.default_collate(npu_batch) | ||
| 46 | ++ self.assertEqual(npu_collated.device.type, "npu") | ||
| 47 | ++ self.assertEqual( | ||
| 48 | ++ npu_collated, torch.tensor([[1, 2], [3, 4]], device="npu") | ||
| 49 | ++ ) | ||
| 50 | ++ | ||
| 51 | + # Should be a no-op | ||
| 52 | + arr = ["a", "b", "c"] | ||
| 53 | + self.assertEqual(arr, dataloader.default_collate(arr)) | ||