Calculate the communication group between all pipeline stages
()
| 116 | |
| 117 | |
| 118 | def _get_pipeline_group(): |
| 119 | """ |
| 120 | |
| 121 | Calculate the communication group between all pipeline stages |
| 122 | |
| 123 | """ |
| 124 | rank = get_rank() |
| 125 | stage_nums = auto_parallel_context().get_pipeline_stages() |
| 126 | device_nums = get_group_size() |
| 127 | per_stage_device_nums = device_nums // stage_nums |
| 128 | local_stage_rank_id = rank % per_stage_device_nums |
| 129 | group = range(0, stage_nums) |
| 130 | rank_list = [local_stage_rank_id + x * per_stage_device_nums for x in group] |
| 131 | rank_str_list = [str(local_stage_rank_id + x * per_stage_device_nums) for x in group] |
| 132 | rank_list_str = "-".join(rank_str_list) |
| 133 | return rank_list, rank_list_str |
| 134 | |
| 135 | |
| 136 | class GlobalNorm(nn.Cell): |