已合并
test: Add test cases for torch.distributed.algorithms.ddp_comm_hooks.default_hooks.fp16_compress_wrapper to cover all scenarios. #42840
创建于 7月26日
test: Add test cases for torch.distributed.algorithms.ddp_comm_hooks.default_hooks.fp16_compress_wrapper to cover all scenarios. #42840
已合并
创建于 7月26日
已删除 :test_fp16_compress_wrapper_v2.12.0合入到Ascend/pytorchv2.12.0
1 个文件变更+441-0
@@ -0,0 +1,441 @@
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 hook wrappers.
21+2. This file validates
22+ torch.distributed.algorithms.ddp_comm_hooks.default_hooks.fp16_compress_wrapper
23+ (extendable).
24+"""
25+ 
26+import copy
27+import os
28+ 
29+import torch
30+import torch.distributed as dist
31+import torch.multiprocessing as mp
32+from torch import nn
33+from torch.distributed.algorithms.ddp_comm_hooks import default_hooks
34+from torch.distributed.algorithms.ddp_comm_hooks import powerSGD_hook as powerSGD
35+from torch.testing._internal.common_utils import TestCase, find_free_port, run_tests
36+ 
37+from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU
38+ 
39+ 
40+device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"
41+ 
42+WORLD_SIZE = 2
43+ 
44+ 
45+class _Fp16WrapperModel(nn.Module):
46+ def __init__(self, shape=(40, 20), dtype=torch.float32):
47+ super().__init__()
48+ self.weight = nn.Parameter(torch.ones(shape, dtype=dtype).to(device_type))
49+ 
50+ def forward(self, input_tensor):
51+ return self.weight * input_tensor
52+ 
53+ 
54+class _MultiBucketModel(nn.Module):
55+ def __init__(self):
56+ super().__init__()
57+ self.layers = nn.ModuleList([nn.Linear(32, 32, bias=False) for _ in range(4)])
58+ 
59+ def forward(self, input_tensor):
60+ return sum(layer(input_tensor) for layer in self.layers)
61+ 
62+ 
63+class _MutableBucket:
64+ def __init__(self, tensor):
65+ self.tensor = tensor
66+ 
67+ def buffer(self):
68+ return self.tensor
69+ 
70+ def set_buffer(self, tensor):
71+ self.tensor = tensor
72+ 
73+ 
74+def _recording_allreduce_hook(
75+ state, bucket: dist.GradBucket
76+) -> torch.futures.Future[torch.Tensor]:
77+ state["inner_called"] = True
78+ state["inner_dtype"] = bucket.buffer().dtype
79+ future = default_hooks.allreduce_hook(state["process_group"], bucket)
80+ state["inner_is_future"] = isinstance(future, torch._C.Future)
81+ 
82+ def validate(fut):
83+ result = fut.value()
84+ state["inner_result_dtype"] = result.dtype
85+ return result
86+ 
87+ return future.then(validate)
88+ 
89+ 
90+def _wrapper_contract_hook(
91+ state, bucket: dist.GradBucket
92+) -> torch.futures.Future[torch.Tensor]:
93+ buffer_shape = bucket.buffer().shape
94+ future = state["wrapped_hook"](state, bucket)
95+ state["wrapper_is_future"] = isinstance(future, torch._C.Future)
96+ 
97+ def validate(fut):
98+ result = fut.value()
99+ state["wrapper_result_shape"] = result.shape
100+ state["wrapper_result_dtype"] = result.dtype
101+ state["wrapper_result_device_type"] = result.device.type
102+ state["original_shape"] = buffer_shape
103+ return result
104+ 
105+ return future.then(validate)
106+ 
107+ 
108+def _counting_allreduce_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+def _tensor_returning_hook(state, bucket):
116+ return bucket.buffer()
117+ 
118+ 
119+def _recording_future_hook(state, bucket):
120+ state["called"] = True
121+ state["dtype"] = bucket.buffer().dtype
122+ future = torch.futures.Future()
123+ future.set_result(bucket.buffer().add(1))
124+ return future
125+ 
126+ 
127+def _raising_hook(state, bucket):
128+ raise RuntimeError("inner hook failure")
129+ 
130+ 
131+class TestFp16CompressWrapper(TestCase):
132+ @staticmethod
133+ def _init_process_group(rank, world_size, port):
134+ os.environ["MASTER_ADDR"] = "127.0.0.1"
135+ os.environ["MASTER_PORT"] = str(port)
136+ torch.accelerator.set_device_index(rank)
137+ dist.init_process_group("hccl", rank=rank, world_size=world_size)
138+ return dist.group.WORLD
139+ 
140+ @staticmethod
141+ def _ddp(
142+ model,
143+ rank,
144+ process_group,
145+ gradient_as_bucket_view=False,
146+ static_graph=False,
147+ bucket_cap_mb=None,
148+ ):
149+ kwargs = {
150+ "device_ids": [rank],
151+ "process_group": process_group,
152+ "gradient_as_bucket_view": gradient_as_bucket_view,
153+ "static_graph": static_graph,
154+ }
155+ if bucket_cap_mb is not None:
156+ kwargs["bucket_cap_mb"] = bucket_cap_mb
157+ return nn.parallel.DistributedDataParallel(model, **kwargs)
158+ 
159+ @staticmethod
160+ def _gradient(model, input_tensor, use_mean=True):
161+ output = model(input_tensor)
162+ loss = output.mean() if use_mean else output.sum()
163+ loss.backward()
164+ return [parameter.grad.detach().clone() for parameter in model.parameters()]
165+ 
166+ @classmethod
167+ def _run_ddp_parity(
168+ cls,
169+ rank,
170+ world_size,
171+ port,
172+ gradient_as_bucket_view=False,
173+ static_graph=False,
174+ use_none_process_group=False,
175+ use_power_sgd=False,
176+ ):
177+ self = cls()
178+ process_group = self._init_process_group(rank, world_size, port)
179+ input_tensor = torch.full((40, 20), rank + 1.0).to(device_type)
180+ base_model = _Fp16WrapperModel()
181+ reference_model = self._ddp(
182+ copy.deepcopy(base_model),
183+ rank,
184+ process_group,
185+ gradient_as_bucket_view,
186+ static_graph,
187+ )
188+ if use_power_sgd:
189+ inner_hook = powerSGD.powerSGD_hook
190+ hook_state = powerSGD.PowerSGDState(
191+ process_group=process_group,
192+ start_powerSGD_iter=2,
193+ )
194+ else:
195+ inner_hook = default_hooks.allreduce_hook
196+ hook_state = None if use_none_process_group else process_group
197+ wrapped_hook = default_hooks.fp16_compress_wrapper(inner_hook)
198+ hook_model = self._ddp(
199+ copy.deepcopy(base_model),
200+ rank,
201+ process_group,
202+ gradient_as_bucket_view,
203+ static_graph,
204+ )
205+ hook_model.register_comm_hook(hook_state, wrapped_hook)
206+ 
207+ iterations = 3 if use_power_sgd else 1
208+ for _ in range(iterations):
209+ reference_model.zero_grad(set_to_none=True)
210+ hook_model.zero_grad(set_to_none=True)
211+ reference_grads = self._gradient(
212+ reference_model, input_tensor, use_mean=False
213+ )
214+ hook_grads = self._gradient(
215+ hook_model, input_tensor, use_mean=False
216+ )
217+ 
218+ if use_power_sgd:
219+ self.assertGreater(hook_state.iter, hook_state.start_powerSGD_iter)
220+ self.assertGreater(hook_state.total_numel_after_compression, 0)
221+ self.assertEqual(hook_grads, reference_grads, rtol=1e-3, atol=1e-3)
222+ else:
223+ self.assertEqual(hook_grads, reference_grads)
224+ dist.destroy_process_group()
225+ 
226+ @staticmethod
227+ def _compressed_average(dtype):
228+ values = []
229+ for value in (1.003, 2.007):
230+ tensor = torch.full((4,), value, dtype=dtype).to(device_type)
231+ values.append(tensor.to(torch.float16).div_(WORLD_SIZE))
232+ return (values[0] + values[1]).to(dtype)
233+ 
234+ @classmethod
235+ def _run_future_dtype_and_state_contract(cls, rank, world_size, port):
236+ self = cls()
237+ process_group = self._init_process_group(rank, world_size, port)
238+ 
239+ # HCCL-supported floating-point gradient dtypes.
240+ for dtype in (torch.float32, torch.float16, torch.bfloat16):
241+ value = (1.003, 2.007)[rank]
242+ input_tensor = torch.full((4,), value, dtype=dtype).to(device_type)
243+ model = self._ddp(_Fp16WrapperModel((4,), dtype), rank, process_group)
244+ wrapped_hook = default_hooks.fp16_compress_wrapper(
245+ _recording_allreduce_hook
246+ )
247+ state = {
248+ "process_group": process_group,
249+ "wrapped_hook": wrapped_hook,
250+ }
251+ model.register_comm_hook(state, _wrapper_contract_hook)
252+ gradients = self._gradient(model, input_tensor, use_mean=False)
253+ 
254+ self.assertEqual(gradients[0], self._compressed_average(dtype))
255+ self.assertEqual(gradients[0].dtype, dtype)
256+ self.assertTrue(state["inner_called"])
257+ self.assertTrue(state["inner_is_future"])
258+ self.assertTrue(state["wrapper_is_future"])
259+ self.assertEqual(state["inner_dtype"], torch.float16)
260+ self.assertEqual(state["inner_result_dtype"], torch.float16)
261+ self.assertEqual(state["wrapper_result_dtype"], torch.float16)
262+ self.assertEqual(
263+ state["wrapper_result_shape"], state["original_shape"]
264+ )
265+ self.assertEqual(state["wrapper_result_device_type"], device_type)
266+ 
267+ dist.destroy_process_group()
268+ 
269+ @classmethod
270+ def _run_custom_subgroup(cls, rank, world_size, port):
271+ self = cls()
272+ self._init_process_group(rank, world_size, port)
273+ subgroups = []
274+ try:
275+ for group_rank in range(world_size):
276+ subgroups.append(dist.new_group([group_rank], backend="hccl"))
277+ subgroup = subgroups[rank]
278+ model = self._ddp(_Fp16WrapperModel((4,)), rank, subgroup)
279+ wrapped_hook = default_hooks.fp16_compress_wrapper(
280+ default_hooks.allreduce_hook
281+ )
282+ model.register_comm_hook(subgroup, wrapped_hook)
283+ input_tensor = torch.full((4,), rank + 1.0).to(device_type)
284+ gradient = self._gradient(model, input_tensor, use_mean=False)[0]
285+ 
286+ self.assertEqual(gradient, torch.full_like(gradient, rank + 1.0))
287+ finally:
288+ try:
289+ for subgroup in reversed(subgroups):
290+ dist.destroy_process_group(subgroup)
291+ finally:
292+ dist.destroy_process_group()
293+ 
294+ @classmethod
295+ def _run_predivide_overflow_boundary(cls, rank, world_size, port):
296+ self = cls()
297+ process_group = self._init_process_group(rank, world_size, port)
298+ model = self._ddp(_Fp16WrapperModel((4,)), rank, process_group)
299+ wrapped_hook = default_hooks.fp16_compress_wrapper(
300+ default_hooks.allreduce_hook
301+ )
302+ model.register_comm_hook(process_group, wrapped_hook)
303+ # Pre-division avoids a 120000 FP16 intermediate during all-reduce.
304+ input_tensor = torch.full((4,), 60000.0).to(device_type)
305+ output = model(input_tensor)
306+ output.backward(torch.ones_like(output))
307+ gradient = next(model.parameters()).grad.detach().clone()
308+ 
309+ self.assertTrue(torch.isfinite(gradient).all().item())
310+ self.assertEqual(gradient, torch.full_like(gradient, 60000.0))
311+ dist.destroy_process_group()
312+ 
313+ @classmethod
314+ def _run_multiple_buckets(cls, rank, world_size, port):
315+ self = cls()
316+ process_group = self._init_process_group(rank, world_size, port)
317+ base_model = _MultiBucketModel().to(device_type)
318+ input_tensor = torch.full((8, 32), rank + 1.0).to(device_type)
319+ reference_model = self._ddp(
320+ copy.deepcopy(base_model), rank, process_group, bucket_cap_mb=0.001
321+ )
322+ hook_model = self._ddp(
323+ copy.deepcopy(base_model), rank, process_group, bucket_cap_mb=0.001
324+ )
325+ state = {"process_group": process_group, "calls": 0}
326+ wrapped_hook = default_hooks.fp16_compress_wrapper(
327+ _counting_allreduce_hook
328+ )
329+ hook_model.register_comm_hook(state, wrapped_hook)
330+ 
331+ self._gradient(reference_model, input_tensor, use_mean=False)
332+ self._gradient(hook_model, input_tensor, use_mean=False)
333+ reference_model.zero_grad(set_to_none=True)
334+ hook_model.zero_grad(set_to_none=True)
335+ state["calls"] = 0
336+ reference_grads = self._gradient(
337+ reference_model, input_tensor, use_mean=False
338+ )
339+ hook_grads = self._gradient(hook_model, input_tensor, use_mean=False)
340+ 
341+ self.assertGreater(state["calls"], 1)
342+ self.assertEqual(hook_grads, reference_grads)
343+ dist.destroy_process_group()
344+ 
345+ @skipIfUnsupportMultiNPU(WORLD_SIZE)
346+ def _spawn(self, worker, *args):
347+ mp.spawn(
348+ worker,
349+ args=(WORLD_SIZE, find_free_port(), *args),
350+ nprocs=WORLD_SIZE,
351+ join=True,
352+ )
353+ 
354+ def test_fp16_compress_wrapper_allreduce(self):
355+ self._spawn(self._run_ddp_parity)
356+ 
357+ def test_fp16_compress_wrapper_allreduce_grad_is_view(self):
358+ self._spawn(self._run_ddp_parity, True)
359+ 
360+ def test_fp16_compress_wrapper_allreduce_static_graph(self):
361+ self._spawn(self._run_ddp_parity, False, True)
362+ 
363+ def test_fp16_compress_wrapper_allreduce_grad_is_view_static_graph(self):
364+ self._spawn(self._run_ddp_parity, True, True)
365+ 
366+ def test_fp16_compress_wrapper_allreduce_none_pg(self):
367+ self._spawn(self._run_ddp_parity, False, False, True)
368+ 
369+ def test_fp16_compress_wrapper_powersgd(self):
370+ self._spawn(self._run_ddp_parity, False, False, False, True)
371+ 
372+ def test_fp16_compress_wrapper_powersgd_grad_is_view(self):
373+ self._spawn(self._run_ddp_parity, True, False, False, True)
374+ 
375+ def test_fp16_compress_wrapper_powersgd_static_graph(self):
376+ self._spawn(self._run_ddp_parity, False, True, False, True)
377+ 
378+ def test_fp16_compress_wrapper_powersgd_grad_is_view_static_graph(self):
379+ self._spawn(self._run_ddp_parity, True, True, False, True)
380+ 
381+ def test_fp16_compress_wrapper_future_dtype_and_state(self):
382+ self._spawn(self._run_future_dtype_and_state_contract)
383+ 
384+ def test_fp16_compress_wrapper_custom_subgroup(self):
385+ self._spawn(self._run_custom_subgroup)
386+ 
387+ def test_fp16_compress_wrapper_predivide_overflow_boundary(self):
388+ self._spawn(self._run_predivide_overflow_boundary)
389+ 
390+ def test_fp16_compress_wrapper_multiple_buckets(self):
391+ self._spawn(self._run_multiple_buckets)
392+ 
393+ @skipIfUnsupportMultiNPU(1)
394+ def test_fp16_compress_wrapper_single_npu_contract(self):
395+ tensor = torch.tensor([1.003, 2.007]).to(device_type)
396+ bucket = _MutableBucket(tensor)
397+ state = {"called": False}
398+ wrapped_hook = default_hooks.fp16_compress_wrapper(
399+ _recording_future_hook
400+ )
401+ 
402+ future = wrapped_hook(state, bucket)
403+ result = future.wait()
404+ 
405+ self.assertTrue(state["called"])
406+ self.assertEqual(state["dtype"], torch.float16)
407+ self.assertIsInstance(future, torch._C.Future)
408+ self.assertEqual(result, tensor.to(torch.float16).add(1))
409+ self.assertEqual(result.dtype, torch.float16)
410+ self.assertEqual(result.device.type, device_type)
411+ 
412+ @skipIfUnsupportMultiNPU(1)
413+ def test_fp16_compress_wrapper_invalid_arguments(self):
414+ tensor = torch.ones(4).to(device_type)
415+ 
416+ with self.assertRaises(TypeError):
417+ default_hooks.fp16_compress_wrapper()
418+ with self.assertRaises(AttributeError):
419+ default_hooks.fp16_compress_wrapper(default_hooks.allreduce_hook)(
420+ None, None
421+ )
422+ with self.assertRaises(TypeError):
423+ default_hooks.fp16_compress_wrapper(None)(
424+ None, _MutableBucket(tensor)
425+ )
426+ with self.assertRaises(AttributeError):
427+ default_hooks.fp16_compress_wrapper(_tensor_returning_hook)(
428+ None, _MutableBucket(tensor)
429+ )
430+ with self.assertRaises(RuntimeError):
431+ default_hooks.fp16_compress_wrapper(_raising_hook)(
432+ None, _MutableBucket(tensor)
433+ )
434+ with self.assertRaises(AttributeError):
435+ default_hooks.fp16_compress_wrapper(default_hooks.allreduce_hook)(
436+ None, (tensor,)
437+ )
438+ 
439+ 
440+if __name__ == "__main__":
441+ run_tests()