[CRL] Optimize P2P communication by introducing dynamic chunk sizing and improved buffer slicing (flagos-ai#319)

Author: leoda1

flagcx/core/group.cc

@@ -233,9 +233,9 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
                                  .proxyConn.connection;
             op->stream = p2p->stream;
             if (op->connection->transport == TRANSPORT_P2P) {
-              op->args.chunkSize = flagcxP2pChunkSize;
+              op->args.chunkSize = computeP2pChunkSize(p2p->bytes);
               op->args.chunkSteps =
-                  (p2p->bytes + flagcxP2pChunkSize - 1) / (flagcxP2pChunkSize);
+                  (p2p->bytes + op->args.chunkSize - 1) / (op->args.chunkSize);
               op->args.sendStepMask = flagcxP2pChunks - 1;
               setP2pSlotInfo(comm->rank, peer, p2p->bytes, p2p->dtype, 1,
                              &op->args.p2pOpHash, &op->args.p2pSlotIdx);
@@ -312,9 +312,9 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
                                  .proxyConn.connection;
             op->stream = p2p->stream;
             if (op->connection->transport == TRANSPORT_P2P) {
-              op->args.chunkSize = flagcxP2pChunkSize;
+              op->args.chunkSize = computeP2pChunkSize(p2p->bytes);
               op->args.chunkSteps =
-                  (p2p->bytes + flagcxP2pChunkSize - 1) / (flagcxP2pChunkSize);
+                  (p2p->bytes + op->args.chunkSize - 1) / (op->args.chunkSize);
               op->args.sendStepMask = flagcxP2pChunks - 1;
               setP2pSlotInfo(comm->rank, peer, p2p->bytes, p2p->dtype, 0,
                              &op->args.p2pOpHash, &op->args.p2pSlotIdx);

flagcx/core/include/p2p.h

@@ -12,6 +12,7 @@
 extern int64_t flagcxP2pBufferSize;
 extern int64_t flagcxP2pChunkSize;
 extern int64_t flagcxP2pChunks;
+size_t computeP2pChunkSize(size_t nbytes);
 #define FLAGCX_P2P_MAX_STEPS 16
 #define FLAGCX_P2P_MAX_OPS                                                     \
   (FLAGCX_P2P_MAX_STEPS * 2) // Maximum number of concurrent P2P operation pairs

flagcx/core/init.cc

@@ -256,7 +256,7 @@ static flagcxResult_t initTransportsRank(flagcxHeteroComm_t comm,
 FLAGCX_PARAM(P2pBufferSize, "P2P_BUFFER_SIZE",
              64L * 1024 * 1024); // default value to 64MB
 FLAGCX_PARAM(P2pChunkSize, "P2P_CHUNK_SIZE",
-             4L * 1024 * 1024); // default value to 4MB
+             16L * 1024 * 1024); // default value to 16MB
 FLAGCX_PARAM(NetBufferSize, "NET_BUFFER_SIZE",
              64L * 1024 * 1024); // default value to 64MB
 FLAGCX_PARAM(NetChunkSize, "NET_CHUNK_SIZE",

flagcx/core/p2p.cc

@@ -12,6 +12,30 @@ int64_t flagcxP2pBufferSize;
 int64_t flagcxP2pChunkSize;
 int64_t flagcxP2pChunks;
 
+size_t computeP2pChunkSize(size_t nbytes) {
+  size_t dynamicBufferSize = flagcxP2pBufferSize;
+  if (nbytes < (size_t)flagcxP2pBufferSize) {
+    size_t msize = nbytes / (1024 * 1024);
+    int adjustFactor = 0;
+    if (msize >= 32)
+      adjustFactor = 1;
+    else if (msize >= 16)
+      adjustFactor = 2;
+    else if (msize >= 8)
+      adjustFactor = 4;
+    else if (msize >= 4)
+      adjustFactor = 8;
+    else if (msize >= 2)
+      adjustFactor = 16;
+    else if (msize >= 1)
+      adjustFactor = 32;
+    else
+      adjustFactor = 64;
+    dynamicBufferSize = flagcxP2pBufferSize / adjustFactor;
+  }
+  return dynamicBufferSize / flagcxP2pChunks;
+}
+
 struct p2pIpcExpInfo {
   flagcxP2pIpcDesc ipcDesc;
   bool legacyIpcCap;
@@ -145,7 +169,7 @@ flagcxResult_t flagcxP2pProxySend(struct flagcxP2pResources *resources,
         args->subs[step].stepSize =
             std::min(args->chunkSize, size - args->totalCopySize);
         args->subs[step].stepBuff =
-            resources->proxyInfo.recvFifo + (flagcxP2pChunkSize * step);
+            resources->proxyInfo.recvFifo + (args->chunkSize * step);
 
         FLAGCXCHECK(deviceAdaptor->deviceMemcpy(
             args->subs[step].stepBuff, (char *)data + args->totalCopySize,
@@ -266,7 +290,7 @@ flagcxResult_t flagcxP2pProxyRecv(struct flagcxP2pResources *resources,
         args->subs[step].stepSize =
             std::min(args->chunkSize, size - args->totalCopySize);
         args->subs[step].stepBuff =
-            resources->proxyInfo.recvFifo + (flagcxP2pChunkSize * step);
+            resources->proxyInfo.recvFifo + (args->chunkSize * step);
 
         FLAGCXCHECK(deviceAdaptor->deviceMemcpy(
             (char *)data + args->totalCopySize, args->subs[step].stepBuff,