(self, check_using_norm)
| 78 | world_size = 2 |
| 79 | |
| 80 | def test(self, check_using_norm): |
| 81 | groups._create_expert_and_data_parallel(2) |
| 82 | |
| 83 | param1 = torch.nn.Parameter(torch.Tensor([0])) |
| 84 | param1.grad = torch.Tensor([1]) |
| 85 | param2 = torch.nn.Parameter(torch.Tensor([0])) |
| 86 | if dist.get_rank() == 0: |
| 87 | param2.grad = torch.Tensor([1]) |
| 88 | else: |
| 89 | param2.grad = torch.Tensor([float("inf")]) |
| 90 | param2.allreduce = False |
| 91 | # param2 is now MoE parameter |
| 92 | parameters = [param1, param2] |
| 93 | if check_using_norm: |
| 94 | grads_group_flat = [_flatten_dense_tensors([p.grad for p in parameters])] |
| 95 | norm = ds_utils.get_weight_norm(grads_group_flat) |
| 96 | overflow_checker = ds_utils.CheckOverflow([parameters]) |
| 97 | overflow = overflow_checker.check_using_norm([norm], reduce_overflow=False) |
| 98 | else: |
| 99 | overflow_checker = ds_utils.CheckOverflow([parameters]) |
| 100 | overflow = overflow_checker.check() |
| 101 | assert overflow |
| 102 | |
| 103 | |
| 104 | @pytest.mark.skipif(not hasattr(torch.autograd.graph, "_get_grad_fn_or_grad_acc"), |
nothing calls this directly
no test coverage detected