Serve zero-chunk relay geometries with a direct fallback instead of failing (meta-pytorch#3678)

Summary:
Pull Request resolved: meta-pytorch#3678

X-link: meta-pytorch/torchrec#4570

When a group's aligned relay chunk rounded down to zero the four sharded relay
collectives returned `ncclInvalidArgument` and required the caller to retry on a
plain collective. In a fused multi-group call that is worse than it sounds: a
single small or awkwardly sized group forced the whole fused call to fail even
when the other groups were large enough to relay, so heterogeneous group sizes
were effectively unsupported.

Give every collective a phase-symmetric per-group direct fallback instead. When
`chunkSize` aligns down to zero the segment is covered by the two direct regions
(`dirA` takes half, `dirB` absorbs the remainder) over the active-to-active link,
the helper scatter/forward is skipped for that group, and the helper-side
reduction is skipped as well. Groups that do have a workable chunk size keep the
existing relay schedule unchanged, so each group in a fused call now picks its
own schedule independently.

- track `dirASizes` per group so the direct regions are sized from the selected
  schedule rather than assumed equal to `chunkSize`
- guard the helper scatter, forward, and reduction on `chunkSize > 0`
- all-gather additionally tracks `dirAOffsets`, which is no longer derivable
  from `relayTotals` once the fallback can move the direct region to offset 0
- correct the minimum helper geometry to H+2 chunks

Also fixes a torchrec test that asserted the A>2 reduce-scatter helper buffer was
`_passthrough_helper_size(...)` when production has always sized it through
`_relay_helper_size()` (`2 * recvCount`); the expectation now matches the
unchanged production contract.

Behavior for buffers that already had a nonzero aligned chunk is unchanged.

Differential Revision: D115998367

Author: srinathb-meta

comms/rcclx/develop/meta/relay/sharded_relay_all_gather.cc

@@ -184,7 +184,8 @@ static ncclResult_t shardedRelayAllGather2Active(
   const int numChunks = numHelpers + 2;
 
   size_t chunkSizes[SHARDED_RELAY_MAX_GROUPS];
-  size_t relayTotals[SHARDED_RELAY_MAX_GROUPS]; // == direct chunk A's offset
+  size_t dirAOffsets[SHARDED_RELAY_MAX_GROUPS];
+  size_t dirASizes[SHARDED_RELAY_MAX_GROUPS];
   size_t dirBOffsets[SHARDED_RELAY_MAX_GROUPS];
   size_t dirBSizes[SHARDED_RELAY_MAX_GROUPS]; // absorbs the remainder
 
@@ -194,22 +195,27 @@ static ncclResult_t shardedRelayAllGather2Active(
     // Zero-count groups are skipped by every loop below.
     if (count == 0) {
       chunkSizes[g] = 0;
-      relayTotals[g] = 0;
+      dirAOffsets[g] = 0;
+      dirASizes[g] = 0;
       dirBOffsets[g] = 0;
       dirBSizes[g] = 0;
       continue;
     }
 
-    // A chunk size that rounds down to zero means the send buffer is too small
-    // to scatter; the caller should fall back to a regular all-gather.
     size_t chunkSize = count / numChunks;
     chunkSize = (chunkSize / CHUNK_ALIGN_ELEMENTS) * CHUNK_ALIGN_ELEMENTS;
+    chunkSizes[g] = chunkSize;
     if (chunkSize == 0) {
-      return ncclInvalidArgument;
+      dirAOffsets[g] = 0;
+      dirASizes[g] = count / 2;
+      dirBOffsets[g] = dirASizes[g];
+      dirBSizes[g] = count - dirASizes[g];
+      continue;
     }
-    chunkSizes[g] = chunkSize;
-    relayTotals[g] = static_cast<size_t>(numHelpers) * chunkSize;
-    dirBOffsets[g] = relayTotals[g] + chunkSize;
+
+    dirAOffsets[g] = static_cast<size_t>(numHelpers) * chunkSize;
+    dirASizes[g] = chunkSize;
+    dirBOffsets[g] = dirAOffsets[g] + dirASizes[g];
     dirBSizes[g] = count - dirBOffsets[g];
   }
 
@@ -265,7 +271,7 @@ static ncclResult_t shardedRelayAllGather2Active(
       char* recvbuff = static_cast<char*>(recvBuffs[g]);
 
       // Scatter: chunk h of my sendBuff goes to helper h.
-      for (int h = 0; h < cfg.numHelpers; h++) {
+      for (int h = 0; h < cfg.numHelpers && chunkSize > 0; h++) {
         NCCLCHECK(ncclSend(
             sendbuff + static_cast<size_t>(h) * chunkSize * elementSize,
             chunkSize,
@@ -279,20 +285,20 @@ static ncclResult_t shardedRelayAllGather2Active(
       // gather slot never overlaps the send source, so this is safe in-place.
       int partner = cfg.activeRanks[1 - cfg.myActiveIndex];
       NCCLCHECK(ncclSend(
-          sendbuff + relayTotals[g] * elementSize,
-          chunkSize,
+          sendbuff + dirAOffsets[g] * elementSize,
+          dirASizes[g],
           datatype,
           partner,
           comm,
           stream));
       NCCLCHECK(ncclRecv(
-          recvbuff + (gatherSlotOffset + relayTotals[g]) * elementSize,
-          chunkSize,
+          recvbuff + (gatherSlotOffset + dirAOffsets[g]) * elementSize,
+          dirASizes[g],
           datatype,
           partner,
           comm,
           stream));
-    } else {
+    } else if (chunkSize > 0) {
       // Helper: receive active rank a's chunk into slot a.
       char* helperBuf = static_cast<char*>(recvBuffs[g]);
       for (int a = 0; a < cfg.nActiveRanks; a++) {
@@ -326,7 +332,7 @@ static ncclResult_t shardedRelayAllGather2Active(
     if (cfg.isActiveRank) {
       char* recvbuff = static_cast<char*>(recvBuffs[g]);
 
-      for (int h = 0; h < cfg.numHelpers; h++) {
+      for (int h = 0; h < cfg.numHelpers && chunkSize > 0; h++) {
         NCCLCHECK(ncclRecv(
             recvbuff +
                 (gatherSlotOffset + static_cast<size_t>(h) * chunkSize) *
@@ -354,7 +360,7 @@ static ncclResult_t shardedRelayAllGather2Active(
           partner,
           comm,
           stream));
-    } else {
+    } else if (chunkSize > 0) {
       const char* helperBuf = static_cast<const char*>(recvBuffs[g]);
       for (int a = 0; a < cfg.nActiveRanks; a++) {
         NCCLCHECK(ncclSend(

comms/rcclx/develop/meta/relay/sharded_relay_all_to_all.cc

@@ -181,6 +181,7 @@ static ncclResult_t shardedRelayAllToAll2Active(
 
   size_t chunkSizes[SHARDED_RELAY_MAX_GROUPS];
   size_t relayTotals[SHARDED_RELAY_MAX_GROUPS]; // == direct chunk A's offset
+  size_t dirASizes[SHARDED_RELAY_MAX_GROUPS];
   size_t dirBOffsets[SHARDED_RELAY_MAX_GROUPS];
   size_t dirBSizes[SHARDED_RELAY_MAX_GROUPS]; // absorbs the remainder
 
@@ -191,21 +192,23 @@ static ncclResult_t shardedRelayAllToAll2Active(
     if (count == 0) {
       chunkSizes[g] = 0;
       relayTotals[g] = 0;
+      dirASizes[g] = 0;
       dirBOffsets[g] = 0;
       dirBSizes[g] = 0;
       continue;
     }
 
-    // A chunk size that rounds down to zero means the segment is too small to
-    // scatter; the caller should fall back to a regular all-to-all.
     size_t chunkSize = count / numChunks;
     chunkSize = (chunkSize / CHUNK_ALIGN_ELEMENTS) * CHUNK_ALIGN_ELEMENTS;
-    if (chunkSize == 0) {
-      return ncclInvalidArgument;
-    }
     chunkSizes[g] = chunkSize;
     relayTotals[g] = static_cast<size_t>(numHelpers) * chunkSize;
-    dirBOffsets[g] = relayTotals[g] + chunkSize;
+    if (chunkSize == 0) {
+      dirASizes[g] = count / 2;
+      dirBOffsets[g] = dirASizes[g];
+    } else {
+      dirASizes[g] = chunkSize;
+      dirBOffsets[g] = relayTotals[g] + chunkSize;
+    }
     dirBSizes[g] = count - dirBOffsets[g];
   }
 
@@ -253,33 +256,35 @@ static ncclResult_t shardedRelayAllToAll2Active(
           static_cast<char*>(recvBuffs[g]) + exchangeSegOffset * elementSize;
 
       // Scatter: chunk h of my exchange segment goes to helper h.
-      for (int h = 0; h < cfg.numHelpers; h++) {
-        NCCLCHECK(ncclSend(
-            sendSeg + static_cast<size_t>(h) * chunkSize * elementSize,
-            chunkSize,
-            datatype,
-            cfg.helperRanks[h],
-            comm,
-            stream));
+      if (chunkSize > 0) {
+        for (int h = 0; h < cfg.numHelpers; h++) {
+          NCCLCHECK(ncclSend(
+              sendSeg + static_cast<size_t>(h) * chunkSize * elementSize,
+              chunkSize,
+              datatype,
+              cfg.helperRanks[h],
+              comm,
+              stream));
+        }
       }
 
       // Direct chunk A over the otherwise-idle active<->active link.
       int partner = cfg.activeRanks[1 - cfg.myActiveIndex];
       NCCLCHECK(ncclSend(
           sendSeg + relayTotals[g] * elementSize,
-          chunkSize,
+          dirASizes[g],
           datatype,
           partner,
           comm,
           stream));
       NCCLCHECK(ncclRecv(
           recvSeg + relayTotals[g] * elementSize,
-          chunkSize,
+          dirASizes[g],
           datatype,
           partner,
           comm,
           stream));
-    } else {
+    } else if (chunkSize > 0) {
       // Helper: receive active rank a's chunk into slot a.
       char* helperBuf = static_cast<char*>(recvBuffs[g]);
       for (int a = 0; a < cfg.nActiveRanks; a++) {
@@ -316,14 +321,16 @@ static ncclResult_t shardedRelayAllToAll2Active(
       char* recvSeg =
           static_cast<char*>(recvBuffs[g]) + exchangeSegOffset * elementSize;
 
-      for (int h = 0; h < cfg.numHelpers; h++) {
-        NCCLCHECK(ncclRecv(
-            recvSeg + static_cast<size_t>(h) * chunkSize * elementSize,
-            chunkSize,
-            datatype,
-            cfg.helperRanks[h],
-            comm,
-            stream));
+      if (chunkSize > 0) {
+        for (int h = 0; h < cfg.numHelpers; h++) {
+          NCCLCHECK(ncclRecv(
+              recvSeg + static_cast<size_t>(h) * chunkSize * elementSize,
+              chunkSize,
+              datatype,
+              cfg.helperRanks[h],
+              comm,
+              stream));
+        }
       }
 
       // Direct chunk B, again over the idle active<->active link.
@@ -342,7 +349,7 @@ static ncclResult_t shardedRelayAllToAll2Active(
           partner,
           comm,
           stream));
-    } else {
+    } else if (chunkSize > 0) {
       const char* helperBuf = static_cast<const char*>(recvBuffs[g]);
       for (int a = 0; a < cfg.nActiveRanks; a++) {
         NCCLCHECK(ncclSend(

comms/rcclx/develop/meta/relay/sharded_relay_allreduce.cc

@@ -674,6 +674,7 @@ static ncclResult_t shardedRelayAllReduce2Active(
 
   size_t chunkSizes[SHARDED_RELAY_MAX_GROUPS];
   size_t relayTotals[SHARDED_RELAY_MAX_GROUPS]; // == direct chunk A's offset
+  size_t dirASizes[SHARDED_RELAY_MAX_GROUPS];
   size_t dirBOffsets[SHARDED_RELAY_MAX_GROUPS];
   size_t dirBSizes[SHARDED_RELAY_MAX_GROUPS]; // absorbs the remainder
 
@@ -685,21 +686,23 @@ static ncclResult_t shardedRelayAllReduce2Active(
     if (count == 0) {
       chunkSizes[g] = 0;
       relayTotals[g] = 0;
+      dirASizes[g] = 0;
       dirBOffsets[g] = 0;
       dirBSizes[g] = 0;
       continue;
     }
 
-    // A chunk size that rounds down to zero means the buffer is too small to
-    // scatter; the caller should fall back to a regular allreduce.
     size_t chunkSize = count / numChunks;
     chunkSize = (chunkSize / CHUNK_ALIGN_ELEMENTS) * CHUNK_ALIGN_ELEMENTS;
-    if (chunkSize == 0) {
-      return ncclInvalidArgument;
-    }
     chunkSizes[g] = chunkSize;
     relayTotals[g] = static_cast<size_t>(numHelpers) * chunkSize;
-    dirBOffsets[g] = relayTotals[g] + chunkSize;
+    if (chunkSize == 0) {
+      dirASizes[g] = count / 2;
+      dirBOffsets[g] = dirASizes[g];
+    } else {
+      dirASizes[g] = chunkSize;
+      dirBOffsets[g] = relayTotals[g] + chunkSize;
+    }
     dirBSizes[g] = count - dirBOffsets[g];
   }
 
@@ -753,7 +756,7 @@ static ncclResult_t shardedRelayAllReduce2Active(
       const char* sendbuff = static_cast<const char*>(sendBuffs[g]);
 
       // Scatter: chunk h goes to helper h.
-      for (int h = 0; h < cfg.numHelpers; h++) {
+      for (int h = 0; h < cfg.numHelpers && chunkSize > 0; h++) {
         NCCLCHECK(ncclSend(
             sendbuff + static_cast<size_t>(h) * chunkSize * elementSize,
             chunkSize,
@@ -765,16 +768,18 @@ static ncclResult_t shardedRelayAllReduce2Active(
 
       // Direct chunk A over the otherwise-idle active<->active link.
       int partner = cfg.activeRanks[1 - cfg.myActiveIndex];
-      NCCLCHECK(ncclSend(
-          sendbuff + relayTotals[g] * elementSize,
-          chunkSize,
-          datatype,
-          partner,
-          comm,
-          stream));
-      NCCLCHECK(
-          ncclRecv(directDst(0), chunkSize, datatype, partner, comm, stream));
-    } else {
+      if (dirASizes[g] > 0) {
+        NCCLCHECK(ncclSend(
+            sendbuff + relayTotals[g] * elementSize,
+            dirASizes[g],
+            datatype,
+            partner,
+            comm,
+            stream));
+        NCCLCHECK(ncclRecv(
+            directDst(0), dirASizes[g], datatype, partner, comm, stream));
+      }
+    } else if (chunkSize > 0) {
       // Helper: receive active rank a's chunk into slot a.
       char* helperBuf = static_cast<char*>(recvBuffs[g]);
       for (int a = 0; a < cfg.nActiveRanks; a++) {
@@ -801,7 +806,7 @@ static ncclResult_t shardedRelayAllReduce2Active(
   // fused add+scale over numHelpers/numChunks of the whole buffer. The work is
   // also spread over every helper GPU rather than piled onto the two actives.
   for (int g = 0; g < nGroups; g++) {
-    if (counts[g] == 0 || configs[g].isActiveRank)
+    if (counts[g] == 0 || chunkSizes[g] == 0 || configs[g].isActiveRank)
       continue;
     char* helperBuf = static_cast<char*>(recvBuffs[g]);
     size_t chunkSize = chunkSizes[g];
@@ -832,7 +837,7 @@ static ncclResult_t shardedRelayAllReduce2Active(
 
       // The helper's chunk is already reduced, so it lands directly in its
       // final place in recvBuff — no active-side reduction for this region.
-      for (int h = 0; h < cfg.numHelpers; h++) {
+      for (int h = 0; h < cfg.numHelpers && chunkSize > 0; h++) {
         NCCLCHECK(ncclRecv(
             recvbuff + static_cast<size_t>(h) * chunkSize * elementSize,
             chunkSize,
@@ -852,8 +857,13 @@ static ncclResult_t shardedRelayAllReduce2Active(
           comm,
           stream));
       NCCLCHECK(ncclRecv(
-          directDst(chunkSize), dirBSizes[g], datatype, partner, comm, stream));
-    } else {
+          directDst(dirASizes[g]),
+          dirBSizes[g],
+          datatype,
+          partner,
+          comm,
+          stream));
+    } else if (chunkSize > 0) {
       // Helper: hand the reduced chunk to both active ranks.
       const char* helperBuf = static_cast<const char*>(recvBuffs[g]);
       for (int a = 0; a < cfg.nActiveRanks; a++) {