12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879 |
- # Copyright (c) Microsoft Corporation.
- # SPDX-License-Identifier: Apache-2.0
- # DeepSpeed Team
- import torch
- from torch._utils import _flatten_dense_tensors
- import deepspeed.comm as dist
- import pytest
- import deepspeed.runtime.utils as ds_utils
- import deepspeed.utils.groups as groups
- from deepspeed.accelerator import get_accelerator
- from unit.common import DistributedTest
- def test_call_to_str():
- c2s = ds_utils.call_to_str
- assert c2s('int') == 'int()'
- assert c2s('int', 3) == 'int(3)'
- assert c2s('int', 3, 'jeff') == 'int(3, \'jeff\')'
- assert c2s('hello', val=3) == 'hello(val=3)'
- assert c2s('hello', 1138, val=3) == 'hello(1138, val=3)'
- class TestClibGradNorm(DistributedTest):
- world_size = 2
- def test(self):
- param1 = torch.nn.Parameter(torch.Tensor([0]))
- param1.grad = torch.Tensor([1])
- param2 = torch.nn.Parameter(torch.Tensor([0]))
- param2.grad = torch.Tensor([dist.get_rank() + 1])
- # param2 is now MoE parameter
- param2.allreduce = False
- parameters = [param1, param2]
- groups._create_expert_and_data_parallel(2)
- norm = ds_utils.clip_grad_norm_(parameters, max_norm=0.1)
- norm = torch.Tensor([norm]).to(get_accelerator().device_name(dist.get_rank()))
- world_size = dist.get_world_size()
- gathered_norm = [torch.zeros(1).to(get_accelerator().device_name()) for i in range(world_size)]
- dist.all_gather(gathered_norm, norm)
- assert gathered_norm[0] == gathered_norm[1], "norm at rank 0 does not match the norm at rank 1"
- @pytest.mark.parametrize("check_using_norm", [(False), (True)])
- class TestCheckOverflow(DistributedTest):
- world_size = 2
- def test(self, check_using_norm):
- groups._create_expert_and_data_parallel(2)
- param1 = torch.nn.Parameter(torch.Tensor([0]))
- param1.grad = torch.Tensor([1])
- param2 = torch.nn.Parameter(torch.Tensor([0]))
- if dist.get_rank() == 0:
- param2.grad = torch.Tensor([1])
- else:
- param2.grad = torch.Tensor([float("inf")])
- param2.allreduce = False
- # param2 is now MoE parameter
- parameters = [param1, param2]
- if check_using_norm:
- grads_group_flat = [_flatten_dense_tensors([p.grad for p in parameters])]
- norm = ds_utils.get_weight_norm(grads_group_flat)
- overflow_checker = ds_utils.CheckOverflow([parameters])
- overflow = overflow_checker.check_using_norm([norm], reduce_overflow=False)
- else:
- overflow_checker = ds_utils.CheckOverflow([parameters])
- overflow = overflow_checker.check()
- assert overflow
|