已合并
Apply backend patch from testing.__init__ to test_schedule #33205
limuan创建于 4月7日
Apply backend patch from testing.__init__ to test_schedule #33205
已合并
共 2 个文件变更+5-4
| @@ -8,6 +8,7 @@ import os | |||
| 8 | from model_registry import MultiMLP | 8 | from model_registry import MultiMLP |
| 9 | 9 | ||
| 10 | import torch | 10 | import torch |
| 11 | +import torch_npu.testing | ||
| 11 | from torch.distributed.pipelining import ( | 12 | from torch.distributed.pipelining import ( |
| 12 | Schedule1F1B, | 13 | Schedule1F1B, |
| 13 | ScheduleGPipe, | 14 | ScheduleGPipe, |
| @@ -219,7 +220,6 @@ class ScheduleTest(TestCase): | |||
| 219 | ) | 220 | ) |
| 220 | with self.assertRaises(RuntimeError): | 221 | with self.assertRaises(RuntimeError): |
| 221 | ScheduleInterleavedZeroBubble([stage], 2) | 222 | ScheduleInterleavedZeroBubble([stage], 2) |
| 222 | - | ||
| 223 | torch.distributed.destroy_process_group() | 223 | torch.distributed.destroy_process_group() |
| 224 | 224 | ||
| 225 | 225 | ||
| @@ -714,7 +714,7 @@ class TestScheduleLowering(TestCase): | |||
| 714 | loss_fn=loss_fn, | 714 | loss_fn=loss_fn, |
| 715 | scale_grads=False, | 715 | scale_grads=False, |
| 716 | ) | 716 | ) |
| 717 | - schedule._load_actions( | 717 | + schedule._prepare_schedule_with_comms( |
| 718 | { | 718 | { |
| 719 | 0: self._parse_actions( | 719 | 0: self._parse_actions( |
| 720 | [ | 720 | [ |
| @@ -825,7 +825,7 @@ class TestScheduleLowering(TestCase): | |||
| 825 | num_microbatches, | 825 | num_microbatches, |
| 826 | loss_fn=loss_fn, | 826 | loss_fn=loss_fn, |
| 827 | ) | 827 | ) |
| 828 | - schedule._load_actions( | 828 | + schedule._prepare_schedule_with_comms( |
| 829 | { | 829 | { |
| 830 | 0: self._parse_actions( | 830 | 0: self._parse_actions( |
| 831 | [ | 831 | [ |
| @@ -31602,5 +31602,6 @@ | |||
| 31602 | "test_index_put_accumulate_large_tensor_npu (__main__.TestIndexingPRIVATEUSE1)": ["", ["910A"]], | 31602 | "test_index_put_accumulate_large_tensor_npu (__main__.TestIndexingPRIVATEUSE1)": ["", ["910A"]], |
| 31603 | "test_module_backcompat (__main__.TestNN)": ["", [""]], | 31603 | "test_module_backcompat (__main__.TestNN)": ["", [""]], |
| 31604 | "test_register_state_dict_pre_hook_backward_compat (__main__.TestNN)": ["", [""]], | 31604 | "test_register_state_dict_pre_hook_backward_compat (__main__.TestNN)": ["", [""]], |
| 31605 | - "test_softmax_forward_64bit_indexing_npu (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]] | 31605 | + "test_softmax_forward_64bit_indexing_npu (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]], |
| 31606 | + "test_grad_with_v_schedule (__main__.TestScheduleLowering)": ["", [""]] | ||
| 31606 | } | 31607 | } |