已合并
bugfix: rankid can not trans to device id, use current device #11816
AtomGit-Bot创建于 2024年5月18日
bugfix: rankid can not trans to device id, use current device #11816
已合并
从refs/pull/11816/head合入到master
共 2 个文件变更+28-12
| @@ -13,6 +13,9 @@ from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU | |||
| 13 | 13 | ||
| 14 | 14 | ||
| 15 | class TestDataParallel(TestCase): | 15 | class TestDataParallel(TestCase): |
| 16 | + def setUp(self): | ||
| 17 | + super().setUp() | ||
| 18 | + os.environ['PYTORCH_NPU_ALLOC_CONF'] = 'expandable_segments:False' | ||
| 16 | 19 | ||
| 17 | 20 | ||
| 18 | def test_data_parallel_rnn(self): | 21 | def test_data_parallel_rnn(self): |
| @@ -279,6 +282,5 @@ class TestDataParallel(TestCase): | |||
| 279 | 282 | ||
| 280 | 283 | ||
| 281 | if __name__ == "__main__": | 284 | if __name__ == "__main__": |
| 282 | - os.environ['PYTORCH_NPU_ALLOC_CONF'] = 'expandable_segments:False' | ||
| 283 | TestCase._default_dtype_check_enabled = True | 285 | TestCase._default_dtype_check_enabled = True |
| 284 | run_tests() | 286 | run_tests() |
| @@ -18,6 +18,7 @@ | |||
| 18 | 18 | ||
| 19 | 19 | ||
| 20 | 20 | ||
| 21 | + | ||
| 21 | 22 | ||
| 22 | 23 | ||
| 23 | 24 | ||
| @@ -1153,17 +1154,30 @@ int64_t ProcessGroupHCCL::getHcclComm(int rankid) | |||
| 1153 | } | 1154 | } |
| 1154 | 1155 | ||
| 1155 | std::string ProcessGroupHCCL::getHcclCommName(int rankid) { | 1156 | std::string ProcessGroupHCCL::getHcclCommName(int rankid) { |
| 1156 | - at::Device device = getDeviceForRank(rankid); | 1157 | + TORCH_CHECK(rankid >= 0, "Invalid rank ", rankid, DIST_ERROR(ErrCode::VALUE)); |
| 1157 | - std::vector<at::Device> devices = {device}; | 1158 | + auto numNPUs = c10_npu::device_count(); |
| 1158 | - const auto key = getKeyFromDevices(devices); | 1159 | + TORCH_CHECK(numNPUs > 0, "Invalid device number", numNPUs, DIST_ERROR(ErrCode::VALUE)); |
| 1159 | - auto& hcclComms = getHCCLComm(key, devices); | 1160 | + c10::DeviceIndex indexFromRank = static_cast<c10::DeviceIndex>(rankid % numNPUs); |
| 1160 | - TORCH_CHECK(hcclComms.size() == 1, "expect hcclComms.size() = 1, but hcclComms.size() = ", | 1161 | + c10::DeviceIndex indexFromCurDevice = c10_npu::current_device(); |
| 1161 | - hcclComms.size(), DIST_ERROR(ErrCode::VALUE)); | 1162 | + if (indexFromRank != indexFromCurDevice) { |
| 1162 | - HcclComm ret_hcom = hcclComms[0]->getHcclComm(); | 1163 | + std::string warning_message = "The indexFromRank " + std::to_string(indexFromRank) + |
| 1163 | - char commName[MAX_GROUP_NAME_LEN]; | 1164 | + "is not equal indexFromCurDevice " + std::to_string(indexFromCurDevice) + |
| 1164 | - HCCL_CHECK_ERROR(at_npu::hccl::HcclGetCommNameFace(ret_hcom, commName)); | 1165 | + " , which might be normal if the number of devices on your collective communication server is inconsistent." + |
| 1165 | - std::string name_str(commName); | 1166 | + "Otherwise, you need to check if the current device is correct when calling the interface." + |
| 1166 | - return name_str; | 1167 | + "If it's incorrect, it might have introduced an error."; |
| 1168 | + TORCH_WARN_ONCE(warning_message); | ||
| 1169 | + } | ||
| 1170 | + | ||
| 1171 | + at::Device device = at::Device(c10::DeviceType::PrivateUse1, indexFromCurDevice); | ||
| 1172 | + std::vector<at::Device> devices = {device}; | ||
| 1173 | + const auto key = getKeyFromDevices(devices); | ||
| 1174 | + auto& hcclComms = getHCCLComm(key, devices); | ||
| 1175 | + TORCH_CHECK(hcclComms.size() == 1, "expect hcclComms.size() = 1, but hcclComms.size() = ", | ||
| 1176 | + hcclComms.size(), DIST_ERROR(ErrCode::VALUE)); | ||
| 1177 | + HcclComm hcom = hcclComms[0]->getHcclComm(); | ||
| 1178 | + char commName[MAX_GROUP_NAME_LEN] = {}; | ||
| 1179 | + HCCL_CHECK_ERROR(at_npu::hccl::HcclGetCommNameFace(hcom, commName)); | ||
| 1180 | + return std::string(commName); | ||
| 1167 | } | 1181 | } |
| 1168 | 1182 | ||
| 1169 | std::string ProcessGroupHCCL::getHcclCommNameWithoutInit(int rankid, std::vector<std::shared_ptr<HCCLComm>>& hcclComms) | 1183 | std::string ProcessGroupHCCL::getHcclCommNameWithoutInit(int rankid, std::vector<std::shared_ptr<HCCLComm>>& hcclComms) |