已合并
test: adapt Generator.clone_state test for NPU #43348
Flipped创建于 7月30日
test: adapt Generator.clone_state test for NPU #43348
已合并
共 1 个文件变更+91-0
| @@ -0,0 +1,91 @@ | |||
| 1 | +diff --git a/test/test_cuda.py b/test/test_cuda.py | ||
| 2 | +index b33f21e..4207559 100644 | ||
| 3 | +--- a/test/test_cuda.py | ||
| 4 | ++++ b/test/test_cuda.py | ||
| 5 | + import torch | ||
| 6 | + import torch.cuda | ||
| 7 | + import torch.nn as nn | ||
| 8 | + from torch import inf, nan | ||
| 9 | ++import torch_npu | ||
| 10 | ++from torch_npu.contrib import transfer_to_npu | ||
| 11 | ++ | ||
| 12 | ++if torch.cuda.get_device_capability() is None: | ||
| 13 | ++ torch.cuda.get_device_capability = lambda device=None: (0, 0) # noqa: E731 | ||
| 14 | + from torch.cuda._memory_viz import ( | ||
| 15 | + _profile_to_snapshot, | ||
| 16 | + profile_plot, | ||
| 17 | + torch.cuda.synchronize() | ||
| 18 | + | ||
| 19 | + self.assertEqual(b.sum().item(), 11000.0) | ||
| 20 | + | ||
| 21 | ++ @skipCUDAMemoryLeakCheckIf(True) | ||
| 22 | ++ @skipCUDANonDefaultStreamIf(True) | ||
| 23 | + @unittest.skipIf( | ||
| 24 | +- not TEST_CUDA_GRAPH, "CUDA >= 11.0 or ROCM >= 5.3 required for graphs" | ||
| 25 | ++ not torch.npu.is_available(), "NPU not available, skipping tests" | ||
| 26 | + ) | ||
| 27 | + def test_graphsafe_set_get_rng_state(self): | ||
| 28 | + # Define a function to create generator states, with optional graph registration | ||
| 29 | + def create_states(generator): | ||
| 30 | +- """Initializes generator states and registers them with a CUDA graph if provided.""" | ||
| 31 | +- # Ensure the CUDA generator is initialized | ||
| 32 | +- torch.rand(1, device="cuda") | ||
| 33 | ++ """Initializes generator states and registers them with a NPU graph if provided.""" | ||
| 34 | ++ # Ensure the NPU generator is initialized | ||
| 35 | ++ torch.rand(1, device="npu") | ||
| 36 | + generator.manual_seed(0) | ||
| 37 | + | ||
| 38 | + # Save the current state of the generator | ||
| 39 | + torch.cuda.synchronize() | ||
| 40 | + | ||
| 41 | + # Generate random numbers with the new generator state | ||
| 42 | + generator.graphsafe_set_state(new_state) | ||
| 43 | +- random_values.append(torch.rand(5, device="cuda", generator=generator)) | ||
| 44 | ++ random_values.append(torch.rand(5, device="npu", generator=generator)) | ||
| 45 | + | ||
| 46 | + # Generate random numbers twice with the old generator state | ||
| 47 | + generator.graphsafe_set_state(old_state) | ||
| 48 | + random_values.extend( | ||
| 49 | +- [torch.rand(5, device="cuda", generator=generator) for _ in range(2)] | ||
| 50 | ++ [torch.rand(5, device="npu", generator=generator) for _ in range(2)] | ||
| 51 | + ) | ||
| 52 | + | ||
| 53 | + return random_values | ||
| 54 | + torch.cuda.synchronize() | ||
| 55 | + new_state_offset = new_state.get_offset() | ||
| 56 | + return old_state_offset, new_state_offset | ||
| 57 | + | ||
| 58 | +- # Set up and test a new CUDA generator | ||
| 59 | +- generator = torch.Generator(device="cuda") | ||
| 60 | ++ # Set up and test a new NPU generator | ||
| 61 | ++ generator = torch.Generator(device="npu") | ||
| 62 | + generator_state = create_states(generator) | ||
| 63 | + | ||
| 64 | +- # Set up and test the default CUDA generator with a CUDA Graph | ||
| 65 | +- g = torch.cuda.CUDAGraph() | ||
| 66 | +- s = torch.cuda.Stream() | ||
| 67 | +- default_generator = torch.cuda.default_generators[0] | ||
| 68 | ++ # Set up and test the default NPU generator with a NPU Graph | ||
| 69 | ++ g = torch.npu.NPUGraph() | ||
| 70 | ++ s = torch.npu.Stream() | ||
| 71 | ++ torch.rand(1, device="npu") | ||
| 72 | ++ default_generator = torch.npu.default_generators[0] | ||
| 73 | + default_generator_state = create_states(default_generator) | ||
| 74 | + register_states_to_graph(default_generator_state, g) | ||
| 75 | + | ||
| 76 | +- # Perform random number generation within a CUDA graph | ||
| 77 | +- with torch.cuda.stream(s): | ||
| 78 | ++ # Perform random number generation within a NPU graph | ||
| 79 | ++ with torch.npu.stream(s): | ||
| 80 | + g.capture_begin() | ||
| 81 | + graphed_random_values = perform_random_generation_steps( | ||
| 82 | + default_generator_state | ||
| 83 | + torch.cuda.synchronize() | ||
| 84 | + g.capture_end() | ||
| 85 | + | ||
| 86 | + # Synchronize the streams and replay the graph | ||
| 87 | +- torch.cuda.current_stream().wait_stream(s) | ||
| 88 | ++ torch.npu.current_stream().wait_stream(s) | ||
| 89 | + for _ in range(3): | ||
| 90 | + random_values = perform_random_generation_steps(generator_state) | ||
| 91 | + g.replay() | ||