fix position bias in tensor parallel (#1714)

* fix position bias in tensor parallel * add symbol ncclCommFinalize
OpenNMT · May 30, 2024 · 5eb5d5a · 5eb5d5a
1 parent 3b248f1
commit 5eb5d5a
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 5 deletions.
diff --git a/src/cuda/nccl_stub.cc b/src/cuda/nccl_stub.cc
@@ -69,9 +69,9 @@ extern "C" {
   return func(comm);
   }
 
-  ncclResult_t ncclCommAbort(ncclComm_t comm) {
+  ncclResult_t ncclCommFinalize(ncclComm_t comm) {
     using Signature = ncclResult_t(*)(ncclComm_t comm);
-    static auto func = ctranslate2::load_symbol<Signature>("ncclCommAbort");
+    static auto func = ctranslate2::load_symbol<Signature>("ncclCommFinalize");
     return func(comm);
   }
 

diff --git a/src/devices.cc b/src/devices.cc
@@ -196,7 +196,7 @@ namespace ctranslate2 {
     for (auto* comm : _nccl_comms) {
         //finalizing NCCL
         if (*comm) {
-          NCCL_CHECK(ncclCommAbort(*comm));
+          NCCL_CHECK(ncclCommFinalize(*comm));
           NCCL_CHECK(ncclCommDestroy(*comm));
         }
     }

diff --git a/src/layers/attention.cc b/src/layers/attention.cc
@@ -1,5 +1,7 @@
 #include "ctranslate2/layers/attention.h"
 #include "ctranslate2/ops/split.h"
+#include "ctranslate2/utils.h"
+
 
 #include <algorithm>
 #include <cmath>
@@ -210,11 +212,20 @@ namespace ctranslate2 {
                                                  is_decoder,
                                                  with_cache ? key_length - 1 : 0);
         }
+        StorageView* position_bias_per_gpu = position_bias;
+        StorageView position_bias_tmp(position_bias->dtype(), position_bias->device());
+        if (ScopedMPISetter::getCurRank() != 0) {
+          const dim_t num_head_per_gpu = SAFE_DIVIDE(position_bias->dim(0), ScopedMPISetter::getNRanks());
+          ops::Slide slide_ops(0, num_head_per_gpu * ScopedMPISetter::getCurRank(),
+                               num_head_per_gpu, true);
+          slide_ops(*position_bias, position_bias_tmp);
+          position_bias_per_gpu = &position_bias_tmp;
+        }
 
         DEVICE_AND_TYPE_DISPATCH(output.device(), output.dtype(),
-                                 primitives<D>::add_batch_broadcast(position_bias->data<T>(),
+                                 primitives<D>::add_batch_broadcast(position_bias_per_gpu->data<T>(),
                                                                     output.data<T>(),
-                                                                    position_bias->size(),
+                                                                    position_bias_per_gpu->size(),
                                                                     output.size()));
       }