reduce transfer size per token.

b4rtaz · Nov 21, 2024 · c2c2fa0 · c2c2fa0
1 parent dffc23d
commit c2c2fa0
Show file tree

Hide file tree

Showing 5 changed files with 29 additions and 23 deletions.
diff --git a/src/app.cpp b/src/app.cpp
@@ -113,7 +113,6 @@ AppArgs AppArgs::parse(int argc, char** argv, bool hasMode) {
 
 TransformerArch TransformerArchFactory::create(TransformerSpec* spec) {
     if (spec->archType == LLAMA) return buildLlamaArch(spec);
-    if (spec->archType == MIXTRAL) return buildMixtralArch(spec);
     printf("Unsupported arch type: %d\n", spec->archType);
     exit(EXIT_FAILURE);
 }

diff --git a/src/llama2-tasks.cpp b/src/llama2-tasks.cpp
@@ -116,7 +116,7 @@ void llamaQuantizeAtt(TASK_ARGS) {
 
 void llamaSyncAtt(TASK_ARGS) {
     TASK_VARIABLES;
-    syncSliceOfSlicedBuffer(nThreads, threadIndex, ctx, TB_SLICED_XBV_QUANTIZED);
+    syncSliceOfSlicedBuffer(nThreads, threadIndex, false, ctx, TB_SLICED_XBV_QUANTIZED);
 }
 
 void llamaDequantizeAtt(TASK_ARGS) {
@@ -195,7 +195,7 @@ void llamaQuantizeFfn2(TASK_ARGS) {
 
 void llamaSyncFfn2(TASK_ARGS) {
     TASK_VARIABLES;
-    syncSliceOfSlicedBuffer(nThreads, threadIndex, ctx, TB_SLICED_XBV_QUANTIZED);
+    syncSliceOfSlicedBuffer(nThreads, threadIndex, false, ctx, TB_SLICED_XBV_QUANTIZED);
 }
 
 void llamaDequantizeFfn2(TASK_ARGS) {
@@ -248,7 +248,7 @@ void llamaFinalize(TASK_ARGS) {
 
 void llamaSyncLogits(TASK_ARGS) {
     TASK_VARIABLES;
-    syncSliceOfSlicedBuffer(nThreads, threadIndex, ctx, TB_SLICED_LOGITS);
+    syncSliceOfSlicedBuffer(nThreads, threadIndex, true, ctx, TB_SLICED_LOGITS);
 }
 
 TransformerArch buildLlamaArch(TransformerSpec* spec) {

diff --git a/src/mixtral-tasks.cpp b/src/mixtral-tasks.cpp
@@ -121,7 +121,7 @@ void grokQuantizeMoeMul(TASK_ARGS) {
 
 void grokSyncMoeMulA(TASK_ARGS) {
     TASK_VARIABLES;
-    syncSliceOfSlicedBuffer(nThreads, threadIndex, ctx, TB_SLICED_HB_QUANTIZED);
+    syncSliceOfSlicedBuffer(nThreads, threadIndex, false, ctx, TB_SLICED_HB_QUANTIZED);
 }
 
 void grokSyncMoeMulRearrange(TASK_ARGS) {
@@ -155,7 +155,7 @@ void grokSyncMoeMulB(TASK_ARGS) {
 
 void grokSyncMoeOutput(TASK_ARGS) {
     TASK_VARIABLES;
-    syncSliceOfSlicedBuffer(nThreads, threadIndex, ctx, TB_SLICED_XB2_QUANTIZED);
+    syncSliceOfSlicedBuffer(nThreads, threadIndex, false, ctx, TB_SLICED_XB2_QUANTIZED);
 }
 
 void grokDequantizeMoeOutput(TASK_ARGS) {

diff --git a/src/tasks.cpp b/src/tasks.cpp
@@ -71,29 +71,36 @@ void syncUnitBuffer(unsigned int nThreads, unsigned int threadIndex, Transformer
     }
 }
 
-void syncSliceOfSlicedBuffer(unsigned int nThreads, unsigned int threadIndex, TransformerContext* ctx, uint8_t bufferIndex) {
+void syncSliceOfSlicedBuffer(unsigned int nThreads, unsigned int threadIndex, bool onlyFromWorkerToRoot, TransformerContext* ctx, uint8_t bufferIndex) {
     unsigned int nSocketsPerThread = ctx->socketPool->nSockets / nThreads + (ctx->socketPool->nSockets % nThreads > threadIndex ? 1 : 0);
     if (nSocketsPerThread == 0) return;
-    size_t sliceBytes = ctx->transformer->buffer->getSlicedBytes(bufferIndex);
 
-    void* mySliceData = ctx->transformer->buffer->getSliced(bufferIndex, ctx->transformer->sliceIndex);
+    size_t sliceBytes = ctx->transformer->buffer->getSlicedBytes(bufferIndex);
 
-    SocketIo writeIos[nSocketsPerThread];
-    SocketIo readIos[nSocketsPerThread];
-    for (unsigned int i = 0; i < nSocketsPerThread; i++) {
-        unsigned int socketIndex = threadIndex + i * nThreads;
-        writeIos[i].socketIndex = socketIndex;
-        writeIos[i].data = mySliceData;
-        writeIos[i].size = sliceBytes;
+    if (!onlyFromWorkerToRoot || ctx->transformer->sliceIndex != 0) {
+        void* mySliceData = ctx->transformer->buffer->getSliced(bufferIndex, ctx->transformer->sliceIndex);
 
-        int sliceIndex = socketIndex >= ctx->transformer->sliceIndex ? socketIndex + 1 : socketIndex;
-        readIos[i].socketIndex = socketIndex;
-        readIos[i].data = ctx->transformer->buffer->getSliced(bufferIndex, sliceIndex);
-        readIos[i].size = sliceBytes;
+        SocketIo writeIos[nSocketsPerThread];
+        for (unsigned int i = 0; i < nSocketsPerThread; i++) {
+            unsigned int socketIndex = threadIndex + i * nThreads;
+            writeIos[i].socketIndex = socketIndex;
+            writeIos[i].data = mySliceData;
+            writeIos[i].size = sliceBytes;
+        }
+        ctx->socketPool->writeManyWithAlignment(nSocketsPerThread, writeIos);
     }
 
-    ctx->socketPool->writeManyWithAlignment(nSocketsPerThread, writeIos);
-    ctx->socketPool->readManyWithAlignment(nSocketsPerThread, readIos);
+    if (!onlyFromWorkerToRoot || ctx->transformer->sliceIndex == 0) {
+        SocketIo readIos[nSocketsPerThread];
+        for (unsigned int i = 0; i < nSocketsPerThread; i++) {
+            unsigned int socketIndex = threadIndex + i * nThreads;
+            int sliceIndex = socketIndex >= ctx->transformer->sliceIndex ? socketIndex + 1 : socketIndex;
+            readIos[i].socketIndex = socketIndex;
+            readIos[i].data = ctx->transformer->buffer->getSliced(bufferIndex, sliceIndex);
+            readIos[i].size = sliceBytes;
+        }
+        ctx->socketPool->readManyWithAlignment(nSocketsPerThread, readIos);
+    }
 }
 
 void quantizeUnitBuffer(unsigned int nThreads, unsigned int threadIndex, TransformerContext* ctx, uint8_t sourceBufferIndex, uint8_t targetBufferIndex) {

diff --git a/src/tasks.hpp b/src/tasks.hpp
@@ -42,7 +42,7 @@ class TransformerArch {
     TransformerSpec* spec = transformer->spec; // printf("%s:%d\n", __FUNCTION__, ctx->currentBlockIndex); fflush(stdout);
 
 void syncUnitBuffer(unsigned int nThreads, unsigned int threadIndex, TransformerContext* ctx, uint8_t bufferIndex);
-void syncSliceOfSlicedBuffer(unsigned int nThreads, unsigned int threadIndex, TransformerContext* ctx, uint8_t bufferIndex);
+void syncSliceOfSlicedBuffer(unsigned int nThreads, unsigned int threadIndex, bool onlyFromWorkerToRoot, TransformerContext* ctx, uint8_t bufferIndex);
 void quantizeUnitBuffer(unsigned int nThreads, unsigned int threadIndex, TransformerContext* ctx, uint8_t sourceBufferIndex, uint8_t targetBufferIndex);
 void quantizeSlicedBuffer(unsigned int nThreads, unsigned int threadIndex, TransformerContext* ctx, uint8_t sourceBufferIndex, uint8_t targetBufferIndex);
 void dequantizeSlicedBuffer(unsigned int nThreads, unsigned int threadIndex, TransformerContext* ctx, bool skipMySlice, uint8_t sourceBufferIndex, uint8_t targetBufferIndex);