已合并
test: Since the PyTorch community's test cases do not provide full coverage, we adopt the method of adding new test cases to supplement and cover all test case scenarios. #42824
创建于 27 天前
test: Since the PyTorch community's test cases do not provide full coverage, we adopt the method of adding new test cases to supplement and cover all test case scenarios. #42824
已合并
从已删除 :test_allreduce_hook_v2.7.1合入到Ascend/pytorchv2.7.1
共 1 个文件变更+362-0
| @@ -0,0 +1,362 @@ | |||
| 1 | +# Copyright (c) 2026 Huawei Technologies Co., Ltd | ||
| 2 | +# All rights reserved. | ||
| 3 | +# | ||
| 4 | +# Licensed under the BSD 3-Clause License (the "License"); | ||
| 5 | +# you may not use this file except in compliance with the License. | ||
| 6 | +# You may obtain a copy of the License at | ||
| 7 | +# | ||
| 8 | +# https://opensource.org/licenses/BSD-3-Clause | ||
| 9 | +# | ||
| 10 | +# Unless required by applicable law or agreed to in writing, software | ||
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, | ||
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 13 | +# See the License for the specific language governing permissions and | ||
| 14 | +# limitations under the License. | ||
| 15 | + | ||
| 16 | +""" | ||
| 17 | +Add validation cases for | ||
| 18 | +torch.distributed.algorithms.ddp_comm_hooks.default_hooks APIs on NPU: | ||
| 19 | +1. PyTorch community lacks sufficient direct validations for some default DDP | ||
| 20 | + communication hooks. | ||
| 21 | +2. This file validates | ||
| 22 | + torch.distributed.algorithms.ddp_comm_hooks.default_hooks.allreduce_hook | ||
| 23 | + (extendable). | ||
| 24 | +""" | ||
| 25 | + | ||
| 26 | +import copy | ||
| 27 | +import os | ||
| 28 | +from datetime import timedelta | ||
| 29 | + | ||
| 30 | +import torch | ||
| 31 | +import torch.distributed as dist | ||
| 32 | +import torch.multiprocessing as mp | ||
| 33 | +from torch import nn | ||
| 34 | +from torch.distributed.algorithms.ddp_comm_hooks import ( | ||
| 35 | + DDPCommHookType, | ||
| 36 | + default_hooks, | ||
| 37 | + register_ddp_comm_hook, | ||
| 38 | +) | ||
| 39 | +from torch.testing._internal.common_utils import TestCase, find_free_port, run_tests | ||
| 40 | +from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU | ||
| 41 | + | ||
| 42 | + | ||
| 43 | +device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu" | ||
| 44 | + | ||
| 45 | +WORLD_SIZE = 2 | ||
| 46 | +PROCESS_GROUP_TIMEOUT = timedelta(minutes=2) | ||
| 47 | + | ||
| 48 | + | ||
| 49 | +class _AllreduceModel(nn.Module): | ||
| 50 | + def __init__(self, shape=(40, 20), dtype=torch.float32): | ||
| 51 | + super().__init__() | ||
| 52 | + self.weight = nn.Parameter(torch.ones(shape, dtype=dtype).to(device_type)) | ||
| 53 | + | ||
| 54 | + def forward(self, input_tensor): | ||
| 55 | + return self.weight * input_tensor | ||
| 56 | + | ||
| 57 | + | ||
| 58 | +class _MultiBucketModel(nn.Module): | ||
| 59 | + def __init__(self): | ||
| 60 | + super().__init__() | ||
| 61 | + self.layers = nn.ModuleList([nn.Linear(32, 32, bias=False) for _ in range(4)]) | ||
| 62 | + | ||
| 63 | + def forward(self, input_tensor): | ||
| 64 | + return sum(layer(input_tensor) for layer in self.layers) | ||
| 65 | + | ||
| 66 | + | ||
| 67 | +class _Bucket: | ||
| 68 | + def __init__(self, tensor): | ||
| 69 | + self.tensor = tensor | ||
| 70 | + | ||
| 71 | + def index(self): | ||
| 72 | + return 0 | ||
| 73 | + | ||
| 74 | + def buffer(self): | ||
| 75 | + return self.tensor | ||
| 76 | + | ||
| 77 | + def gradients(self): | ||
| 78 | + return [self.tensor] | ||
| 79 | + | ||
| 80 | + def parameters(self): | ||
| 81 | + return [] | ||
| 82 | + | ||
| 83 | + def is_last(self): | ||
| 84 | + return True | ||
| 85 | + | ||
| 86 | + def set_buffer(self, tensor): | ||
| 87 | + self.tensor = tensor | ||
| 88 | + | ||
| 89 | + | ||
| 90 | +def _future_contract_hook( | ||
| 91 | + state, bucket: dist.GradBucket | ||
| 92 | +) -> torch.futures.Future[torch.Tensor]: | ||
| 93 | + buffer = bucket.buffer() | ||
| 94 | + future = default_hooks.allreduce_hook(state["process_group"], bucket) | ||
| 95 | + state["is_future"] = isinstance(future, torch._C.Future) | ||
| 96 | + state["buffer_shape"] = buffer.shape | ||
| 97 | + | ||
| 98 | + def validate(fut): | ||
| 99 | + result = fut.value() | ||
| 100 | + state["result_shape"] = result.shape | ||
| 101 | + state["result_dtype"] = result.dtype | ||
| 102 | + state["result_device_type"] = result.device.type | ||
| 103 | + return result | ||
| 104 | + | ||
| 105 | + return future.then(validate) | ||
| 106 | + | ||
| 107 | + | ||
| 108 | +def _counting_hook( | ||
| 109 | + state, bucket: dist.GradBucket | ||
| 110 | +) -> torch.futures.Future[torch.Tensor]: | ||
| 111 | + state["calls"] += 1 | ||
| 112 | + return default_hooks.allreduce_hook(state["process_group"], bucket) | ||
| 113 | + | ||
| 114 | + | ||
| 115 | +class TestAllreduceHook(TestCase): | ||
| 116 | + | ||
| 117 | + def _init_process_group(rank, world_size, port): | ||
| 118 | + os.environ["MASTER_ADDR"] = "127.0.0.1" | ||
| 119 | + os.environ["MASTER_PORT"] = str(port) | ||
| 120 | + torch.accelerator.set_device_index(rank) | ||
| 121 | + dist.init_process_group( | ||
| 122 | + "hccl", | ||
| 123 | + rank=rank, | ||
| 124 | + world_size=world_size, | ||
| 125 | + timeout=PROCESS_GROUP_TIMEOUT, | ||
| 126 | + ) | ||
| 127 | + return dist.group.WORLD | ||
| 128 | + | ||
| 129 | + | ||
| 130 | + def _ddp( | ||
| 131 | + model, | ||
| 132 | + rank, | ||
| 133 | + process_group, | ||
| 134 | + gradient_as_bucket_view=False, | ||
| 135 | + static_graph=False, | ||
| 136 | + bucket_cap_mb=None, | ||
| 137 | + ): | ||
| 138 | + kwargs = { | ||
| 139 | + "device_ids": [rank], | ||
| 140 | + "process_group": process_group, | ||
| 141 | + "gradient_as_bucket_view": gradient_as_bucket_view, | ||
| 142 | + "static_graph": static_graph, | ||
| 143 | + } | ||
| 144 | + if bucket_cap_mb is not None: | ||
| 145 | + kwargs["bucket_cap_mb"] = bucket_cap_mb | ||
| 146 | + return nn.parallel.DistributedDataParallel(model, **kwargs) | ||
| 147 | + | ||
| 148 | + | ||
| 149 | + def _gradient(model, input_tensor, use_mean=True): | ||
| 150 | + output = model(input_tensor) | ||
| 151 | + loss = output.mean() if use_mean else output.sum() | ||
| 152 | + loss.backward() | ||
| 153 | + return [parameter.grad.detach().clone() for parameter in model.parameters()] | ||
| 154 | + | ||
| 155 | + | ||
| 156 | + def _run_ddp_parity( | ||
| 157 | + cls, | ||
| 158 | + rank, | ||
| 159 | + world_size, | ||
| 160 | + port, | ||
| 161 | + gradient_as_bucket_view=False, | ||
| 162 | + static_graph=False, | ||
| 163 | + use_none_process_group=False, | ||
| 164 | + use_registration_helper=False, | ||
| 165 | + ): | ||
| 166 | + self = cls() | ||
| 167 | + process_group = self._init_process_group(rank, world_size, port) | ||
| 168 | + input_tensor = torch.full((40, 20), rank + 1.0).to(device_type) | ||
| 169 | + base_model = _AllreduceModel() | ||
| 170 | + | ||
| 171 | + reference_model = self._ddp( | ||
| 172 | + copy.deepcopy(base_model), | ||
| 173 | + rank, | ||
| 174 | + process_group, | ||
| 175 | + gradient_as_bucket_view, | ||
| 176 | + static_graph, | ||
| 177 | + ) | ||
| 178 | + reference_grads = self._gradient(reference_model, input_tensor) | ||
| 179 | + | ||
| 180 | + hook_model = self._ddp( | ||
| 181 | + copy.deepcopy(base_model), | ||
| 182 | + rank, | ||
| 183 | + process_group, | ||
| 184 | + gradient_as_bucket_view, | ||
| 185 | + static_graph, | ||
| 186 | + ) | ||
| 187 | + hook_state = None if use_none_process_group else process_group | ||
| 188 | + if use_registration_helper: | ||
| 189 | + register_ddp_comm_hook(DDPCommHookType.ALLREDUCE, hook_model, hook_state) | ||
| 190 | + else: | ||
| 191 | + hook_model.register_comm_hook(hook_state, default_hooks.allreduce_hook) | ||
| 192 | + hook_grads = self._gradient(hook_model, input_tensor) | ||
| 193 | + | ||
| 194 | + self.assertEqual(hook_grads, reference_grads) | ||
| 195 | + dist.destroy_process_group() | ||
| 196 | + | ||
| 197 | + | ||
| 198 | + def _run_future_and_dtype_contract(cls, rank, world_size, port): | ||
| 199 | + self = cls() | ||
| 200 | + process_group = self._init_process_group(rank, world_size, port) | ||
| 201 | + | ||
| 202 | + # HCCL-supported floating-point gradient dtypes. | ||
| 203 | + for dtype in (torch.float32, torch.float16, torch.bfloat16): | ||
| 204 | + input_tensor = torch.full((4,), rank + 1.0, dtype=dtype).to(device_type) | ||
| 205 | + model = self._ddp(_AllreduceModel((4,), dtype), rank, process_group) | ||
| 206 | + state = {"process_group": process_group} | ||
| 207 | + model.register_comm_hook(state, _future_contract_hook) | ||
| 208 | + gradients = self._gradient(model, input_tensor, use_mean=False) | ||
| 209 | + expected = torch.full((4,), 1.5, dtype=dtype).to(device_type) | ||
| 210 | + | ||
| 211 | + self.assertEqual(gradients[0], expected) | ||
| 212 | + self.assertTrue(state["is_future"]) | ||
| 213 | + self.assertEqual(state["result_shape"], state["buffer_shape"]) | ||
| 214 | + self.assertEqual(state["result_dtype"], dtype) | ||
| 215 | + self.assertEqual(state["result_device_type"], device_type) | ||
| 216 | + | ||
| 217 | + dist.destroy_process_group() | ||
| 218 | + | ||
| 219 | + | ||
| 220 | + def _run_overflow_boundary(cls, rank, world_size, port): | ||
| 221 | + self = cls() | ||
| 222 | + process_group = self._init_process_group(rank, world_size, port) | ||
| 223 | + model = self._ddp(_AllreduceModel((4,), torch.float16), rank, process_group) | ||
| 224 | + model.register_comm_hook(process_group, default_hooks.allreduce_hook) | ||
| 225 | + input_tensor = torch.full((4,), 60000.0, dtype=torch.float16).to(device_type) | ||
| 226 | + gradient = self._gradient(model, input_tensor, use_mean=False)[0] | ||
| 227 | + | ||
| 228 | + self.assertTrue(torch.isfinite(gradient).all().item()) | ||
| 229 | + self.assertEqual(gradient, torch.full_like(gradient, 60000.0)) | ||
| 230 | + dist.destroy_process_group() | ||
| 231 | + | ||
| 232 | + | ||
| 233 | + def _run_custom_subgroup(cls, rank, world_size, port): | ||
| 234 | + self = cls() | ||
| 235 | + self._init_process_group(rank, world_size, port) | ||
| 236 | + subgroups = [] | ||
| 237 | + try: | ||
| 238 | + for group_rank in range(world_size): | ||
| 239 | + subgroups.append( | ||
| 240 | + dist.new_group( | ||
| 241 | + [group_rank], backend="hccl", timeout=PROCESS_GROUP_TIMEOUT | ||
| 242 | + ) | ||
| 243 | + ) | ||
| 244 | + subgroup = subgroups[rank] | ||
| 245 | + model = self._ddp(_AllreduceModel((4,)), rank, subgroup) | ||
| 246 | + model.register_comm_hook(subgroup, default_hooks.allreduce_hook) | ||
| 247 | + input_tensor = torch.full((4,), rank + 1.0).to(device_type) | ||
| 248 | + gradient = self._gradient(model, input_tensor, use_mean=False)[0] | ||
| 249 | + | ||
| 250 | + self.assertEqual(gradient, torch.full_like(gradient, rank + 1.0)) | ||
| 251 | + finally: | ||
| 252 | + try: | ||
| 253 | + for subgroup in reversed(subgroups): | ||
| 254 | + dist.destroy_process_group(subgroup) | ||
| 255 | + finally: | ||
| 256 | + dist.destroy_process_group() | ||
| 257 | + | ||
| 258 | + | ||
| 259 | + def _run_multiple_buckets(cls, rank, world_size, port): | ||
| 260 | + self = cls() | ||
| 261 | + process_group = self._init_process_group(rank, world_size, port) | ||
| 262 | + base_model = _MultiBucketModel().to(device_type) | ||
| 263 | + input_tensor = torch.full((8, 32), rank + 1.0).to(device_type) | ||
| 264 | + reference_model = self._ddp( | ||
| 265 | + copy.deepcopy(base_model), rank, process_group, bucket_cap_mb=0.001 | ||
| 266 | + ) | ||
| 267 | + hook_model = self._ddp( | ||
| 268 | + copy.deepcopy(base_model), rank, process_group, bucket_cap_mb=0.001 | ||
| 269 | + ) | ||
| 270 | + state = {"process_group": process_group, "calls": 0} | ||
| 271 | + hook_model.register_comm_hook(state, _counting_hook) | ||
| 272 | + | ||
| 273 | + self._gradient(reference_model, input_tensor) | ||
| 274 | + self._gradient(hook_model, input_tensor) | ||
| 275 | + reference_model.zero_grad(set_to_none=True) | ||
| 276 | + hook_model.zero_grad(set_to_none=True) | ||
| 277 | + state["calls"] = 0 | ||
| 278 | + reference_grads = self._gradient(reference_model, input_tensor) | ||
| 279 | + hook_grads = self._gradient(hook_model, input_tensor) | ||
| 280 | + | ||
| 281 | + self.assertGreater(state["calls"], 1) | ||
| 282 | + self.assertEqual(hook_grads, reference_grads) | ||
| 283 | + dist.destroy_process_group() | ||
| 284 | + | ||
| 285 | + | ||
| 286 | + def _spawn(self, worker, *args): | ||
| 287 | + mp.spawn( | ||
| 288 | + worker, | ||
| 289 | + args=(WORLD_SIZE, find_free_port(), *args), | ||
| 290 | + nprocs=WORLD_SIZE, | ||
| 291 | + join=True, | ||
| 292 | + ) | ||
| 293 | + | ||
| 294 | + def test_allreduce_hook(self): | ||
| 295 | + self._spawn(self._run_ddp_parity) | ||
| 296 | + | ||
| 297 | + def test_allreduce_hook_grad_is_view(self): | ||
| 298 | + self._spawn(self._run_ddp_parity, True) | ||
| 299 | + | ||
| 300 | + def test_allreduce_hook_static_graph(self): | ||
| 301 | + self._spawn(self._run_ddp_parity, False, True) | ||
| 302 | + | ||
| 303 | + def test_allreduce_hook_grad_is_view_static_graph(self): | ||
| 304 | + self._spawn(self._run_ddp_parity, True, True) | ||
| 305 | + | ||
| 306 | + def test_allreduce_hook_none_pg(self): | ||
| 307 | + self._spawn(self._run_ddp_parity, False, False, True) | ||
| 308 | + | ||
| 309 | + def test_allreduce_hook_registration_helper(self): | ||
| 310 | + self._spawn(self._run_ddp_parity, False, False, False, True) | ||
| 311 | + | ||
| 312 | + def test_allreduce_hook_future_and_dtypes(self): | ||
| 313 | + self._spawn(self._run_future_and_dtype_contract) | ||
| 314 | + | ||
| 315 | + def test_allreduce_hook_overflow_boundary(self): | ||
| 316 | + self._spawn(self._run_overflow_boundary) | ||
| 317 | + | ||
| 318 | + def test_allreduce_hook_custom_subgroup(self): | ||
| 319 | + self._spawn(self._run_custom_subgroup) | ||
| 320 | + | ||
| 321 | + def test_allreduce_hook_multiple_buckets(self): | ||
| 322 | + self._spawn(self._run_multiple_buckets) | ||
| 323 | + | ||
| 324 | + | ||
| 325 | + def test_allreduce_hook_single_npu_contract(self): | ||
| 326 | + os.environ["MASTER_ADDR"] = "127.0.0.1" | ||
| 327 | + os.environ["MASTER_PORT"] = str(find_free_port()) | ||
| 328 | + torch.accelerator.set_device_index(0) | ||
| 329 | + dist.init_process_group( | ||
| 330 | + "hccl", rank=0, world_size=1, timeout=PROCESS_GROUP_TIMEOUT | ||
| 331 | + ) | ||
| 332 | + try: | ||
| 333 | + process_group = dist.group.WORLD | ||
| 334 | + tensor = torch.tensor([1.003, 2.007]).to(device_type) | ||
| 335 | + expected = tensor.clone() | ||
| 336 | + bucket = _Bucket(tensor) | ||
| 337 | + | ||
| 338 | + future = default_hooks.allreduce_hook(process_group, bucket) | ||
| 339 | + result = future.wait() | ||
| 340 | + | ||
| 341 | + self.assertIsInstance(future, torch._C.Future) | ||
| 342 | + self.assertEqual(result, expected) | ||
| 343 | + self.assertEqual(result.shape, tensor.shape) | ||
| 344 | + self.assertEqual(result.dtype, tensor.dtype) | ||
| 345 | + self.assertEqual(result.device.type, device_type) | ||
| 346 | + finally: | ||
| 347 | + dist.destroy_process_group() | ||
| 348 | + | ||
| 349 | + | ||
| 350 | + def test_allreduce_hook_invalid_arguments(self): | ||
| 351 | + bucket = _Bucket(torch.ones(4).to(device_type)) | ||
| 352 | + | ||
| 353 | + with self.assertRaises(TypeError): | ||
| 354 | + default_hooks.allreduce_hook() | ||
| 355 | + with self.assertRaisesRegex(AttributeError, "has no attribute 'buffer'"): | ||
| 356 | + default_hooks.allreduce_hook(None, None) | ||
| 357 | + with self.assertRaisesRegex(AttributeError, "has no attribute 'size'"): | ||
| 358 | + default_hooks.allreduce_hook("invalid", bucket) | ||
| 359 | + | ||
| 360 | + | ||
| 361 | +if __name__ == "__main__": | ||
| 362 | + run_tests() | ||