已合并
add ut for torch.futures #7891
misty-rain-typhoid创建于 2023年11月29日
add ut for torch.futures #7891
已合并
从refs/pull/7891/head合入到master
共 1 个文件变更+76-0
| @@ -0,0 +1,76 @@ | |||
| 1 | +import copy | ||
| 2 | +import os | ||
| 3 | + | ||
| 4 | +import torch | ||
| 5 | +from torch.distributed.algorithms.ddp_comm_hooks import default_hooks | ||
| 6 | +import torch.distributed as dist | ||
| 7 | +import torch.multiprocessing as mp | ||
| 8 | +from torch import nn | ||
| 9 | + | ||
| 10 | +import torch_npu | ||
| 11 | +from torch_npu.testing.testcase import TestCase, run_tests | ||
| 12 | +from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU | ||
| 13 | + | ||
| 14 | + | ||
| 15 | +class TestDdpCommHook(nn.Module): | ||
| 16 | + def __init__(self): | ||
| 17 | + super().__init__() | ||
| 18 | + self.p = nn.Parameter(torch.randn(40, 20)) | ||
| 19 | + | ||
| 20 | + def forward(self, x): | ||
| 21 | + return self.p * x | ||
| 22 | + | ||
| 23 | + | ||
| 24 | +class HcomAllReduceTest(TestCase): | ||
| 25 | + | ||
| 26 | + def _init_dist_hccl(cls, rank, world_size): | ||
| 27 | + os.environ['MASTER_ADDR'] = '127.0.0.1' | ||
| 28 | + os.environ['MASTER_PORT'] = '29500' | ||
| 29 | + os.environ['HCCL_WHITELIST_DISABLE'] = '1' | ||
| 30 | + torch_npu.npu.set_device(rank) | ||
| 31 | + return dist.init_process_group(backend='hccl', world_size=world_size, rank=rank) | ||
| 32 | + | ||
| 33 | + | ||
| 34 | + def _get_grad(cls, model, train_data): | ||
| 35 | + output = model(train_data) | ||
| 36 | + output.mean().backward() | ||
| 37 | + param = next(model.parameters()) | ||
| 38 | + return param.grad | ||
| 39 | + | ||
| 40 | + | ||
| 41 | + def _test_hook(cls, rank, world_size, hook_type): | ||
| 42 | + torch.npu.manual_seed(0) | ||
| 43 | + torch.manual_seed(0) | ||
| 44 | + pg = HcomAllReduceTest._init_dist_hccl(rank, world_size) | ||
| 45 | + torch.npu.set_device(rank) | ||
| 46 | + origin_model = TestDdpCommHook() | ||
| 47 | + train_data = torch.randn(40, 20).npu() | ||
| 48 | + no_hook_model = nn.parallel.DistributedDataParallel(copy.deepcopy(origin_model).npu(), device_ids=[rank]) | ||
| 49 | + no_hook_grad = HcomAllReduceTest._get_grad(no_hook_model, train_data) | ||
| 50 | + | ||
| 51 | + hook_model = nn.parallel.DistributedDataParallel(copy.deepcopy(origin_model).npu(), device_ids=[rank]) | ||
| 52 | + hook_model.register_comm_hook(state=pg, hook=hook_type) | ||
| 53 | + hook_grad = HcomAllReduceTest._get_grad(hook_model, train_data) | ||
| 54 | + TestCase().assertEqual(hook_grad, no_hook_grad) | ||
| 55 | + | ||
| 56 | + | ||
| 57 | + def test_fp16_compress_hook(self): | ||
| 58 | + # CI currently supports only 2 devices | ||
| 59 | + world_size = 2 | ||
| 60 | + mp.spawn(HcomAllReduceTest._test_hook, | ||
| 61 | + args=(world_size, default_hooks.fp16_compress_hook,), | ||
| 62 | + nprocs=world_size, | ||
| 63 | + join=True) | ||
| 64 | + | ||
| 65 | + | ||
| 66 | + def test_allreduce_hook(self): | ||
| 67 | + # CI currently supports only 2 devices | ||
| 68 | + world_size = 2 | ||
| 69 | + mp.spawn(HcomAllReduceTest._test_hook, | ||
| 70 | + args=(world_size, default_hooks.allreduce_hook,), | ||
| 71 | + nprocs=world_size, | ||
| 72 | + join=True) | ||
| 73 | + | ||
| 74 | + | ||
| 75 | +if __name__ == '__main__': | ||
| 76 | + run_tests() | ||