Add MSCCLPP user buffer registration APIs and integrate with RCCL #1477

nusislam · 2025-01-02T15:42:05Z

Details

This PR adds user buffer registration APIs in MSCCLPP and integrates them with the corresponding RCCL APIs.

Work item: "Internal", or link to GitHub issue (if applicable).

What were the changes?
A patch has been added in ext-src. This patch will be applied to MSCCLPP during RCCL build. Also MSCCLPP related changes have been added to the RCCL user buffer registration code.

Why were the changes made?
These changes are required to use MSCCLPP in non-HIP-graph mode.

How was the outcome achieved?
Technical details behind the work. Explain any publicly-available hardware peculiarities.

Additional Documentation:
What else should the reviewer know?

Approval Checklist

Do not approve until these items are satisfied.

Verify the CHANGELOG has been updated, if
- there are any NCCL API version changes,
- any changes impact library users, and/or
- any changes impact any other ROCm library.

…s registered

mustafabar · 2025-01-03T16:03:43Z

ext-src/mem-reg.patch

+}
+
+NCCL_API ncclResult_t ncclCommDeregister(ncclComm_t comm, void* handle) {
+        if (comm && handle) {


Shouldn't we free any resources associated with handle here?
Example:

auto it = comm->channelScratchInfos.find(buffKey); if (it != comm->channelScratchInfos.end()) { comm->channelScratchInfos.erase(it); }

Same thing can be done for channelOut and channelIn

mustafabar · 2025-01-03T16:44:01Z

ext-src/mem-reg.patch

+NCCL_API ncclResult_t ncclCommDeregister(ncclComm_t comm, void* handle) {
+        if (comm && handle) {
+                handle = nullptr;
+        }


The implementation nullifies the handle inside the function, but since it's passed as void*, the change doesn't reflect in the caller's context. This can lead to dangling pointers. I suggest passing handle as a void** to resolve this issue, if API change is possible

Suggested change

+NCCL_API ncclResult_t ncclCommDeregister(ncclComm_t comm, void* handle) {

+ if (comm && handle) {

+ handle = nullptr;

+ }

+NCCL_API ncclResult_t ncclCommDeregister(ncclComm_t comm, void** handle) {

+ if (comm && handle && *handle) {

+ *handle = nullptr;

+ }

This is a NCCL API and cannot be changed. I am still working on some of these fixes.

mustafabar · 2025-01-03T16:49:12Z

ext-src/mem-reg.patch

+  std::vector<mscclpp::RegisteredMemory> remoteMemories;
+
+  // Creating the channels
+  auto buffIt = comm->channelScratchInfos.find(buffKey);


If there's no guarantee this will be called from a single thread, we probably need to lock access to channelScratchInfos and the rest of the maps. I suggest using one lock if this function doesn't fall on the critical path.

This map is supposed to be checked by the host process and MSCCLPP is invoked only when using one GPU per process case.

nusislam · 2025-01-06T16:18:28Z

ext-src/mem-reg.patch

 /*
  * Send
  *
 diff --git a/apps/nccl/src/nccl.cu b/apps/nccl/src/nccl.cu
 index a697be2..d8497e7 100644
 --- a/apps/nccl/src/nccl.cu
 +++ b/apps/nccl/src/nccl.cu
-@@ -577,6 +577,58 @@ NCCL_API ncclResult_t ncclAllGather(const void* sendbuff, void* recvbuff, size_t
+@@ -577,6 +577,67 @@ NCCL_API ncclResult_t ncclAllGather(const void* sendbuff, void* recvbuff, size_t


This change does not seem to be necessary

nusislam · 2025-01-06T19:03:39Z

src/register.cc

-  cache->population -= 1;
+
+  #ifdef ENABLE_MSCCLPP
+    if (comm->mscclCompatible){


Here you should do the same checks that you did before registering the buffer.
if (comm->mscclCompatible && size > 0 && (size & 31) == 0 && size <= comm->mscclpp_threshold){

corey-derochie-amd · 2025-01-06T23:38:23Z

ext-src/mem-reg.patch

+ncclResult_t ncclCommRegister(ncclComm_t comm, void* buff, size_t size, void** handle);
+ncclResult_t ncclCommDeregister(ncclComm_t comm, void* handle);
+ncclResult_t ncclBuffIsRegistered(ncclComm_t comm, const void* buff, size_t count, bool* registered);
+size_t ncclBufferSize(ncclComm_t comm, void* handle);


Since you're inserting functions into mscclpp, you can just give them names such as mscclppCommRegister, and then there's no need to rename via mscclpp_nccl_syms.txt.

ncclCommRegister() and ncclCommDeRegister() are NCCL APIs. We need to keep these names so that a workload can also be run with the mscclpp backend.

I am curious why ncclCommRegister and ncclCommDeRegister are not implement in mscclpp? The stub functions already exist: https://github.com/microsoft/mscclpp/blob/6d26b92665383b60ffae7e4dd9ab4cc047d34721/apps/nccl/src/nccl.cu#L809C1-L817C2

@nusislam Should these changes be submitted as a mscclpp PR?

Eventually, yes. But at the moment, will use this patch.

nusislam · 2025-01-08T15:37:10Z

@corey-derochie-amd - Any more reviews?

isaki001 · 2025-01-08T22:29:05Z

There is a typo I need to fix. Will push commit shortly

corey-derochie-amd · 2025-01-10T18:14:15Z

src/register.cc

+      else{
+        WARN("MSCCL++: Cannot register user-buffers on managed memory");
+      }


This will do nothing, issue a warning, then fall through and return ncclSuccess. Is that the right behaviour? Should it not return an error?

I think the expectation is that if no registration happens in MSCCL++ due to the buffer being in managed memory, no MSCCL++ or RCCL registration will happen and execution will fall back to RCCL kernel without UBR.

I actually need to do a small adjustment to not cause a RCCL deregistration error when MSCCL++ detects no registered buffer because of this. I just tested a simple fix for this.

But since we always fallback to RCCL kernel when MSCCL++ is not enabled, I think a warning without return an error is OK. Also, I haven't seen any errors returned in any UBR related routines, except when RCCL cannot find the registration handle during deregistration. Of course, this definitely makes it harder to see if anything goes wrong.

… back to rccl deRegister which will return an error

nusislam · 2025-01-10T19:25:34Z

@isaki001 - I think the helper routines do not have to be named as NCCL APIs; the return type can also be changed. Other than this, the PR looks good to me.

corey-derochie-amd

Left some suggestions.

src/register.cc

cmake/MSCCLPP.cmake

src/register.cc

isaki001 · 2025-01-10T19:54:31Z

Thanks @corey-derochie-amd I'll add your suggestions and I'll change the names as @nusislam indicated.

Whitespace suggestions and reducing diffs to avoid future merge conflicts Co-authored-by: corey-derochie-amd <[email protected]>

…stration with a buffer in managed memory

corey-derochie-amd · 2025-01-13T16:20:03Z

src/register.cc

@@ -165,11 +165,11 @@ ncclResult_t ncclCommRegister_impl(const ncclComm_t comm, void* buff, size_t siz
      if(!isManagedBuffer){
        INFO(NCCL_INIT, "MSCCL++: ncclCommRegister");
        NCCLCHECK(mscclpp_ncclCommRegister(comm->mscclpp_comm, buff, size, handle));
+        return ncclSuccess;


I'm debating whether this is the right behaviour... Are these buffers reused in ways that might call MSCCL++ sometimes and RCCL other times? Are buffers always matched to message size, or could they use a buffer that's larger than the given message size? If so, the safest thing to do might be to register for BOTH RCCL and MSCCL++. Thoughts, @nusislam ?

The cases when mscclppRegister should be used vs. rcclRegister should be used are mutually exclusive.

ext-src: add MSCCLPP memory registration APIs

6c32a32

nusislam requested review from wenkaidu, gilbertlee-amd, akolliasAMD, edgargabriel, PedramAlizadeh, nileshnegi, KawtharShafie, AtlantaPepsi, mberenjk, corey-derochie-amd, mustafabar, thananon and haripriya-amd as code owners January 2, 2025 15:42

nusislam marked this pull request as draft January 2, 2025 17:03

isaki001 added 2 commits January 2, 2025 18:12

update mem-reg patch with mscclpp helper routine to check if buffer i…

a552427

…s registered

RCCL integration of MSCCL++ user-buffer registration APIs

184cea2

nusislam changed the title ~~ext-src: add MSCCLPP memory registration APIs~~ Add MSCCLPP user buffer registration APIs and integrate with RCCL Jan 3, 2025

nusislam marked this pull request as ready for review January 3, 2025 03:26

only include mscclpp_nccl header if ENABLE_MSCCLPP is defined

8617629

mustafabar reviewed Jan 3, 2025

View reviewed changes

ext-src: update mscclpp mem-reg patch

8bfaf30

nusislam commented Jan 6, 2025

View reviewed changes

isaki001 added 2 commits January 6, 2025 15:08

add helper routine to patch

4f2dff1

check handle before MSCCL++ deregister

8568024

corey-derochie-amd reviewed Jan 6, 2025

View reviewed changes

Merge branch 'develop' into mem-reg

de5fa5c

isaki001 requested a review from JhaShweta1 as a code owner January 8, 2025 21:53

fix typo to replace send buff with recv buff

46c0438

corey-derochie-amd reviewed Jan 10, 2025

View reviewed changes

in case of no mscclpp registration, dduring deRegister call, ont fall…

ad036d4

… back to rccl deRegister which will return an error

corey-derochie-amd approved these changes Jan 10, 2025

View reviewed changes

corey-derochie-amd reviewed Jan 10, 2025

View reviewed changes

src/register.cc Outdated Show resolved Hide resolved

isaki001 and others added 4 commits January 10, 2025 16:08

Apply suggestions from code review

5afa3e5

Whitespace suggestions and reducing diffs to avoid future merge conflicts Co-authored-by: corey-derochie-amd <[email protected]>

rename helper functions and change their return type

4a9aa09

resolve conflict

65dacd5

set RCCL user-buffer registration to occur if attempting MSCCL++ regi…

373b9f9

…stration with a buffer in managed memory

corey-derochie-amd reviewed Jan 13, 2025

View reviewed changes

isaki001 merged commit e9b6bbc into ROCm:develop Jan 14, 2025
24 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add MSCCLPP user buffer registration APIs and integrate with RCCL #1477

Add MSCCLPP user buffer registration APIs and integrate with RCCL #1477

nusislam commented Jan 2, 2025 •

edited

Loading

mustafabar Jan 3, 2025

mustafabar Jan 3, 2025

nusislam Jan 3, 2025

mustafabar Jan 3, 2025

nusislam Jan 3, 2025

nusislam Jan 6, 2025

nusislam Jan 6, 2025

corey-derochie-amd Jan 6, 2025

nusislam Jan 6, 2025 •

edited

Loading

dsidler Jan 7, 2025

corey-derochie-amd Jan 10, 2025

nusislam Jan 10, 2025

nusislam commented Jan 8, 2025

isaki001 commented Jan 8, 2025

corey-derochie-amd Jan 10, 2025

isaki001 Jan 10, 2025

nusislam commented Jan 10, 2025

corey-derochie-amd left a comment

isaki001 commented Jan 10, 2025

corey-derochie-amd Jan 13, 2025

nusislam Jan 13, 2025

Add MSCCLPP user buffer registration APIs and integrate with RCCL #1477

Add MSCCLPP user buffer registration APIs and integrate with RCCL #1477

Conversation

nusislam commented Jan 2, 2025 • edited Loading

Details

Approval Checklist

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nusislam Jan 6, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nusislam commented Jan 8, 2025

isaki001 commented Jan 8, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nusislam commented Jan 10, 2025

corey-derochie-amd left a comment

Choose a reason for hiding this comment

isaki001 commented Jan 10, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nusislam commented Jan 2, 2025 •

edited

Loading

nusislam Jan 6, 2025 •

edited

Loading