已合并
test: verify community 2.11.0 features and fixes #36311
kuhn7创建于 5月21日
test: verify community 2.11.0 features and fixes #36311
已合并
共 5 个文件变更+887-7
| @@ -1,3 +1,6 @@ | |||
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates | ||
| 2 | +# Owner(s): ["oncall: distributed"] | ||
| 3 | + | ||
| 1 | import unittest | 4 | import unittest |
| 2 | 5 | ||
| 3 | import torch | 6 | import torch |
| @@ -6,14 +9,15 @@ import torch.nn.functional as F | |||
| 6 | from numpy.testing import assert_array_equal | 9 | from numpy.testing import assert_array_equal |
| 7 | from torch.distributed._functional_collectives import AsyncCollectiveTensor | 10 | from torch.distributed._functional_collectives import AsyncCollectiveTensor |
| 8 | import torch.distributed._functional_collectives as funcol | 11 | import torch.distributed._functional_collectives as funcol |
| 9 | - | 12 | +from torch.distributed.device_mesh import init_device_mesh |
| 10 | -from torch.distributed._tensor import ( | 13 | +from torch.distributed.tensor import ( |
| 11 | DeviceMesh, | 14 | DeviceMesh, |
| 12 | distribute_tensor, | 15 | distribute_tensor, |
| 13 | DTensor, | 16 | DTensor, |
| 14 | - init_device_mesh, | 17 | + Partial, |
| 18 | + Replicate, | ||
| 19 | + Shard, | ||
| 15 | ) | 20 | ) |
| 16 | -from torch.distributed._tensor.placement_types import _Partial, Replicate, Shard | ||
| 17 | from torch.distributed.tensor.parallel import ( | 21 | from torch.distributed.tensor.parallel import ( |
| 18 | ColwiseParallel, | 22 | ColwiseParallel, |
| 19 | parallelize_module, | 23 | parallelize_module, |
| @@ -177,7 +181,7 @@ class DTensorTest(DTensorTestBase): | |||
| 177 | ddp_tensor = DTensor.from_local(local_tensor, device_mesh, replica_spec) | 181 | ddp_tensor = DTensor.from_local(local_tensor, device_mesh, replica_spec) |
| 178 | self.assertEqual(ddp_tensor.size(), local_tensor.size()) | 182 | self.assertEqual(ddp_tensor.size(), local_tensor.size()) |
| 179 | 183 | ||
| 180 | - partial_spec = [_Partial()] | 184 | + partial_spec = [Partial()] |
| 181 | partial_tensor = DTensor.from_local(local_tensor, device_mesh, partial_spec) | 185 | partial_tensor = DTensor.from_local(local_tensor, device_mesh, partial_spec) |
| 182 | self.assertEqual(partial_tensor.size(), local_tensor.size()) | 186 | self.assertEqual(partial_tensor.size(), local_tensor.size()) |
| 183 | 187 | ||
| @@ -336,7 +340,7 @@ class DTensorTest(DTensorTestBase): | |||
| 336 | 340 | ||
| 337 | sharded_dtensor = distribute_tensor(global_tensor, device_mesh, placements) | 341 | sharded_dtensor = distribute_tensor(global_tensor, device_mesh, placements) |
| 338 | local_out = sharded_dtensor.redistribute(placements=[Replicate()]).to_local( | 342 | local_out = sharded_dtensor.redistribute(placements=[Replicate()]).to_local( |
| 339 | - grad_placements=[_Partial()] | 343 | + grad_placements=[Partial()] |
| 340 | ) | 344 | ) |
| 341 | local_out.sum().backward() | 345 | local_out.sum().backward() |
| 342 | 346 | ||
| @@ -363,12 +367,51 @@ class DTensorTest(DTensorTestBase): | |||
| 363 | global_tensor = torch.ones(8, 3, requires_grad=True) | 367 | global_tensor = torch.ones(8, 3, requires_grad=True) |
| 364 | 368 | ||
| 365 | sharded_dtensor = distribute_tensor(global_tensor, device_mesh, placements) | 369 | sharded_dtensor = distribute_tensor(global_tensor, device_mesh, placements) |
| 366 | - local_out = sharded_dtensor.full_tensor(grad_placements=[_Partial()]) | 370 | + local_out = sharded_dtensor.full_tensor(grad_placements=[Partial()]) |
| 367 | local_out.sum().backward() | 371 | local_out.sum().backward() |
| 368 | 372 | ||
| 369 | replica_grad = sharded_dtensor.grad.full_tensor() | 373 | replica_grad = sharded_dtensor.grad.full_tensor() |
| 370 | self.assertEqual(replica_grad, global_tensor * self.world_size) | 374 | self.assertEqual(replica_grad, global_tensor * self.world_size) |
| 371 | 375 | ||
| 376 | + | ||
| 377 | + | ||
| 378 | + def test_to_local_from_local_backward(self): | ||
| 379 | + """ | ||
| 380 | + Test that to_local() followed by from_local() with Partial placement | ||
| 381 | + produces correct gradients. This validates the gradient placement mapping: | ||
| 382 | + Partial forward → Replicate gradient (see DTensor.from_local docstring). | ||
| 383 | + """ | ||
| 384 | + device_mesh = self.build_device_mesh() | ||
| 385 | + | ||
| 386 | + # Create a sharded tensor [1., 2., ...] across ranks | ||
| 387 | + # rank 0 gets [1.], rank 1 gets [2.], etc. | ||
| 388 | + global_tensor = torch.arange( | ||
| 389 | + 1.0, self.world_size + 1, device=self.device_type, requires_grad=True | ||
| 390 | + ) | ||
| 391 | + sharded_tensor = distribute_tensor(global_tensor, device_mesh, [Shard(0)]) | ||
| 392 | + | ||
| 393 | + # sum() produces a Partial DTensor | ||
| 394 | + out = sharded_tensor.sum() | ||
| 395 | + | ||
| 396 | + # to_local() + from_local() round-trip with same Partial placement | ||
| 397 | + out = DTensor.from_local( | ||
| 398 | + out.to_local(), | ||
| 399 | + out.device_mesh, | ||
| 400 | + out.placements, | ||
| 401 | + run_check=False, | ||
| 402 | + ) | ||
| 403 | + | ||
| 404 | + # full_tensor() reduces the Partial, then compute loss | ||
| 405 | + loss = out.full_tensor().sum() | ||
| 406 | + loss.backward() | ||
| 407 | + | ||
| 408 | + # Expected: each element contributes equally to the sum, so gradient is 1.0 | ||
| 409 | + # The full gradient should be [1., 1., ...] (world_size elements) | ||
| 410 | + expected_grad = torch.ones(self.world_size, device=self.device_type) | ||
| 411 | + actual_grad = sharded_tensor.grad.full_tensor() | ||
| 412 | + | ||
| 413 | + self.assertEqual(actual_grad, expected_grad) | ||
| 414 | + | ||
| 372 | 415 | ||
| 373 | 416 | ||
| 374 | def test_dtensor_new_empty_strided(self): | 417 | def test_dtensor_new_empty_strided(self): |
| @@ -16,6 +16,12 @@ from torch.distributed._tensor.placement_types import ( | |||
| 16 | Replicate, | 16 | Replicate, |
| 17 | Shard, | 17 | Shard, |
| 18 | ) | 18 | ) |
| 19 | +from torch.distributed.tensor.debug import CommDebugMode | ||
| 20 | +from torch.testing._internal.common_utils import ( | ||
| 21 | + instantiate_parametrized_tests, | ||
| 22 | + parametrize, | ||
| 23 | + run_tests, | ||
| 24 | +) | ||
| 19 | from torch.distributed.distributed_c10d import ReduceOp | 25 | from torch.distributed.distributed_c10d import ReduceOp |
| 20 | from torch.testing._internal.common_utils import run_tests | 26 | from torch.testing._internal.common_utils import run_tests |
| 21 | from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase | 27 | from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase |
| @@ -293,5 +299,30 @@ class DistElementwiseOpsTest(DTensorTestBase): | |||
| 293 | self.assertEqual(dist_res.to_local().dtype, dst_dtype) | 299 | self.assertEqual(dist_res.to_local().dtype, dst_dtype) |
| 294 | self.assertEqual(dist_res.to_local(), local_result) | 300 | self.assertEqual(dist_res.to_local(), local_result) |
| 295 | 301 | ||
| 302 | + | ||
| 303 | + | ||
| 304 | + def test_partial_propagation(self, op, reduce_op): | ||
| 305 | + # Test that torch.maximum/minimum preserves Partial("max"/"min") placements | ||
| 306 | + # since max(max(a), max(b)) == max(a, b) and min(min(a), min(b)) == min(a, b) | ||
| 307 | + device_mesh = self.build_device_mesh() | ||
| 308 | + comm_mode = CommDebugMode() | ||
| 309 | + | ||
| 310 | + input1 = torch.rand(8, 8) * self.rank | ||
| 311 | + input2 = torch.rand(8, 8) * (self.world_size - self.rank) | ||
| 312 | + | ||
| 313 | + d_input1 = DTensor.from_local(input1, device_mesh, [Partial(reduce_op)]) | ||
| 314 | + d_input2 = DTensor.from_local(input2, device_mesh, [Partial(reduce_op)]) | ||
| 315 | + | ||
| 316 | + with comm_mode: | ||
| 317 | + result = op(d_input1, d_input2) | ||
| 318 | + | ||
| 319 | + # Should not require any communication | ||
| 320 | + self.assertEqual(comm_mode.get_total_counts(), 0) | ||
| 321 | + # Result should still be Partial with the same reduce_op | ||
| 322 | + self.assertEqual(result.placements, (Partial(reduce_op),)) | ||
| 323 | + | ||
| 324 | +instantiate_parametrized_tests(DistElementwiseOpsTest) | ||
| 325 | + | ||
| 326 | + | ||
| 296 | if __name__ == "__main__": | 327 | if __name__ == "__main__": |
| 297 | run_tests() | 328 | run_tests() |
| @@ -0,0 +1,28 @@ | |||
| 1 | +# Owner(s): ["oncall: distributed"] | ||
| 2 | + | ||
| 3 | +import torch | ||
| 4 | +from torch._subclasses.fake_tensor import FakeTensorMode | ||
| 5 | +from torch.testing._internal.common_utils import run_tests, TestCase | ||
| 6 | +from torch.testing._internal.distributed.fake_pg import FakeStore | ||
| 7 | + | ||
| 8 | + | ||
| 9 | +class TestFakeDTensor(TestCase): | ||
| 10 | + def test_fake_collectives(self): | ||
| 11 | + # Ensure that we can run (non-functional) collectives under FakeTensorMode. | ||
| 12 | + # This requires the meta impls for non-functional collectives | ||
| 13 | + # to be registered at impor | ||
| 14 | + fake_mode = FakeTensorMode() | ||
| 15 | + world_size = 4 | ||
| 16 | + | ||
| 17 | + fake_store = FakeStore() | ||
| 18 | + torch.distributed.init_process_group( | ||
| 19 | + "fake", store=fake_store, rank=0, world_size=world_size | ||
| 20 | + ) | ||
| 21 | + default_pg = torch.distributed.distributed_c10d._get_default_group() | ||
| 22 | + with fake_mode: | ||
| 23 | + x = torch.randn(2, 2, device="npu") | ||
| 24 | + torch.distributed.all_reduce(x, group=default_pg) | ||
| 25 | + | ||
| 26 | + | ||
| 27 | +if __name__ == "__main__": | ||
| 28 | + run_tests() | ||
| @@ -0,0 +1,600 @@ | |||
| 1 | +# Owner(s): ["oncall: distributed"] | ||
| 2 | + | ||
| 3 | +import sys | ||
| 4 | +from functools import partial, wraps | ||
| 5 | + | ||
| 6 | +import torch | ||
| 7 | +import torch.distributed as dist | ||
| 8 | +from torch.distributed import _functional_collectives as fcols | ||
| 9 | +from torch.testing._internal.common_distributed import ( | ||
| 10 | + MultiThreadedTestCase, | ||
| 11 | + TEST_SKIPS, | ||
| 12 | +) | ||
| 13 | +from torch.testing._internal.common_utils import ( | ||
| 14 | + instantiate_parametrized_tests, | ||
| 15 | + parametrize, | ||
| 16 | + run_tests, | ||
| 17 | +) | ||
| 18 | + | ||
| 19 | + | ||
| 20 | +if not dist.is_available(): | ||
| 21 | + print("Distributed not available, skipping tests", file=sys.stderr) | ||
| 22 | + sys.exit(0) | ||
| 23 | + | ||
| 24 | + | ||
| 25 | +# Determine available devices | ||
| 26 | +DEVICE = "npu" | ||
| 27 | +devices = ["cpu"] | ||
| 28 | +if acc := torch.accelerator.current_accelerator(True): | ||
| 29 | + devices += [acc.type] | ||
| 30 | + | ||
| 31 | + | ||
| 32 | +def with_comms(func=None): | ||
| 33 | + if func is None: | ||
| 34 | + return partial(with_comms) | ||
| 35 | + | ||
| 36 | + | ||
| 37 | + def wrapper(self, *args, **kwargs): | ||
| 38 | + if ( | ||
| 39 | + torch.npu.is_available() | ||
| 40 | + and torch.accelerator.device_count() < self.world_size | ||
| 41 | + ): | ||
| 42 | + sys.exit(TEST_SKIPS[f"multi-gpu-{self.world_size}"].exit_code) | ||
| 43 | + | ||
| 44 | + self.pg = self.create_pg(device=DEVICE) | ||
| 45 | + self.device = DEVICE | ||
| 46 | + try: | ||
| 47 | + return func(self, *args, **kwargs) | ||
| 48 | + finally: | ||
| 49 | + torch.distributed.destroy_process_group() | ||
| 50 | + | ||
| 51 | + return wrapper | ||
 | |||
| 52 | + | ||
| 53 | + | ||
| 54 | + | ||
| 55 | +class TestFunctionalDifferentials(MultiThreadedTestCase): | ||
| 56 | + | ||
| 57 | + def world_size(self): | ||
| 58 | + return 4 | ||
| 59 | + | ||
| 60 | + def setUp(self): | ||
| 61 | + super().setUp() | ||
| 62 | + self._spawn_threads() | ||
| 63 | + | ||
| 64 | + # ============================================================ | ||
| 65 | + # Forward Correctness Tests | ||
| 66 | + # ============================================================ | ||
| 67 | + | ||
| 68 | + | ||
| 69 | + def test_all_reduce_forward(self, device): | ||
| 70 | + """Test all_reduce does all_reduce in forward. | ||
| 71 | + | ||
| 72 | + Tensor is VARYING (different across ranks). | ||
| 73 | + Forward aggregates varying tensors via all_reduce(sum). | ||
| 74 | + """ | ||
| 75 | + group_name = dist.group.WORLD.group_name | ||
| 76 | + rank = dist.get_rank() | ||
| 77 | + | ||
| 78 | + # Each rank contributes its rank value (tensor is varying) | ||
| 79 | + input_tensor = torch.full((3, 3), fill_value=float(rank), device=device) | ||
| 80 | + output = fcols.all_reduce(input_tensor, "sum", group=group_name) | ||
| 81 | + | ||
| 82 | + # Forward does all_reduce: sum of 0+1+2+3 = 6 | ||
| 83 | + expected = torch.full( | ||
| 84 | + (3, 3), | ||
| 85 | + fill_value=self.world_size * (self.world_size - 1) / 2, | ||
| 86 | + device=device, | ||
| 87 | + ) | ||
| 88 | + self.assertEqual(output, expected) | ||
| 89 | + | ||
| 90 | + | ||
| 91 | + | ||
| 92 | + def test_all_gather_tensor_forward(self, device, gather_dim): | ||
| 93 | + """Test all_gather_tensor produces correct output shape. | ||
| 94 | + | ||
| 95 | + Tensor is VARYING (different across ranks). | ||
| 96 | + Forward gathers tensors from all ranks along gather_dim. | ||
| 97 | + """ | ||
| 98 | + group_name = dist.group.WORLD.group_name | ||
| 99 | + rank = dist.get_rank() | ||
| 100 | + | ||
| 101 | + # Each rank has tensor with its rank value | ||
| 102 | + input_tensor = torch.full((3, 3, 3), fill_value=float(rank), device=device) | ||
| 103 | + output = fcols.all_gather_tensor( | ||
| 104 | + input_tensor, gather_dim=gather_dim, group=group_name | ||
| 105 | + ) | ||
| 106 | + | ||
| 107 | + # Verify output shape | ||
| 108 | + expected_shape = list(input_tensor.shape) | ||
| 109 | + expected_shape[gather_dim] *= self.world_size | ||
| 110 | + self.assertEqual(list(output.shape), expected_shape) | ||
| 111 | + | ||
| 112 | + # Verify output contains all ranks' data | ||
| 113 | + # Check each chunk along gather_dim contains the correct rank value | ||
| 114 | + for r in range(self.world_size): | ||
| 115 | + chunk = output.narrow(gather_dim, r * 3, 3) | ||
| 116 | + expected_chunk = torch.full((3, 3, 3), fill_value=float(r), device=device) | ||
| 117 | + self.assertEqual(chunk, expected_chunk) | ||
| 118 | + | ||
| 119 | + | ||
| 120 | + | ||
| 121 | + def test_reduce_scatter_tensor_forward(self, device, scatter_dim): | ||
| 122 | + """Test reduce_scatter_tensor produces correct output shape. | ||
| 123 | + | ||
| 124 | + Tensor is VARYING (different across ranks). | ||
| 125 | + Forward reduces and scatters chunks to ranks. | ||
| 126 | + """ | ||
| 127 | + group_name = dist.group.WORLD.group_name | ||
| 128 | + rank = dist.get_rank() | ||
| 129 | + | ||
| 130 | + # Create input with appropriate size | ||
| 131 | + if scatter_dim == 0: | ||
| 132 | + input_tensor = torch.full( | ||
| 133 | + (4 * self.world_size, 3), fill_value=float(rank), device=device | ||
| 134 | + ) | ||
| 135 | + else: # scatter_dim == 1 | ||
| 136 | + input_tensor = torch.full( | ||
| 137 | + (3, 4 * self.world_size), fill_value=float(rank), device=device | ||
| 138 | + ) | ||
| 139 | + | ||
| 140 | + output = fcols.reduce_scatter_tensor( | ||
| 141 | + input_tensor, "sum", scatter_dim=scatter_dim, group=group_name | ||
| 142 | + ) | ||
| 143 | + | ||
| 144 | + # Verify output shape | ||
| 145 | + expected_shape = list(input_tensor.shape) | ||
| 146 | + expected_shape[scatter_dim] //= self.world_size | ||
| 147 | + self.assertEqual(list(output.shape), expected_shape) | ||
| 148 | + | ||
| 149 | + # Each rank should receive sum of all ranks' values: 0+1+2+3 = 6 | ||
| 150 | + expected_value = self.world_size * (self.world_size - 1) / 2 | ||
| 151 | + expected = torch.full_like(output, fill_value=expected_value) | ||
| 152 | + self.assertEqual(output, expected) | ||
| 153 | + | ||
| 154 | + | ||
| 155 | + def test_all_to_all_single_forward(self, device): | ||
| 156 | + """Test all_to_all_single with uniform splits. | ||
| 157 | + | ||
| 158 | + Tensor is VARYING (different across ranks). | ||
| 159 | + Forward exchanges tensor chunks between ranks. | ||
| 160 | + """ | ||
| 161 | + group_name = dist.group.WORLD.group_name | ||
| 162 | + rank = dist.get_rank() | ||
| 163 | + | ||
| 164 | + # Each rank contributes its rank value | ||
| 165 | + input_tensor = torch.full( | ||
| 166 | + (2 * self.world_size, 3), fill_value=float(rank), device=device | ||
| 167 | + ) | ||
| 168 | + | ||
| 169 | + # Uniform split | ||
| 170 | + output = fcols.all_to_all_single( | ||
| 171 | + input_tensor, | ||
| 172 | + output_split_sizes=None, | ||
| 173 | + input_split_sizes=None, | ||
| 174 | + group=group_name, | ||
| 175 | + ) | ||
| 176 | + | ||
| 177 | + # Output should have same shape as input for uniform splits | ||
| 178 | + self.assertEqual(output.shape, input_tensor.shape) | ||
| 179 | + | ||
| 180 | + # Verify each rank receives data from all other ranks | ||
| 181 | + for r in range(self.world_size): | ||
| 182 | + chunk = output[r * 2 : (r + 1) * 2, :] | ||
| 183 | + expected_chunk = torch.full((2, 3), fill_value=float(r), device=device) | ||
| 184 | + self.assertEqual(chunk, expected_chunk) | ||
| 185 | + | ||
| 186 | + | ||
| 187 | + def test_all_reduce_coalesced_forward(self, device): | ||
| 188 | + """Test all_reduce_coalesced does all_reduce on each tensor. | ||
| 189 | + | ||
| 190 | + Tensors are VARYING (different across ranks). | ||
| 191 | + Forward aggregates varying tensors via all_reduce(sum). | ||
| 192 | + """ | ||
| 193 | + group_name = dist.group.WORLD.group_name | ||
| 194 | + rank = dist.get_rank() | ||
| 195 | + | ||
| 196 | + # Each rank contributes its rank value | ||
| 197 | + input_tensors = [ | ||
| 198 | + torch.full((3, 3), fill_value=float(rank), device=device), | ||
| 199 | + torch.full((2, 2), fill_value=float(rank), device=device), | ||
| 200 | + ] | ||
| 201 | + outputs = fcols.all_reduce_coalesced(input_tensors, "sum", group=group_name) | ||
| 202 | + | ||
| 203 | + # Forward does all_reduce: sum of 0+1+2+3 = 6 | ||
| 204 | + expected_value = self.world_size * (self.world_size - 1) / 2 | ||
| 205 | + for output, input_tensor in zip(outputs, input_tensors): | ||
| 206 | + expected = torch.full_like(input_tensor, fill_value=expected_value) | ||
| 207 | + self.assertEqual(output, expected) | ||
| 208 | + | ||
| 209 | + | ||
| 210 | + def test_all_gather_into_tensor_coalesced_forward(self, device): | ||
| 211 | + """Test all_gather_into_tensor_coalesced gathers each tensor. | ||
| 212 | + | ||
| 213 | + Tensors are VARYING (different across ranks). | ||
| 214 | + Forward gathers tensors from all ranks. | ||
| 215 | + """ | ||
| 216 | + group_name = dist.group.WORLD.group_name | ||
| 217 | + rank = dist.get_rank() | ||
| 218 | + | ||
| 219 | + # Each rank has tensors with its rank value | ||
| 220 | + input_tensors = [ | ||
| 221 | + torch.full((3, 3), fill_value=float(rank), device=device), | ||
| 222 | + torch.full((2, 2), fill_value=float(rank), device=device), | ||
| 223 | + ] | ||
| 224 | + outputs = fcols.all_gather_into_tensor_coalesced( | ||
| 225 | + input_tensors, group=group_name | ||
| 226 | + ) | ||
| 227 | + | ||
| 228 | + # Verify output shapes | ||
| 229 | + for output, input_tensor in zip(outputs, input_tensors): | ||
| 230 | + expected_shape = list(input_tensor.shape) | ||
| 231 | + expected_shape[0] *= self.world_size | ||
| 232 | + self.assertEqual(list(output.shape), expected_shape) | ||
| 233 | + | ||
| 234 | + | ||
| 235 | + def test_reduce_scatter_tensor_coalesced_forward(self, device): | ||
| 236 | + """Test reduce_scatter_tensor_coalesced reduces and scatters each tensor. | ||
| 237 | + | ||
| 238 | + Tensors are VARYING (different across ranks). | ||
| 239 | + Forward reduces and scatters chunks to ranks. | ||
| 240 | + """ | ||
| 241 | + group_name = dist.group.WORLD.group_name | ||
| 242 | + rank = dist.get_rank() | ||
| 243 | + | ||
| 244 | + # Create inputs with appropriate size (divisible by world_size) | ||
| 245 | + input_tensors = [ | ||
| 246 | + torch.full((4 * self.world_size, 3), fill_value=float(rank), device=device), | ||
| 247 | + torch.full((2 * self.world_size, 2), fill_value=float(rank), device=device), | ||
| 248 | + ] | ||
| 249 | + scatter_dims = [0, 0] | ||
| 250 | + | ||
| 251 | + outputs = fcols.reduce_scatter_tensor_coalesced( | ||
| 252 | + input_tensors, "sum", scatter_dims, group=group_name | ||
| 253 | + ) | ||
| 254 | + | ||
| 255 | + # Each rank should receive sum of all ranks' values: 0+1+2+3 = 6 | ||
| 256 | + expected_value = self.world_size * (self.world_size - 1) / 2 | ||
| 257 | + for output, input_tensor in zip(outputs, input_tensors): | ||
| 258 | + expected_shape = list(input_tensor.shape) | ||
| 259 | + expected_shape[0] //= self.world_size | ||
| 260 | + self.assertEqual(list(output.shape), expected_shape) | ||
| 261 | + expected = torch.full_like(output, fill_value=expected_value) | ||
| 262 | + self.assertEqual(output, expected) | ||
| 263 | + | ||
| 264 | + # ============================================================ | ||
| 265 | + # Backward Correctness Tests | ||
| 266 | + # ============================================================ | ||
| 267 | + | ||
| 268 | + | ||
| 269 | + def test_all_reduce_backward(self, device): | ||
| 270 | + """Test all_reduce backward does all_reduce. | ||
| 271 | + | ||
| 272 | + Both tensor AND gradients are VARYING (different across ranks). | ||
| 273 | + Backward aggregates gradients via all_reduce(sum). | ||
| 274 | + """ | ||
| 275 | + group_name = dist.group.WORLD.group_name | ||
| 276 | + | ||
| 277 | + input_tensor = torch.randn(3, 3, requires_grad=True, device=device) | ||
| 278 | + output = fcols.all_reduce(input_tensor, "sum", group=group_name) | ||
| 279 | + | ||
| 280 | + # Backward with ones | ||
| 281 | + output.sum().backward() | ||
| 282 | + | ||
| 283 | + # Gradient should be aggregated (backward is all_reduce) | ||
| 284 | + expected_grad = torch.full( | ||
| 285 | + (3, 3), fill_value=float(self.world_size), device=device | ||
| 286 | + ) | ||
| 287 | + self.assertEqual(input_tensor.grad, expected_grad) | ||
| 288 | + | ||
| 289 | + # Backward is all_reduce (sum) | ||
| 290 | + grad_outputs = torch.rand_like(output, device=device) | ||
| 291 | + (grad_input,) = torch.autograd.grad( | ||
| 292 | + output, input_tensor, grad_outputs=grad_outputs | ||
| 293 | + ) | ||
| 294 | + expected_grad_input = fcols.all_reduce(grad_outputs, "sum", group=group_name) | ||
| 295 | + self.assertEqual(grad_input, expected_grad_input) | ||
| 296 | + | ||
| 297 | + | ||
| 298 | + | ||
| 299 | + def test_all_gather_tensor_backward(self, device, gather_dim): | ||
| 300 | + """Test all_gather_tensor backward does reduce_scatter. | ||
| 301 | + | ||
| 302 | + Both tensor AND gradients are VARYING (different across ranks). | ||
| 303 | + Forward gathers tensors, backward reduces and scatters gradients. | ||
| 304 | + """ | ||
| 305 | + group_name = dist.group.WORLD.group_name | ||
| 306 | + | ||
| 307 | + input_tensor = torch.randn(3, 3, 3, requires_grad=True, device=device) | ||
| 308 | + output = fcols.all_gather_tensor( | ||
| 309 | + input_tensor, gather_dim=gather_dim, group=group_name | ||
| 310 | + ) | ||
| 311 | + | ||
| 312 | + # Backward with ones | ||
| 313 | + output.sum().backward() | ||
| 314 | + | ||
| 315 | + # Gradient should be reduce_scatter of ones | ||
| 316 | + self.assertIsNotNone(input_tensor.grad) | ||
| 317 | + # Gradient should be all world_size (sum from all ranks) | ||
| 318 | + expected_grad = torch.full( | ||
| 319 | + (3, 3, 3), fill_value=float(self.world_size), device=device | ||
| 320 | + ) | ||
| 321 | + self.assertEqual(input_tensor.grad, expected_grad) | ||
| 322 | + | ||
| 323 | + # Backward is reduce_scatter (sum) | ||
| 324 | + grad_outputs = torch.rand_like(output, device=device) | ||
| 325 | + (grad_input,) = torch.autograd.grad( | ||
| 326 | + output, input_tensor, grad_outputs=grad_outputs | ||
| 327 | + ) | ||
| 328 | + expected_grad_input = fcols.reduce_scatter_tensor( | ||
| 329 | + grad_outputs, "sum", gather_dim, group=group_name | ||
| 330 | + ) | ||
| 331 | + self.assertEqual(grad_input, expected_grad_input) | ||
| 332 | + | ||
| 333 | + | ||
| 334 | + | ||
| 335 | + def test_reduce_scatter_tensor_backward(self, device, scatter_dim): | ||
| 336 | + """Test reduce_scatter_tensor backward does all_gather. | ||
| 337 | + | ||
| 338 | + Both tensor AND gradients are VARYING (different across ranks). | ||
| 339 | + Forward reduces and scatters, backward gathers gradients. | ||
| 340 | + """ | ||
| 341 | + group_name = dist.group.WORLD.group_name | ||
| 342 | + | ||
| 343 | + # Create input with appropriate size | ||
| 344 | + if scatter_dim == 0: | ||
| 345 | + input_tensor = torch.randn( | ||
| 346 | + 4 * self.world_size, 3, requires_grad=True, device=device | ||
| 347 | + ) | ||
| 348 | + else: | ||
| 349 | + input_tensor = torch.randn( | ||
| 350 | + 3, 4 * self.world_size, requires_grad=True, device=device | ||
| 351 | + ) | ||
| 352 | + | ||
| 353 | + output = fcols.reduce_scatter_tensor( | ||
| 354 | + input_tensor, "sum", scatter_dim=scatter_dim, group=group_name | ||
| 355 | + ) | ||
| 356 | + | ||
| 357 | + # Backward with ones | ||
| 358 | + output.sum().backward() | ||
| 359 | + | ||
| 360 | + # Gradient should be all_gather of ones | ||
| 361 | + self.assertIsNotNone(input_tensor.grad) | ||
| 362 | + self.assertEqual(input_tensor.grad.shape, input_tensor.shape) | ||
| 363 | + | ||
| 364 | + # All gradients should be 1 (gathered from all ranks) | ||
| 365 | + expected_grad = torch.ones_like(input_tensor) | ||
| 366 | + self.assertEqual(input_tensor.grad, expected_grad) | ||
| 367 | + | ||
| 368 | + # Backward is all_gather (sum) | ||
| 369 | + grad_outputs = torch.rand_like(output, device=device) | ||
| 370 | + (grad_input,) = torch.autograd.grad( | ||
| 371 | + output, input_tensor, grad_outputs=grad_outputs | ||
| 372 | + ) | ||
| 373 | + expected_grad_input = fcols.all_gather_tensor( | ||
| 374 | + grad_outputs, scatter_dim, group=group_name | ||
| 375 | + ) | ||
| 376 | + self.assertEqual(grad_input, expected_grad_input) | ||
| 377 | + | ||
| 378 | + | ||
| 379 | + def test_all_to_all_single_backward(self, device): | ||
| 380 | + """Test all_to_all_single backward reverses split sizes. | ||
| 381 | + | ||
| 382 | + Both tensor AND gradients are VARYING (different across ranks). | ||
| 383 | + Forward does all_to_all, backward does all_to_all with reversed splits. | ||
| 384 | + """ | ||
| 385 | + group_name = dist.group.WORLD.group_name | ||
| 386 | + | ||
| 387 | + input_tensor = torch.randn( | ||
| 388 | + 4 * self.world_size, 3, requires_grad=True, device=device | ||
| 389 | + ) | ||
| 390 | + output = fcols.all_to_all_single( | ||
| 391 | + input_tensor, | ||
| 392 | + output_split_sizes=None, | ||
| 393 | + input_split_sizes=None, | ||
| 394 | + group=group_name, | ||
| 395 | + ) | ||
| 396 | + | ||
| 397 | + # Backward | ||
| 398 | + output.sum().backward() | ||
| 399 | + | ||
| 400 | + # Gradient should have same shape as input | ||
| 401 | + self.assertIsNotNone(input_tensor.grad) | ||
| 402 | + self.assertEqual(input_tensor.grad.shape, input_tensor.shape) | ||
| 403 | + expected_grad = torch.ones_like(input_tensor) | ||
| 404 | + self.assertEqual(input_tensor.grad, expected_grad) | ||
| 405 | + | ||
| 406 | + # Backward is all_gather (sum) | ||
| 407 | + grad_outputs = torch.rand_like(output, device=device) | ||
| 408 | + (grad_input,) = torch.autograd.grad( | ||
| 409 | + output, input_tensor, grad_outputs=grad_outputs | ||
| 410 | + ) | ||
| 411 | + expected_grad_input = fcols.all_to_all_single( | ||
| 412 | + grad_outputs, None, None, group=group_name | ||
| 413 | + ) | ||
| 414 | + self.assertEqual(grad_input, expected_grad_input) | ||
| 415 | + | ||
| 416 | + | ||
| 417 | + def test_all_reduce_coalesced_backward(self, device): | ||
| 418 | + """Test all_reduce_coalesced backward does all_reduce on each gradient. | ||
| 419 | + | ||
| 420 | + Tensors AND gradients are VARYING (different across ranks). | ||
| 421 | + Backward aggregates each gradient via all_reduce(sum). | ||
| 422 | + """ | ||
| 423 | + group_name = dist.group.WORLD.group_name | ||
| 424 | + | ||
| 425 | + input_tensors = [ | ||
| 426 | + torch.randn(3, 3, requires_grad=True, device=device), | ||
| 427 | + torch.randn(2, 2, requires_grad=True, device=device), | ||
| 428 | + ] | ||
| 429 | + outputs = fcols.all_reduce_coalesced(input_tensors, "sum", group=group_name) | ||
| 430 | + | ||
| 431 | + # Backward with ones | ||
| 432 | + loss = sum(output.sum() for output in outputs) | ||
| 433 | + loss.backward() | ||
| 434 | + | ||
| 435 | + # Each gradient should be aggregated (backward is all_reduce) | ||
| 436 | + for input_tensor in input_tensors: | ||
| 437 | + self.assertIsNotNone(input_tensor.grad) | ||
| 438 | + expected_grad = torch.full_like( | ||
| 439 | + input_tensor, fill_value=float(self.world_size) | ||
| 440 | + ) | ||
| 441 | + self.assertEqual(input_tensor.grad, expected_grad) | ||
| 442 | + | ||
| 443 | + | ||
| 444 | + def test_all_gather_into_tensor_coalesced_backward(self, device): | ||
| 445 | + """Test all_gather_into_tensor_coalesced backward does reduce_scatter on each gradient. | ||
| 446 | + | ||
| 447 | + Tensors AND gradients are VARYING (different across ranks). | ||
| 448 | + Forward gathers each tensor, backward reduce_scatters each gradient. | ||
| 449 | + """ | ||
| 450 | + group_name = dist.group.WORLD.group_name | ||
| 451 | + | ||
| 452 | + input_tensors = [ | ||
| 453 | + torch.randn(3, 3, requires_grad=True, device=device), | ||
| 454 | + torch.randn(2, 2, requires_grad=True, device=device), | ||
| 455 | + ] | ||
| 456 | + outputs = fcols.all_gather_into_tensor_coalesced( | ||
| 457 | + input_tensors, group=group_name | ||
| 458 | + ) | ||
| 459 | + | ||
| 460 | + # Backward with ones | ||
| 461 | + loss = sum(output.sum() for output in outputs) | ||
| 462 | + loss.backward() | ||
| 463 | + | ||
| 464 | + # Each gradient should be reduce_scatter of ones | ||
| 465 | + for input_tensor in input_tensors: | ||
| 466 | + self.assertIsNotNone(input_tensor.grad) | ||
| 467 | + expected_grad = torch.full_like( | ||
| 468 | + input_tensor, fill_value=float(self.world_size) | ||
| 469 | + ) | ||
| 470 | + self.assertEqual(input_tensor.grad, expected_grad) | ||
| 471 | + | ||
| 472 | + | ||
| 473 | + def test_reduce_scatter_tensor_coalesced_backward(self, device): | ||
| 474 | + """Test reduce_scatter_tensor_coalesced backward does all_gather on each gradient. | ||
| 475 | + | ||
| 476 | + Tensors AND gradients are VARYING (different across ranks). | ||
| 477 | + Forward reduces and scatters each tensor, backward gathers each gradient. | ||
| 478 | + """ | ||
| 479 | + group_name = dist.group.WORLD.group_name | ||
| 480 | + | ||
| 481 | + input_tensors = [ | ||
| 482 | + torch.randn(4 * self.world_size, 3, requires_grad=True, device=device), | ||
| 483 | + torch.randn(2 * self.world_size, 2, requires_grad=True, device=device), | ||
| 484 | + ] | ||
| 485 | + scatter_dims = [0, 0] | ||
| 486 | + | ||
| 487 | + outputs = fcols.reduce_scatter_tensor_coalesced( | ||
| 488 | + input_tensors, "sum", scatter_dims, group=group_name | ||
| 489 | + ) | ||
| 490 | + | ||
| 491 | + # Backward with ones | ||
| 492 | + loss = sum(output.sum() for output in outputs) | ||
| 493 | + loss.backward() | ||
| 494 | + | ||
| 495 | + # Each gradient should be all_gather of ones | ||
| 496 | + for input_tensor in input_tensors: | ||
| 497 | + self.assertIsNotNone(input_tensor.grad) | ||
| 498 | + expected_grad = torch.ones_like(input_tensor) | ||
| 499 | + self.assertEqual(input_tensor.grad, expected_grad) | ||
| 500 | + | ||
| 501 | + # ============================================================ | ||
| 502 | + # torch.library.opcheck Tests | ||
| 503 | + # ============================================================ | ||
| 504 | + | ||
| 505 | + test_utils = [ | ||
| 506 | + "test_schema", | ||
| 507 | + "test_autograd_registration", | ||
| 508 | + "test_faketensor", | ||
| 509 | + # "test_aot_dispatch_dynamic" - Open issue with check: TBD | ||
| 510 | + ] | ||
| 511 | + | ||
| 512 | + | ||
| 513 | + def test_all_reduce_opcheck(self, test_utils): | ||
| 514 | + """Test all_reduce op registration with torch.library.opcheck. | ||
| 515 | + | ||
| 516 | + Verifies all aspects of op registration including: | ||
| 517 | + - Fake tensor support | ||
| 518 | + - Autograd support (backward does all_reduce) | ||
| 519 | + - Schema validation | ||
| 520 | + """ | ||
| 521 | + group_name = dist.group.WORLD.group_name | ||
| 522 | + | ||
| 523 | + input_tensor = torch.ones(3, 3, requires_grad=True) | ||
| 524 | + | ||
| 525 | + # opcheck verifies all aspects of op registration | ||
| 526 | + torch.library.opcheck( | ||
| 527 | + torch.ops._c10d_functional.all_reduce, | ||
| 528 | + (input_tensor, "sum", group_name), | ||
| 529 | + test_utils=test_utils, | ||
| 530 | + ) | ||
| 531 | + | ||
| 532 | + | ||
| 533 | + def test_all_gather_into_tensor_opcheck(self, test_utils): | ||
| 534 | + """Test all_gather_into_tensor op registration with torch.library.opcheck. | ||
| 535 | + | ||
| 536 | + Verifies all aspects of op registration including: | ||
| 537 | + - Fake tensor support | ||
| 538 | + - Autograd support (backward does reduce_scatter) | ||
| 539 | + - Schema validation | ||
| 540 | + """ | ||
| 541 | + group_name = dist.group.WORLD.group_name | ||
| 542 | + | ||
| 543 | + input_tensor = torch.ones(3, 3, 3, requires_grad=True) | ||
| 544 | + | ||
| 545 | + # opcheck verifies all aspects of op registration | ||
| 546 | + torch.library.opcheck( | ||
| 547 | + torch.ops._c10d_functional.all_gather_into_tensor, | ||
| 548 | + (input_tensor, self.world_size, group_name), | ||
| 549 | + test_utils=test_utils, | ||
| 550 | + ) | ||
| 551 | + | ||
| 552 | + | ||
| 553 | + def test_reduce_scatter_tensor_opcheck(self, test_utils): | ||
| 554 | + """Test reduce_scatter_tensor op registration with torch.library.opcheck. | ||
| 555 | + | ||
| 556 | + Verifies all aspects of op registration including: | ||
| 557 | + - Fake tensor support | ||
| 558 | + - Autograd support (backward does all_gather) | ||
| 559 | + - Schema validation | ||
| 560 | + """ | ||
| 561 | + group_name = dist.group.WORLD.group_name | ||
| 562 | + | ||
| 563 | + # Input should be divisible by world_size | ||
| 564 | + input_tensor = torch.ones(4 * self.world_size, 3, requires_grad=True) | ||
| 565 | + | ||
| 566 | + # opcheck verifies all aspects of op registration | ||
| 567 | + torch.library.opcheck( | ||
| 568 | + torch.ops._c10d_functional.reduce_scatter_tensor, | ||
| 569 | + (input_tensor, "sum", self.world_size, group_name), | ||
| 570 | + test_utils=test_utils, | ||
| 571 | + ) | ||
| 572 | + | ||
| 573 | + | ||
| 574 | + def test_all_to_all_single_opcheck(self, test_utils): | ||
| 575 | + """Test all_to_all_single op registration with torch.library.opcheck. | ||
| 576 | + | ||
| 577 | + Verifies all aspects of op registration including: | ||
| 578 | + - Fake tensor support | ||
| 579 | + - Autograd support (backward reverses split sizes) | ||
| 580 | + - Schema validation | ||
| 581 | + """ | ||
| 582 | + group_name = dist.group.WORLD.group_name | ||
| 583 | + group_size = dist.group.WORLD.size() | ||
| 584 | + | ||
| 585 | + # Input should be divisible by world_size | ||
| 586 | + input_tensor = torch.ones(4 * self.world_size, 3, requires_grad=True) | ||
| 587 | + | ||
| 588 | + output_split_sizes = [input_tensor.shape[0] // group_size] * group_size | ||
| 589 | + input_split_sizes = output_split_sizes | ||
| 590 | + | ||
| 591 | + # opcheck verifies all aspects of op registration | ||
| 592 | + torch.library.opcheck( | ||
| 593 | + torch.ops._c10d_functional.all_to_all_single, | ||
| 594 | + (input_tensor, output_split_sizes, input_split_sizes, group_name), | ||
| 595 | + test_utils=test_utils, | ||
| 596 | + ) | ||
| 597 | + | ||
| 598 | + | ||
| 599 | +if __name__ == "__main__": | ||
| 600 | + run_tests() | ||
| @@ -0,0 +1,178 @@ | |||
| 1 | +# Owner(s): ["oncall: distributed"] | ||
| 2 | + | ||
| 3 | +import os | ||
| 4 | +import sys | ||
| 5 | +from datetime import timedelta | ||
| 6 | + | ||
| 7 | +import torch | ||
| 8 | +import torch_npu | ||
| 9 | +import torch.distributed as c10d | ||
| 10 | +from torch._C._distributed_c10d import _ProcessGroupWrapper | ||
| 11 | +from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU | ||
| 12 | + | ||
| 13 | + | ||
| 14 | +if not c10d.is_available(): | ||
| 15 | + print("c10d not available, skipping tests", file=sys.stderr) | ||
| 16 | + sys.exit(0) | ||
| 17 | + | ||
| 18 | +from torch.testing._internal.common_distributed import ( | ||
| 19 | + MultiProcessTestCase, | ||
| 20 | + with_dist_debug_levels, | ||
| 21 | +) | ||
| 22 | +from torch.testing._internal.common_utils import ( | ||
| 23 | + run_tests, | ||
| 24 | + TEST_WITH_DEV_DBG_ASAN, | ||
| 25 | +) | ||
| 26 | + | ||
| 27 | + | ||
| 28 | +device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu" | ||
| 29 | +backend = c10d.get_default_backend_for_device(device_type) | ||
| 30 | + | ||
| 31 | + | ||
| 32 | +class AbstractProcessGroupWrapperTest(MultiProcessTestCase): | ||
| 33 | + def setUp(self): | ||
| 34 | + super().setUp() | ||
| 35 | + self._spawn_processes() | ||
| 36 | + | ||
| 37 | + def _validate_error(self, exception, op_type, rank, tensor, verify_diff=True): | ||
| 38 | + err = str(exception) | ||
| 39 | + self.assertTrue( | ||
| 40 | + op_type in err, f"Got {err} but expected {op_type} to be in error." | ||
| 41 | + ) | ||
| 42 | + # User doesn't call barrier with tensor. | ||
| 43 | + if op_type != "BARRIER": | ||
| 44 | + self.assertTrue( | ||
| 45 | + f"{list(tensor.shape)}" in err, | ||
| 46 | + f"Did not find shapes {list(tensor.shape)} in error {err}", | ||
| 47 | + ) | ||
| 48 | + # For NPU, only assert on device type, not index | ||
| 49 | + if device_type in str(tensor.device): | ||
| 50 | + self.assertTrue( | ||
| 51 | + device_type in err, | ||
| 52 | + f"Did not find {device_type} device in error {err}", | ||
| 53 | + ) | ||
| 54 | + else: | ||
| 55 | + self.assertTrue( | ||
| 56 | + str(tensor.device) in err, | ||
| 57 | + f"Did not find tensor device {str(tensor.device)} in error {err}", | ||
| 58 | + ) | ||
| 59 | + # C++ and python type strings are not exactly the same. | ||
| 60 | + if "float" in str(tensor.dtype): | ||
| 61 | + self.assertTrue("Float" in err, "Expected Float type") | ||
| 62 | + elif "int" in str(tensor.dtype): | ||
| 63 | + self.assertTrue("Long" in err, "Expected Long type") | ||
| 64 | + else: | ||
| 65 | + self.fail(f"Unexpected dtype {str(tensor.dtype)} for error {err}") | ||
| 66 | + | ||
| 67 | + # Ensure sequence number is logged in error | ||
| 68 | + self.assertTrue("SequenceNumber" in err) | ||
| 69 | + # Ensure info about how collectives diff is in the error. | ||
| 70 | + if verify_diff: | ||
| 71 | + self.assertTrue( | ||
| 72 | + "Collectives differ in the following" in err, f"Got error {err}" | ||
| 73 | + ) | ||
| 74 | + | ||
| 75 | + | ||
| 76 | +# ASAN is not safe since we are spawning processes. | ||
| 77 | +if not TEST_WITH_DEV_DBG_ASAN: | ||
| 78 | + class ProcessGroupHCCLWrapperTest(AbstractProcessGroupWrapperTest): | ||
| 79 | + def setUp(self): | ||
| 80 | + super(AbstractProcessGroupWrapperTest, self).setUp() | ||
| 81 | + self._spawn_processes() | ||
| 82 | + # TORCH_HCCL_BLOCKING_WAIT overrides TORCH_HCCL_ASYNC_ERROR_HANDLING hence tests | ||
| 83 | + # that use TORCH_HCCL_BLOCKING_WAIT will test it as expected. | ||
| 84 | + os.environ["TORCH_HCCL_ASYNC_ERROR_HANDLING"] = "1" | ||
| 85 | + | ||
| 86 | + | ||
| 87 | + def world_size(self) -> int: | ||
| 88 | + return 2 | ||
| 89 | + | ||
| 90 | + def _create_wrapper_pg(self, with_new_group=False, timeout=10.0): | ||
| 91 | + store = c10d.FileStore(self.file_name, self.world_size) | ||
| 92 | + c10d.init_process_group( | ||
| 93 | + backend=backend, | ||
| 94 | + rank=self.rank, | ||
| 95 | + world_size=self.world_size, | ||
| 96 | + store=store, | ||
| 97 | + timeout=timedelta(seconds=timeout), | ||
| 98 | + ) | ||
| 99 | + if with_new_group: | ||
| 100 | + pg = c10d.new_group(backend=backend, timeout=timedelta(seconds=timeout)) | ||
| 101 | + else: | ||
| 102 | + if device_type == "xpu": | ||
| 103 | + _pg = c10d.ProcessGroupXCCL( | ||
| 104 | + store, | ||
| 105 | + self.rank, | ||
| 106 | + self.world_size, | ||
| 107 | + ) | ||
| 108 | + else: | ||
| 109 | + _pg = torch_npu._C._distributed_c10d.ProcessGroupHCCL( | ||
| 110 | + store, | ||
| 111 | + self.rank, | ||
| 112 | + self.world_size, | ||
| 113 | + timeout=timedelta(seconds=timeout), | ||
| 114 | + ) | ||
| 115 | + pg = c10d._create_process_group_wrapper( | ||
| 116 | + _pg, | ||
| 117 | + "unused", | ||
| 118 | + store, | ||
| 119 | + self.rank, | ||
| 120 | + self.world_size, | ||
| 121 | + timeout=timeout, | ||
| 122 | + ) | ||
| 123 | + return pg | ||
| 124 | + | ||
| 125 | + | ||
| 126 | + | ||
| 127 | + def test_wrapper_forwards_hccl_methods(self): | ||
| 128 | + """ | ||
| 129 | + Tests that ProcessGroupWrapper correctly forwards HCCL-specific | ||
| 130 | + utility methods to the wrapped backend. See issue #173538. | ||
| 131 | + """ | ||
| 132 | + torch.npu.set_device(self.rank) | ||
| 133 | + device = torch.device(f"npu:{self.rank}") | ||
| 134 | + wrapper = self._create_wrapper_pg(with_new_group=False) | ||
| 135 | + | ||
| 136 | + # Verify we're testing the wrapper | ||
| 137 | + self.assertIsInstance(wrapper, _ProcessGroupWrapper) | ||
| 138 | + unwrapped = wrapper.wrapped_pg | ||
| 139 | + | ||
| 140 | + # Verify wrapper forwards property/method calls to wrapped backend | ||
| 141 | + self.assertEqual(wrapper.supports_splitting, unwrapped.supports_splitting) | ||
| 142 | + self.assertEqual(wrapper.supports_coalescing, unwrapped.supports_coalescing) | ||
| 143 | + self.assertEqual( | ||
| 144 | + wrapper.supports_time_estimate, unwrapped.supports_time_estimate | ||
| 145 | + ) | ||
| 146 | + self.assertEqual( | ||
| 147 | + wrapper.supports_tensor_alloc(device), | ||
| 148 | + unwrapped.supports_tensor_alloc(device), | ||
| 149 | + ) | ||
| 150 | + if hasattr(unwrapped, 'get_error'): | ||
| 151 | + self.assertEqual(wrapper.get_error(), unwrapped.get_error()) | ||
| 152 | + try: | ||
| 153 | + wrapper_options = wrapper.options | ||
| 154 | + except RuntimeError as e: | ||
| 155 | + self.assertIn("does not implement getBackendOptions", str(e)) | ||
| 156 | + wrapper_options = None | ||
| 157 | + | ||
| 158 | + if wrapper_options is not None: | ||
| 159 | + self.assertIs(wrapper_options, unwrapped.options) | ||
| 160 | + | ||
| 161 | + # Test eager_connect_single_device forwarding (should not raise) | ||
| 162 | + wrapper.eager_connect_single_device(device) | ||
| 163 | + | ||
| 164 | + # HCCL does not support mem_allocator and allocate_tensor | ||
| 165 | + if wrapper.supports_tensor_alloc(device): | ||
| 166 | + self.assertIs(wrapper.mem_allocator, unwrapped.mem_allocator) | ||
| 167 | + tensor = wrapper.allocate_tensor( | ||
| 168 | + 1024, dtype=torch.float32, device=device | ||
| 169 | + ) | ||
| 170 | + self.assertEqual(tensor.shape, torch.Size([1024])) | ||
| 171 | + | ||
| 172 | + | ||
| 173 | +if __name__ == "__main__": | ||
| 174 | + assert not (torch.npu.is_initialized() or torch.xpu.is_initialized()), ( | ||
| 175 | + "test_pg_wrapper must not have initialized NPU/XPU context on main process" | ||
| 176 | + ) | ||
| 177 | + | ||
| 178 | + run_tests() | ||
🟠 High Priority
代码库中所有其他调用均为 current_accelerator() 无参数形式（见 test_pg_wrapper.py:28、test_accelerator.py:22 等）。PyTorch 标准 API 签名为 current_accelerator(device=None)，传入布尔值 True 不是有效的 device 参数。该行在模块顶层执行，NPU 可用时将直接触发 TypeError，导致整个测试文件无法导入运行。
同一段代码其他问题