Added new unit tests for AllReduce with Bias API

atulkulk · atulkulk · commit 159cb6fe5bda · 2025-11-12T21:17:33.000-06:00
diff --git a/test/AllReduceTests.cpp b/test/AllReduceTests.cpp
diff --git a/test/common/CallCollectiveForked.cpp b/test/common/CallCollectiveForked.cpp
@@ -40,22 +40,18 @@ void callCollective(ncclUniqueId id, int collID, int rank, int nranks, const std
         default:
             ERROR("This collective is not implemented for callCollective routine");
     }
-    
+
     HIPCALL(hipSetDevice(rank));
     hipStream_t stream;
     HIPCALL(hipStreamCreate(&stream));
     ncclComm_t comm;
-    
-    
 
     NCCLCHECK(ncclCommInitRank(&comm, nranks, id, rank));
     int *sendbuff;
     int *recvbuff;
     void *sendRegHandle;
     void *recvRegHandle;
-    
 
-    
     size_t sendSize = 0;
     size_t recvSize = 0;
 
@@ -78,13 +74,12 @@ void callCollective(ncclUniqueId id, int collID, int rank, int nranks, const std
     else{
       HIPCALL(hipMallocManaged((void **)&sendbuff, sendSize * sizeof(int)));
       HIPCALL(hipMallocManaged((void **)&recvbuff, recvSize * sizeof(int)));
-    }    
-   
+    }
+
     NCCLCHECK(ncclCommRegister(comm, sendbuff, sendSize * sizeof(int), &sendRegHandle));
     NCCLCHECK(ncclCommRegister(comm, recvbuff, recvSize * sizeof(int), &recvRegHandle));
 
     HIPCALL(hipMemcpy(sendbuff, send.data(), sizeof(int) * sendSize, hipMemcpyHostToDevice));
-    HIPCALL(hipMemcpy(recvbuff, recv.data(), sizeof(int) *recvSize, hipMemcpyHostToDevice));
 
     switch(collID){
       case ncclCollAllReduce:
@@ -98,12 +93,13 @@ void callCollective(ncclUniqueId id, int collID, int rank, int nranks, const std
 
     HIPCALL(hipStreamSynchronize(stream));
     HIPCALL(hipMemcpy(recv.data(), recvbuff, sizeof(int) * recvSize, hipMemcpyDeviceToHost));
-    
+
     NCCLCHECK(ncclCommDeregister(comm, sendRegHandle));
     NCCLCHECK(ncclCommDeregister(comm, recvRegHandle));
 
     HIPCALL(hipFree(sendbuff));
     HIPCALL(hipFree(recvbuff));
+
     ncclCommDestroy(comm);
   }
 
@@ -115,7 +111,7 @@ void callCollectiveForked(int nranks,  int collID, const std::vector<int>& sendB
     for(int r = 0; r < nranks; ++r){
       if(pipe(childPipes[r].data()) == -1)
         ERROR("child %i pipe Failed\n", r);
-    } 
+    }
 
     auto createNCCLid = [&](int rank){
         ncclGetUniqueId(&id);
@@ -162,9 +158,10 @@ void callCollectiveForked(int nranks,  int collID, const std::vector<int>& sendB
     }
 
     getAndDistributeNCCLid(nranks);
-    
+
     for(int r = 0; r < nranks; ++r)
       wait(NULL); // Wait for all children
 }
 
 }
+
diff --git a/test/common/CallCollectiveForked.hpp b/test/common/CallCollectiveForked.hpp
@@ -15,3 +15,5 @@ namespace RcclUnitTesting
 }
 
 #endif
+
+
diff --git a/test/common/CollectiveArgs.cpp b/test/common/CollectiveArgs.cpp
@@ -39,6 +39,11 @@ namespace RcclUnitTesting
     this->streamIdx         = streamIdx;
     this->options           = optionalColArgs;
 
+    // Initialize bias fields
+    this->numBiasElements = 0;
+    this->numBiasBytesAllocated = 0;
+    this->biasRegHandle = nullptr;
+
     if (this->options.scalarMode != -1)
     {
       size_t const numBytes = DataTypeToBytes(dataType);
@@ -102,6 +107,23 @@ namespace RcclUnitTesting
       CHECK_CALL(this->expected.AllocateCpuMem(this->numOutputBytesAllocated));
     }
     CHECK_CALL(this->outputCpu.AllocateCpuMem(this->numOutputBytesAllocated));
+
+    // Allocate bias buffers if bias is enabled
+    if (this->options.useBias)
+    {
+      this->numBiasElements = this->options.biasNumElements;
+      this->numBiasBytesAllocated = this->numBiasElements * DataTypeToBytes(this->dataType);
+      CHECK_CALL(this->biasGpu.AllocateGpuMem(this->numBiasBytesAllocated, useManagedMem, userRegistered));
+      CHECK_CALL(this->biasCpu.AllocateCpuMem(this->numBiasBytesAllocated));
+      this->biasRegHandle = nullptr;
+    }
+    else
+    {
+      this->numBiasElements = 0;
+      this->numBiasBytesAllocated = 0;
+      this->biasRegHandle = nullptr;
+    }
+
     return TEST_SUCCESS;
   }
 
@@ -155,6 +177,15 @@ namespace RcclUnitTesting
       if (this->options.scalarMode == 1) CHECK_HIP(hipHostFree(this->localScalar.ptr));
       this->localScalar.Attach(nullptr);
     }
+
+    // Deallocate bias buffers if they were allocated
+    if (this->options.useBias && this->numBiasBytesAllocated > 0)
+    {
+      this->biasGpu.FreeGpuMem(this->userRegistered);
+      this->biasCpu.FreeCpuMem();
+      this->biasRegHandle = nullptr;
+    }
+
     return TEST_SUCCESS;
   }
 
diff --git a/test/common/CollectiveArgs.hpp b/test/common/CollectiveArgs.hpp
@@ -82,6 +82,15 @@ namespace RcclUnitTesting
     ScalarTransport scalarTransport;        // Used for custom reduction operators
     int             scalarMode = -1;        // -1 if scalar not used
 
+    // Bias support for fused AllReduce+Bias operations
+    bool            useBias = false;        // Enable bias addition
+    void*           biasPtr = nullptr;      // Pointer to bias buffer (GPU memory)
+    size_t          biasNumElements = 0;    // Number of elements in bias buffer
+    int             biasConstantValue = -1; // If >= 0, use constant value for all bias elements (instead of incremental pattern)
+
+    // Input data pattern control (useful for ncclProd to avoid overflow at high rank counts)
+    int             inputConstantValue = -1; // If >= 0, use constant value for all input elements (instead of rank-based pattern)
+
     // allToAllv args
     size_t          sendcounts[MAX_RANKS*MAX_RANKS];
     size_t          sdispls[MAX_RANKS*MAX_RANKS];
@@ -122,6 +131,13 @@ namespace RcclUnitTesting
     size_t         numInputElementsAllocated;
     size_t         numOutputElementsAllocated;
 
+    // Bias data for fused AllReduce+Bias operations
+    PtrUnion       biasGpu;           // Bias buffer on GPU
+    PtrUnion       biasCpu;           // Bias buffer on CPU (for initialization/validation)
+    void*          biasRegHandle;     // Handle for registered bias buffer
+    size_t         numBiasElements;   // Number of elements in bias buffer
+    size_t         numBiasBytesAllocated;  // Number of bytes allocated for bias
+
     // Set collective arguments
     ErrCode SetArgs(int             const globalRank,
                     int             const totalRanks,
diff --git a/test/common/PrepDataFuncs.cpp b/test/common/PrepDataFuncs.cpp
@@ -47,6 +47,22 @@ namespace RcclUnitTesting
             collArgs.numOutputElements, collArgs.numOutputElementsAllocated);
       return TEST_FAIL;
     }
+
+    // Check bias allocation if bias is enabled
+    if (collArgs.options.useBias)
+    {
+      if (collArgs.numBiasElements == 0 || collArgs.numBiasBytesAllocated == 0)
+      {
+        ERROR("Bias is enabled but bias buffers are not allocated\n");
+        return TEST_FAIL;
+      }
+      if (collArgs.numBiasElements != collArgs.numOutputElements)
+      {
+        ERROR("Number of bias elements (%lu) must match number of output elements (%lu)\n",
+              collArgs.numBiasElements, collArgs.numOutputElements);
+        return TEST_FAIL;
+      }
+    }
     return TEST_SUCCESS;
   }
 
@@ -108,7 +124,22 @@ namespace RcclUnitTesting
     for (int rank = 0; rank < collArgs.totalRanks; ++rank)
     {
       // Generate temporary input for this rank
-      CHECK_CALL(tempInputCpu.FillPattern(collArgs.dataType, collArgs.numInputElements, rank, false));
+      if (collArgs.options.inputConstantValue >= 0)
+      {
+        // Use constant value for all input elements across all ranks
+        // This is useful for ncclProd at high rank counts to avoid factorial overflow
+        for (size_t i = 0; i < collArgs.numInputElements; i++)
+        {
+          CHECK_CALL(tempInputCpu.Set(collArgs.dataType, i,
+                                      collArgs.options.inputConstantValue,
+                                      (double)collArgs.options.inputConstantValue));
+        }
+      }
+      else
+      {
+        // Use rank-based pattern: value[rank][i] = (rank + i) % 256 (default behavior)
+        CHECK_CALL(tempInputCpu.FillPattern(collArgs.dataType, collArgs.numInputElements, rank, false));
+      }
 
       // Copy the pre-scaled input into GPU memory for the correct rank
       if (rank == collArgs.globalRank)
@@ -144,6 +175,38 @@ namespace RcclUnitTesting
     {
       CHECK_CALL(result.DivideByInt(collArgs.dataType, collArgs.numInputElements, collArgs.totalRanks));
     }
+
+    // Add bias to expected output if bias is enabled
+    if (collArgs.options.useBias && (isAllReduce || collArgs.options.root == collArgs.globalRank))
+    {
+      // Initialize bias data on CPU
+      if (collArgs.options.biasConstantValue >= 0)
+      {
+        // Use constant value for all bias elements (useful for ncclProd to avoid overflow)
+        for (size_t i = 0; i < collArgs.numBiasElements; i++)
+        {
+          CHECK_CALL(collArgs.biasCpu.Set(collArgs.dataType, i,
+                                          collArgs.options.biasConstantValue,
+                                          (double)collArgs.options.biasConstantValue));
+        }
+      }
+      else
+      {
+        // Use incremental pattern: bias[i] = i (default behavior)
+        CHECK_CALL(collArgs.biasCpu.FillPattern(collArgs.dataType, collArgs.numBiasElements, 0, false));
+      }
+
+      // Copy bias data to GPU
+      size_t const biasBytes = collArgs.numBiasBytesAllocated;
+      CHECK_HIP(hipMemcpy(collArgs.biasGpu.ptr, collArgs.biasCpu.ptr, biasBytes, hipMemcpyHostToDevice));
+
+      // Apply bias to expected output using the SAME reduction operation as AllReduce
+      CHECK_CALL(result.Reduce(collArgs.dataType, collArgs.numInputElements, collArgs.biasCpu, tempOp));
+
+      // Update the biasPtr in options to point to the GPU buffer
+      collArgs.options.biasPtr = collArgs.biasGpu.ptr;
+    }
+
     return TEST_SUCCESS;
   }
 
diff --git a/test/common/TestBed.cpp b/test/common/TestBed.cpp
@@ -636,12 +636,16 @@ namespace RcclUnitTesting
                                std::vector<bool>           const& inPlaceList,
                                std::vector<bool>           const& managedMemList,
                                std::vector<bool>           const& useHipGraphList,
-                               bool                        const& enableSweep)
+                               bool                        const& enableSweep,
+                               OptionalColArgs*            const  optionalArgsPtr)
   {
     // Sort numElements in descending order to cut down on # of allocations
     std::vector<int> sortedN = numElements;
     std::sort(sortedN.rbegin(), sortedN.rend());
-    OptionalColArgs optionalArgs;
+
+    // Use provided OptionalColArgs or create default one
+    OptionalColArgs defaultArgs;
+    OptionalColArgs& optionalArgs = (optionalArgsPtr != nullptr) ? *optionalArgsPtr : defaultArgs;
     // Filter out any unsupported datatypes, in case only subset has been compiled for
     std::vector<ncclDataType_t> const& supportedDataTypes = this->GetAllSupportedDataTypes();
     std::vector<ncclDataType_t> dataTypes;
@@ -718,6 +722,11 @@ namespace RcclUnitTesting
                                                     &numOutputElements);
           optionalArgs.redOp = redOps[rdIdx];
           optionalArgs.root = roots[rtIdx] % this->numActiveRanks;
+          // Set biasNumElements if bias is enabled
+          if (optionalArgs.useBias)
+          {
+            optionalArgs.biasNumElements = numOutputElements;
+          }
           this->SetCollectiveArgs(funcTypes[ftIdx],
                                   dataTypes[dtIdx],
                                   numInputElements,
diff --git a/test/common/TestBed.hpp b/test/common/TestBed.hpp
@@ -41,7 +41,7 @@ namespace RcclUnitTesting
                    std::vector<int>              const& numStreamsPerGroup,
                    int                           const  numGroupCalls = 1,
                    bool                          const  useBlocking   = true);
- 
+
     // Prepare TestBed for use with GPUs across multiple child processes
     void InitComms(std::vector<std::vector<int>> const& deviceIdsPerChild,
                    int  const numCollectivesInGroup = 1,
@@ -96,7 +96,7 @@ namespace RcclUnitTesting
     // Execute all collectives on all test children
     // Blocks until collective is completed
     void ExecuteCollectives(std::vector<int> const &currentRanks = {},
-                            int              const groupId       = -1, 
+                            int              const groupId       = -1,
                             bool             const useHipGraph   = false);
 
     // Perform results validation - compare output to expected
@@ -140,7 +140,7 @@ namespace RcclUnitTesting
                                                           int const numGpus,
                                                           int const ranksPerGpu,
                                                           const std::vector<int>& gpuPriorityOrder);
-                                                          
+
     static std::vector<std::vector<int>> GetDeviceIdsList(int const numProcesses,
                                                           int const numGpus,
                                                           const std::vector<int>& gpuPriorityOrder);
@@ -166,7 +166,8 @@ namespace RcclUnitTesting
                         std::vector<bool>           const& inPlaceList,
                         std::vector<bool>           const& managedMemList,
                         std::vector<bool>           const& useHipGraphList,
-                        bool                        const& enableSweep = true);
+                        bool                        const& enableSweep = true,
+                        OptionalColArgs*            const  optionalArgs = nullptr);
 
     // Wait for user-input if in interactive mode
     void InteractiveWait(std::string message);
diff --git a/test/common/TestBedChild.cpp b/test/common/TestBedChild.cpp

Original file line number	Diff line number	Diff line change
`@@ -15,3 +15,5 @@ namespace RcclUnitTesting`
`15`	`15`	`}`
`16`	`16`
`17`	`17`	`#endif`
	`18`	`+`
	`19`	`+`