Merge pull request #174 from PerkzZheng/localBatchSize_fix

fix: don't split batch into micro batches when pipeline_para_size = 1
NVIDIA · Dec 20, 2021 · 8a1c109 · 8a1c109
2 parents 0311bc6 + 0df48e9
commit 8a1c109
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/fastertransformer/utils/nccl_utils.cc b/src/fastertransformer/utils/nccl_utils.cc
@@ -170,6 +170,9 @@ template void ftNcclAllGather(
 size_t getLocalBatchSize(const size_t batch_size, const size_t seq_len, const size_t pipeline_para_size)
 {
     size_t local_batch_size = batch_size;
+    if (pipeline_para_size == 1) {
+        return local_batch_size;
+    }
     if (local_batch_size % pipeline_para_size == 0) {
         local_batch_size /= pipeline_para_size;
     }