| 34 | return res |
| 35 | |
| 36 | def specific_group_split(group_spec, share_backbone_group_ids, \ |
| 37 | share_neck_group_ids, share_decoder_group_ids): |
| 38 | ## sanity check |
| 39 | assert type(group_spec) is list |
| 40 | assert all(map(lambda x: type(x) is int, group_spec)) |
| 41 | |
| 42 | num_groups = len(group_spec) |
| 43 | splits = np.sum(group_spec) |
| 44 | |
| 45 | world_size = dist.get_world_size() |
| 46 | rank = dist.get_rank() |
| 47 | |
| 48 | assert world_size % splits == 0 |
| 49 | unit = int(world_size / splits) |
| 50 | |
| 51 | ## split |
| 52 | group_sizes = [x*unit for x in group_spec] # [8,8,8] / [32, 16] |
| 53 | groups = [] |
| 54 | roots = [] |
| 55 | last = 0 |
| 56 | task_info = edict() |
| 57 | all_ranks = [] |
| 58 | for i,gs in enumerate(group_sizes): |
| 59 | ranks = list(map(int, np.arange(last, last+gs))) #[0...8], [9...15], ... |
| 60 | groups.append(dist.new_group(ranks=ranks)) |
| 61 | roots.append(last) # 0, 8, 16 |
| 62 | all_ranks.append(ranks) |
| 63 | if rank in ranks: # if current gpu rank in traversed rank task group |
| 64 | printlog(f">> task_info.group[{i}] ranks {ranks}") |
| 65 | task_info.group = groups[-1] # subordinate to what group |
| 66 | task_info.task_size = gs # 8 |
| 67 | task_info.task_id = i |
| 68 | task_info.task_rank = rank - last |
| 69 | task_info.task_root_rank = last |
| 70 | last += gs |
| 71 | task_info.root_group = dist.new_group(ranks=roots) |
| 72 | printlog(f">> task_info.root_group ranks {roots}") |
| 73 | task_info.task_sizes = group_sizes |
| 74 | task_info.task_root_ranks = roots |
| 75 | task_info.task_num = num_groups |
| 76 | |
| 77 | ## share_backbone_group spec |
| 78 | if share_backbone_group_ids is not None: # *[0,0,0]*(default) | [0,1,0]task ids |
| 79 | # group size must equal within a share_group |
| 80 | backboneshareid2idx = {} |
| 81 | for idx, this_id in enumerate(share_backbone_group_ids): |
| 82 | if this_id not in backboneshareid2idx: |
| 83 | backboneshareid2idx[this_id] = list() |
| 84 | backboneshareid2idx[this_id].append(idx) # {0: [0,1,2]}| {0: [0,2], 1: [1]} |
| 85 | |
| 86 | ## create backbone share group |
| 87 | for idxs in backboneshareid2idx.values(): # idxs = [0, 1, 2] |
| 88 | this_group_ranks = flat([all_ranks[i] for i in idxs]) |
| 89 | this_share_group = dist.new_group(ranks=this_group_ranks) |
| 90 | this_group_size = len(this_group_ranks) |
| 91 | if rank in this_group_ranks: |
| 92 | task_info.backbone_share_group = this_share_group |
| 93 | printlog(f">> task_info.backbone_share_group[{idxs}] ranks {this_group_ranks}") |