已合并
Bugfix: add tuple support for device_ids in transfer_to_npu #41617
yinglinwei创建于 7月14日
Bugfix: add tuple support for device_ids in transfer_to_npu #41617
已合并
共 3 个文件变更+34-9
| @@ -102,9 +102,12 @@ class TestTransferToNpu(TestCase): | |||
| 102 | self.assertEqual(device.type, 'meta') | 102 | self.assertEqual(device.type, 'meta') |
| 103 | 103 | ||
| 104 | def test_set_default_device(self): | 104 | def test_set_default_device(self): |
| 105 | - torch.set_default_device("cuda") | 105 | + try: |
| 106 | - a = torch.tensor(1) | 106 | + torch.set_default_device("cuda") |
| 107 | - self.assertEqual(a.device.type, 'npu') | 107 | + a = torch.tensor(1) |
| 108 | + self.assertEqual(a.device.type, 'npu') | ||
| 109 | + finally: | ||
| 110 | + torch.set_default_device(None) | ||
| 108 | 111 | ||
| 109 | def test_device_context(self): | 112 | def test_device_context(self): |
| 110 | device = torch.device('cuda') | 113 | device = torch.device('cuda') |
| @@ -404,5 +407,20 @@ class TestTransferToNpu(TestCase): | |||
| 404 | self.assertEqual(torch.cuda.default_generators, torch_npu.npu.default_generators) | 407 | self.assertEqual(torch.cuda.default_generators, torch_npu.npu.default_generators) |
| 405 | self.assertNotEqual(torch.cuda.default_generators, ()) | 408 | self.assertNotEqual(torch.cuda.default_generators, ()) |
| 406 | 409 | ||
| 410 | + def test_wrapper_cuda_device_ids_tuple(self): | ||
| 411 | + | ||
| 412 | + def mock_function(*args, **kwargs): | ||
| 413 | + return kwargs | ||
| 414 | + | ||
| 415 | + kwargs_input = {'device_ids': ('cuda:0', 'cuda:1')} | ||
| 416 | + expected_kwargs_output = {'device_ids': ('npu:0', 'npu:1')} | ||
| 417 | + kwargs_output = mock_function(**kwargs_input) | ||
| 418 | + self.assertEqual(kwargs_output, expected_kwargs_output) | ||
| 419 | + | ||
| 420 | + kwargs_input = {'device_ids': ['cuda:0', 'cuda:1']} | ||
| 421 | + expected_kwargs_output = {'device_ids': ['npu:0', 'npu:1']} | ||
| 422 | + kwargs_output = mock_function(**kwargs_input) | ||
| 423 | + self.assertEqual(kwargs_output, expected_kwargs_output) | ||
| 424 | + | ||
| 407 | if __name__ == "__main__": | 425 | if __name__ == "__main__": |
| 408 | run_tests() | 426 | run_tests() |
| @@ -103,9 +103,12 @@ class TestTransferToNpu(TestCase): | |||
| 103 | self.assertEqual(device.type, 'meta') | 103 | self.assertEqual(device.type, 'meta') |
| 104 | 104 | ||
| 105 | def test_set_default_device(self): | 105 | def test_set_default_device(self): |
| 106 | - torch.set_default_device("cuda") | 106 | + try: |
| 107 | - a = torch.tensor(1) | 107 | + torch.set_default_device("cuda") |
| 108 | - self.assertEqual(a.device.type, 'npu') | 108 | + a = torch.tensor(1) |
| 109 | + self.assertEqual(a.device.type, 'npu') | ||
| 110 | + finally: | ||
| 111 | + torch.set_default_device(None) | ||
| 109 | 112 | ||
| 110 | def test_device_context(self): | 113 | def test_device_context(self): |
| 111 | device = torch.device('cuda') | 114 | device = torch.device('cuda') |
| @@ -186,8 +186,10 @@ def _wrapper_cuda(fn): | |||
| 186 | if device is not None: | 186 | if device is not None: |
| 187 | _replace_cuda_to_npu_in_kwargs(kwargs, device_arg, device) | 187 | _replace_cuda_to_npu_in_kwargs(kwargs, device_arg, device) |
| 188 | device_ids = kwargs.get('device_ids', None) | 188 | device_ids = kwargs.get('device_ids', None) |
| 189 | - if type(device_ids) == list: | 189 | + if isinstance(device_ids, list): |
| 190 | - device_ids = _replace_cuda_to_npu_in_list(device_ids, replace_int) | 190 | + kwargs["device_ids"] = _replace_cuda_to_npu_in_list(device_ids, replace_int) |
| 191 | + elif isinstance(device_ids, tuple): | ||
| 192 | + kwargs["device_ids"] = tuple(_replace_cuda_to_npu_in_list(list(device_ids), replace_int)) | ||
| 191 | return fn(*args, **kwargs) | 193 | return fn(*args, **kwargs) |
| 192 | 194 | ||
| 193 | return decorated | 195 | return decorated |
| @@ -417,7 +419,9 @@ def _patch_nametuple(nametuple): | |||
| 417 | _replace_cuda_to_npu_in_kwargs(kwargs, device_arg, device) | 419 | _replace_cuda_to_npu_in_kwargs(kwargs, device_arg, device) |
| 418 | device_ids = kwargs.get('device_ids', None) | 420 | device_ids = kwargs.get('device_ids', None) |
| 419 | if isinstance(device_ids, list): | 421 | if isinstance(device_ids, list): |
| 420 | - device_ids = _replace_cuda_to_npu_in_list(device_ids, False) | 422 | + kwargs["device_ids"] = _replace_cuda_to_npu_in_list(device_ids, False) |
| 423 | + elif isinstance(device_ids, tuple): | ||
| 424 | + kwargs["device_ids"] = tuple(_replace_cuda_to_npu_in_list(list(device_ids), False)) | ||
| 421 | return original__new__(cls, *args, **kwargs) | 425 | return original__new__(cls, *args, **kwargs) |
| 422 | nametuple.__new__ = new_nametuple__new__ | 426 | nametuple.__new__ = new_nametuple__new__ |
| 423 | 427 | ||