Fixing the bug in allreduce1 (#220)

saeedmaleki · web-flow · commit 70eb6d73287c · 2023-11-18T10:34:52.000-08:00
diff --git a/python/benchmark/allreduce.cu b/python/benchmark/allreduce.cu
@@ -176,10 +176,12 @@ __device__ void allreduce1_helper(mscclpp::SmChannelDeviceHandle* smChans, TYPE*
       TYPE val = smChans[peerIdx].read<TYPE>(idx);
       tmp += val;
     }
-    for (int index = 0; index < nPeer; ++index) {
-      int peerIdx = (index + rank);
-      if (peerIdx >= nPeer) peerIdx -= nPeer;
-      smChans[peerIdx].write<TYPE>(idx, tmp);
+    if (READ_ONLY == 0) {
+      for (int index = 0; index < nPeer; ++index) {
+        int peerIdx = (index + rank);
+        if (peerIdx >= nPeer) peerIdx -= nPeer;
+        smChans[peerIdx].write<TYPE>(idx, tmp);
+      }
     }
     buff[idx] = tmp;
   }
@@ -198,6 +200,7 @@ __device__ void allreduce1_helper(mscclpp::SmChannelDeviceHandle* smChans, TYPE*
   }
 
   if (READ_ONLY) {
+    deviceSyncer.sync(gridDim.x);
     for (int i = 0; i < nPeer; ++i) {
       int peerIdx = (i + rank);
       if (peerIdx >= nPeer) peerIdx -= nPeer;

Original file line number	Diff line number	Diff line change
`@@ -176,10 +176,12 @@ __device__ void allreduce1_helper(mscclpp::SmChannelDeviceHandle* smChans, TYPE*`
`176`	`176`	`TYPE val = smChans[peerIdx].read<TYPE>(idx);`
`177`	`177`	`tmp += val;`
`178`	`178`	`}`
`179`		`- for (int index = 0; index < nPeer; ++index) {`
`180`		`- int peerIdx = (index + rank);`
`181`		`- if (peerIdx >= nPeer) peerIdx -= nPeer;`
`182`		`- smChans[peerIdx].write<TYPE>(idx, tmp);`
	`179`	`+ if (READ_ONLY == 0) {`
	`180`	`+ for (int index = 0; index < nPeer; ++index) {`
	`181`	`+ int peerIdx = (index + rank);`
	`182`	`+ if (peerIdx >= nPeer) peerIdx -= nPeer;`
	`183`	`+ smChans[peerIdx].write<TYPE>(idx, tmp);`
	`184`	`+ }`
`183`	`185`	`}`
`184`	`186`	`buff[idx] = tmp;`
`185`	`187`	`}`
`@@ -198,6 +200,7 @@ __device__ void allreduce1_helper(mscclpp::SmChannelDeviceHandle* smChans, TYPE*`
`198`	`200`	`}`
`199`	`201`
`200`	`202`	`if (READ_ONLY) {`
	`203`	`+ deviceSyncer.sync(gridDim.x);`
`201`	`204`	`for (int i = 0; i < nPeer; ++i) {`
`202`	`205`	`int peerIdx = (i + rank);`
`203`	`206`	`if (peerIdx >= nPeer) peerIdx -= nPeer;`