fix: AllReduce minor

flexflow · Nov 23, 2024 · 1f6dab4 · 1f6dab4
1 parent 1c2875f
commit 1f6dab4
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/parallel_ops/kernels/allreduce_kernels.cu b/src/parallel_ops/kernels/allreduce_kernels.cu
@@ -142,6 +142,9 @@ void inference_kernel_wrapper(Context ctx,
   int device_id = m->handle.device_id;
   ncclComm_t ncclComm = m->handle.ncclComm;
   DataType dtype = input.data_type;
+  if (num_elements == 0) {
+    return;
+  }
 
   tensorrt_llm::AllReduceStrategyType strategy =
       tensorrt_llm::SelectImplementation(