Calculate the communication group of model parallel dim in one pipeline stage
(mp)
| 96 | |
| 97 | |
| 98 | def _get_model_parallel_group(mp): |
| 99 | """ |
| 100 | |
| 101 | Calculate the communication group of model parallel dim in one pipeline stage |
| 102 | |
| 103 | """ |
| 104 | rank = get_rank() |
| 105 | stage_nums = auto_parallel_context().get_pipeline_stages() |
| 106 | device_nums = get_group_size() |
| 107 | per_stage_device_nums = device_nums // stage_nums |
| 108 | stage_id = rank // per_stage_device_nums |
| 109 | local_stage_rank_id = rank % per_stage_device_nums |
| 110 | index = local_stage_rank_id // mp |
| 111 | group = range(0, mp) |
| 112 | rank_str_list = [str(x + index * mp + stage_id * per_stage_device_nums) for x in group] |
| 113 | rank_list_str = "-".join(rank_str_list) |
| 114 | rank_list = [x + index * mp + stage_id * per_stage_device_nums for x in group] |
| 115 | return rank_list, rank_list_str |
| 116 | |
| 117 | |
| 118 | def _get_pipeline_group(): |