(devices, model, net, param, use_nccl=False, control_input=None)
| 1038 | |
| 1039 | |
| 1040 | def _AllReduce(devices, model, net, param, use_nccl=False, control_input=None): |
| 1041 | blobs_group = list(model._device_grouped_blobs[param].values()) |
| 1042 | if model._device_type == caffe2_pb2.CUDA and use_nccl: |
| 1043 | # TODO: for _shared_model, do only NCCLReduce |
| 1044 | model.NCCLAllreduce( |
| 1045 | blobs_group, blobs_group, control_input=control_input |
| 1046 | ) |
| 1047 | return |
| 1048 | |
| 1049 | if model._device_type == workspace.GpuDeviceType: |
| 1050 | p2p_access_pattern = workspace.GetGpuPeerAccessPattern() |
| 1051 | else: |
| 1052 | p2p_access_pattern = None |
| 1053 | |
| 1054 | def sumN(*dev_indices): |
| 1055 | """Create a Sum op for 2 or more blobs on different devices. |
| 1056 | Saves the result on the first device. |
| 1057 | |
| 1058 | Args: |
| 1059 | dev_indices -- a list of device indices, which can be translated into |
| 1060 | CUDA identifiers with model._devices |
| 1061 | """ |
| 1062 | devices = [model._devices[idx] for idx in dev_indices] |
| 1063 | blobs = [blobs_group[idx] for idx in dev_indices] |
| 1064 | device_opt = core.DeviceOption(model._device_type, devices[0]) |
| 1065 | with core.DeviceScope(device_opt): |
| 1066 | for i, peer in enumerate(devices): |
| 1067 | if i == 0: |
| 1068 | continue # Skip the first device |
| 1069 | if p2p_access_pattern is not None and p2p_access_pattern.size and not p2p_access_pattern[ |
| 1070 | devices[0], peer |
| 1071 | ]: |
| 1072 | # Copy from peer to d0 |
| 1073 | blobs[i] = model.Copy( |
| 1074 | blobs[i], |
| 1075 | 'gpu_{}/{}_gpu{}_copy'.format(devices[0], param, peer) |
| 1076 | ) |
| 1077 | net.Sum(blobs, [blobs[0]], name='dpm') |
| 1078 | |
| 1079 | if len(devices) == 16: |
| 1080 | # Special tree reduction for 16 gpus, TODO generalize like in muji.py |
| 1081 | for j in range(8): |
| 1082 | sumN(j * 2, j * 2 + 1) |
| 1083 | for j in range(4): |
| 1084 | sumN(j * 4, j * 4 + 2) |
| 1085 | for j in range(2): |
| 1086 | sumN(j * 8, j * 8 + 4) |
| 1087 | sumN(0, 8) |
| 1088 | elif len(devices) == 8: |
| 1089 | for j in range(4): |
| 1090 | sumN(j * 2, j * 2 + 1) |
| 1091 | for j in range(2): |
| 1092 | sumN(j * 4, j * 4 + 2) |
| 1093 | sumN(0, 4) |
| 1094 | elif len(devices) == 4: |
| 1095 | sumN(0, 1) |
| 1096 | sumN(2, 3) |
| 1097 | sumN(0, 2) |
no test coverage detected
searching dependent graphs…