已合并
Add support for creating sub-communicators from the global domain in the rootinfo scenario. #34151
limuan创建于 4月22日
Add support for creating sub-communicators from the global domain in the rootinfo scenario. #34151
已合并
共 2 个文件变更+63-30
| @@ -2666,34 +2666,45 @@ bool ProcessGroupHCCL::createHCCLCommEx( | |||
| 2666 | return false; | 2666 | return false; |
| 2667 | } | 2667 | } |
| 2668 | c10_npu::OptionalNPUGuard npuGuard; | 2668 | c10_npu::OptionalNPUGuard npuGuard; |
| 2669 | - // global process group | 2669 | + // global process group only; sub comm derivation is handled by createHCCLCommSub |
| 2670 | - if (options_->global_ranks_in_group.empty() && commType == HcclCommType::DEFAULT) { | 2670 | + if (!(options_->global_ranks_in_group.empty() && commType == HcclCommType::DEFAULT)) { |
| 2671 | - auto startTime = std::chrono::steady_clock::now(); | 2671 | + TORCH_NPU_HCCL_LOGI("createHCCLCommEx only handles global comm with ranktable, skip for sub/P2P comm."); |
| 2672 | - for (size_t i = 0; i < devices.size(); ++i) { | 2672 | + return false; |
| 2673 | - int rank = getRank() * static_cast<int>(devices.size()) + static_cast<int>(i); | ||
| 2674 | - | ||
| 2675 | - npuGuard.set_index(devices[i].index()); | ||
| 2676 | - HcclCommConfig config; | ||
| 2677 | - if (commConfig == nullptr) { | ||
| 2678 | - config = createHcclCommConfigWithOptions(); | ||
| 2679 | - commConfig = &config; | ||
| 2680 | - } | ||
| 2681 | - auto comm = HCCLComm::createGlobalHcclComm(rankTableFile.c_str(), rank, commConfig); | ||
| 2682 | - if (comm == nullptr) { | ||
| 2683 | - TORCH_NPU_HCCL_LOGI("Create global hccl comm with ranktable failed, switch to original interface."); | ||
| 2684 | - return false; | ||
| 2685 | - } | ||
| 2686 | - hcclComms[i] = comm; | ||
| 2687 | - // Creates the HCCL streams | ||
| 2688 | - streamVal.push_back(getHcclNPUStream(devices[i])); | ||
| 2689 | - } | ||
| 2690 | - auto endTime = std::chrono::steady_clock::now(); | ||
| 2691 | - auto timeElapsed = std::chrono::duration_cast<std::chrono::milliseconds>(endTime - startTime); | ||
| 2692 | - TORCH_NPU_HCCL_LOGI("Create global hccl comm with ranktable success, take %d milliseconds", timeElapsed.count()); | ||
| 2693 | - return true; | ||
| 2694 | } | 2673 | } |
| 2674 | + auto startTime = std::chrono::steady_clock::now(); | ||
| 2675 | + for (size_t i = 0; i < devices.size(); ++i) { | ||
| 2676 | + int rank = getRank() * static_cast<int>(devices.size()) + static_cast<int>(i); | ||
| 2695 | 2677 | ||
| 2696 | - // sub process group | 2678 | + npuGuard.set_index(devices[i].index()); |
| 2679 | + HcclCommConfig config; | ||
| 2680 | + if (commConfig == nullptr) { | ||
| 2681 | + config = createHcclCommConfigWithOptions(); | ||
| 2682 | + commConfig = &config; | ||
| 2683 | + } | ||
| 2684 | + auto comm = HCCLComm::createGlobalHcclComm(rankTableFile.c_str(), rank, commConfig); | ||
| 2685 | + if (comm == nullptr) { | ||
| 2686 | + TORCH_NPU_HCCL_LOGI("Create global hccl comm with ranktable failed, switch to original interface."); | ||
| 2687 | + return false; | ||
| 2688 | + } | ||
| 2689 | + hcclComms[i] = comm; | ||
| 2690 | + // Creates the HCCL streams | ||
| 2691 | + streamVal.push_back(getHcclNPUStream(devices[i])); | ||
| 2692 | + } | ||
| 2693 | + auto endTime = std::chrono::steady_clock::now(); | ||
| 2694 | + auto timeElapsed = std::chrono::duration_cast<std::chrono::milliseconds>(endTime - startTime); | ||
| 2695 | + TORCH_NPU_HCCL_LOGI("Create global hccl comm with ranktable success, take %d milliseconds", static_cast<int>(timeElapsed.count())); | ||
| 2696 | + return true; | ||
| 2697 | +} | ||
| 2698 | + | ||
| 2699 | +bool ProcessGroupHCCL::createHCCLCommSub( | ||
| 2700 | + const std::string& devicesKey, | ||
| 2701 | + const std::vector<at::Device>& devices, | ||
| 2702 | + HcclCommType commType, | ||
| 2703 | + HcclCommConfig* commConfig, | ||
| 2704 | + std::vector<std::shared_ptr<HCCLComm>> &hcclComms, | ||
| 2705 | + std::vector<c10_npu::NPUStream> &streamVal, | ||
| 2706 | + int p2pRank) | ||
| 2707 | +{ | ||
| 2697 | if (!hcclCreateSubCommConfigExist()) { | 2708 | if (!hcclCreateSubCommConfigExist()) { |
| 2698 | TORCH_NPU_HCCL_LOGI("The hcclCreateSubCommConfig is not exist, switch to original interface."); | 2709 | TORCH_NPU_HCCL_LOGI("The hcclCreateSubCommConfig is not exist, switch to original interface."); |
| 2699 | return false; | 2710 | return false; |
| @@ -2706,14 +2717,15 @@ bool ProcessGroupHCCL::createHCCLCommEx( | |||
| 2706 | try { | 2717 | try { |
| 2707 | globalHcclComm = global_->getHcclCommByDevices(devices); | 2718 | globalHcclComm = global_->getHcclCommByDevices(devices); |
| 2708 | } catch (const std::exception& e) { | 2719 | } catch (const std::exception& e) { |
| 2709 | - TORCH_NPU_HCCL_LOGI("create the global HCCL Communicator failed, the exception info is %s, switch to original interface.", e.what()); | 2720 | + TORCH_NPU_HCCL_LOGI("Get global HCCL communicator failed: %s, switch to original interface.", e.what()); |
| 2710 | return false; | 2721 | return false; |
| 2711 | } | 2722 | } |
| 2712 | if (!globalHcclComm) { | 2723 | if (!globalHcclComm) { |
| 2713 | - TORCH_NPU_HCCL_LOGI("Create sub hccl comm by hcclCreateSubCommConfig failed, globalHcclComm is nullptr, switch to original interface."); | 2724 | + TORCH_NPU_HCCL_LOGI("Create sub hccl comm failed, globalHcclComm is nullptr, switch to original interface."); |
| 2714 | return false; | 2725 | return false; |
| 2715 | } | 2726 | } |
| 2716 | 2727 | ||
| 2728 | + c10_npu::OptionalNPUGuard npuGuard; | ||
| 2717 | uint64_t hcclid = (std::hash<string>{}(options_->group_id)); | 2729 | uint64_t hcclid = (std::hash<string>{}(options_->group_id)); |
| 2718 | auto subStartTime = std::chrono::steady_clock::now(); | 2730 | auto subStartTime = std::chrono::steady_clock::now(); |
| 2719 | for (size_t i = 0; i < devices.size(); ++i) { | 2731 | for (size_t i = 0; i < devices.size(); ++i) { |
| @@ -2780,7 +2792,7 @@ bool ProcessGroupHCCL::createHCCLCommEx( | |||
| 2780 | auto subEndTime = std::chrono::steady_clock::now(); | 2792 | auto subEndTime = std::chrono::steady_clock::now(); |
| 2781 | auto subTimeElapsed = std::chrono::duration_cast<std::chrono::milliseconds>(subEndTime - subStartTime); | 2793 | auto subTimeElapsed = std::chrono::duration_cast<std::chrono::milliseconds>(subEndTime - subStartTime); |
| 2782 | TORCH_NPU_HCCL_LOGI("Create sub hccl comm by hcclCreateSubCommConfig success, group id is %s, subCommId is %llu, devicesKey is %s, use %d ms.", | 2794 | TORCH_NPU_HCCL_LOGI("Create sub hccl comm by hcclCreateSubCommConfig success, group id is %s, subCommId is %llu, devicesKey is %s, use %d ms.", |
| 2783 | - options_->group_id.c_str(), hcclid, devicesKey.c_str(), subTimeElapsed.count()); | 2795 | + options_->group_id.c_str(), hcclid, devicesKey.c_str(), static_cast<int>(subTimeElapsed.count())); |
| 2784 | return true; | 2796 | return true; |
| 2785 | } | 2797 | } |
| 2786 | 2798 | ||
| @@ -2861,7 +2873,19 @@ std::vector<std::shared_ptr<HCCLComm>>& ProcessGroupHCCL::createHCCLComm( | |||
| 2861 | }; | 2873 | }; |
| 2862 | at_npu::native::OpCommand::RunOpApiV3("hcclGroupEnd", hccl_call); | 2874 | at_npu::native::OpCommand::RunOpApiV3("hcclGroupEnd", hccl_call); |
| 2863 | } | 2875 | } |
| 2864 | - if (!createHCCLCommEx(devicesKey, devices, commType, commConfig, hcclComms, streamVal, p2pRank)) { | 2876 | + bool isSubComm = !(options_->global_ranks_in_group.empty() && commType == HcclCommType::DEFAULT); |
| 2877 | + bool created = false; | ||
| 2878 | + if (isSubComm) { | ||
| 2879 | + // Sub comm: derive from global_ via hcclCreateSubCommConfig (works for both ranktable and rootinfo) | ||
| 2880 | + created = createHCCLCommSub(devicesKey, devices, commType, commConfig, hcclComms, streamVal, p2pRank); | ||
| 2881 | + if (!created) { | ||
| 2882 | + TORCH_NPU_HCCL_LOGI("Sub comm derivation failed, fallback to original interface."); | ||
| 2883 | + } | ||
| 2884 | + } else { | ||
| 2885 | + // Global comm: try ranktable path first | ||
| 2886 | + created = createHCCLCommEx(devicesKey, devices, commType, commConfig, hcclComms, streamVal, p2pRank); | ||
| 2887 | + } | ||
| 2888 | + if (!created) { | ||
| 2865 | createHCCLCommOrigin(devicesKey, devices, commType, commConfig, hcclComms, streamVal, p2pRank); | 2889 | createHCCLCommOrigin(devicesKey, devices, commType, commConfig, hcclComms, streamVal, p2pRank); |
| 2866 | } | 2890 | } |
| 2867 | // restart the HcclGroupStart | 2891 | // restart the HcclGroupStart |
| @@ -1273,6 +1273,15 @@ private: | |||
| 1273 | std::vector<c10_npu::NPUStream> &streamVal, | 1273 | std::vector<c10_npu::NPUStream> &streamVal, |
| 1274 | int p2pRank); | 1274 | int p2pRank); |
| 1275 | 1275 | ||
| 1276 | + bool createHCCLCommSub( | ||
| 1277 | + const std::string& devicesKey, | ||
| 1278 | + const std::vector<at::Device>& devices, | ||
| 1279 | + HcclCommType commType, | ||
| 1280 | + HcclCommConfig* commConfig, | ||
| 1281 | + std::vector<std::shared_ptr<HCCLComm>> &hcclComms, | ||
| 1282 | + std::vector<c10_npu::NPUStream> &streamVal, | ||
| 1283 | + int p2pRank); | ||
| 1284 | + | ||
| 1276 | void createHCCLCommForZeroCopy( | 1285 | void createHCCLCommForZeroCopy( |
| 1277 | std::vector<std::shared_ptr<HCCLComm>> &hcclComms, | 1286 | std::vector<std::shared_ptr<HCCLComm>> &hcclComms, |
| 1278 | std::unordered_map<std::string, std::string> &envMap); | 1287 | std::unordered_map<std::string, std::string> &envMap); |
新增代码逻辑,补充UT防护