frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Enables NCCL symmetric memory kernels through mempool registration (#155134)

Browse Source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/155134
Approved by: https://github.com/kwen2501

Co-authored-by: Ke Wen <kw2501@meta.com>
This commit is contained in:
Syed Tousif Ahmed
2025-06-21 13:08:13 -07:00
committed by PyTorch MergeBot
parent 9e132b770e
commit f70c80105e
9 changed files with 146 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
torch/csrc/distributed/c10d/NCCLUtils.hpp
View File

@ -63,6 +63,10 @@ static_assert(
#define NCCL_HAS_COMM_REGISTER
#endif
#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 27, 0)
#define NCCL_HAS_COMM_WINDOW_REGISTER
#endif
#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 19, 0)
#define NCCL_HAS_MEM_ALLOC
#endif
@ -341,9 +345,10 @@ class NCCLComm {
ncclResult_t registerSegment(
void* ptr,
size_t size,
bool errorOnRereg = true);
bool errorOnRereg = true,
bool window = false);
ncclResult_t deregisterSegment(void* ptr);
ncclResult_t deregisterSegment(void* ptr, bool window = false);
std::string repr() const;