Fix for stage3 when setting different communication data type (#4540)

Co-authored-by: Olatunji Ruwase <[email protected]>
deepspeedai · Dec 11, 2023 · b186816 · b186816
1 parent 2ce6bf8
commit b186816
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py
@@ -1215,7 +1215,7 @@ def __reduce_and_partition_ipg_grads(self, safe_mode: bool = False) -> None:
     @instrument_w_nvtx
     def __avg_scatter_contiguous_grads(self, buffer_to_reduce: Tensor) -> List[Tensor]:
         dtype = buffer_to_reduce.dtype
-        if self.communication_data_type == self.dtype:
+        if self.communication_data_type != dtype:
             buffer_to_reduce = buffer_to_reduce.to(self.communication_data_type)
         if self.postscale_gradients and self.gradient_predivide_factor != 1.0:
             buffer_to_reduce = buffer_to_reduce.div_(self.gradient_predivide_factor)