diff --git a/astra-sim-alibabacloud/astra-sim/system/MockNcclGroup.cc b/astra-sim-alibabacloud/astra-sim/system/MockNcclGroup.cc index 4daf68d5..365dcd29 100644 --- a/astra-sim-alibabacloud/astra-sim/system/MockNcclGroup.cc +++ b/astra-sim-alibabacloud/astra-sim/system/MockNcclGroup.cc @@ -106,25 +106,23 @@ namespace MockNccl { std::set EPnodes; for (int i = 0; i < TP_nums / _EP_size; i++){ TP_idx = i*_EP_size; - for(int j =0;j<_EP_size;j++){ - for(int k = 0;k DP_EP_nodes; for (int i = 0; i < TP_nums / _DP_EP_size; i++){ TP_idx = i; - for (int j = 0; j < _DP_EP_size; j++){ - for (int k = 0; k < AllTPGroups[TP_idx].Ranks.size(); k++){ - ranks.clear(); - DP_EP_nodes.clear(); - for (int l = TP_idx; l < TP_idx + _DP_EP_size * _EP_size; l += _EP_size){ - int tmp_rank = AllTPGroups[l].Ranks[k]; - int node_idx = tmp_rank / _gpus_per_nodes; - ranks.push_back(tmp_rank); - GroupIndex[std::make_pair(tmp_rank, DP_EP)] = all_group_idx; - DP_EP_nodes.insert(node_idx); - } - NVSwitchs.clear(); - for (int idx : DP_EP_nodes){ - NVSwitchs.push_back(_NVSwitch[idx]); - GroupIndex[std::make_pair(_NVSwitch[idx], DP_EP)] = all_group_idx; - } - AllGroups[all_group_idx] = GroupInfo(all_group_idx, DP_EP, DP_EP_nodes.size(), _DP_EP_size, ranks, NVSwitchs); - all_group_idx++; + for (int k = 0; k < AllTPGroups[TP_idx].Ranks.size(); k++){ + ranks.clear(); + DP_EP_nodes.clear(); + for (int l = TP_idx; l < TP_idx + _DP_EP_size * _EP_size; l += _EP_size){ + int tmp_rank = AllTPGroups[l].Ranks[k]; + int node_idx = tmp_rank / _gpus_per_nodes; + ranks.push_back(tmp_rank); + GroupIndex[std::make_pair(tmp_rank, DP_EP)] = all_group_idx; + DP_EP_nodes.insert(node_idx); + } + NVSwitchs.clear(); + for (int idx : DP_EP_nodes){ + NVSwitchs.push_back(_NVSwitch[idx]); + GroupIndex[std::make_pair(_NVSwitch[idx], DP_EP)] = all_group_idx; } + AllGroups[all_group_idx] = GroupInfo(all_group_idx, DP_EP, DP_EP_nodes.size(), _DP_EP_size, ranks, NVSwitchs); + all_group_idx++; } } } @@ -2100,4 +2096,4 @@ namespace MockNccl { return info; } } -} \ No newline at end of file +}