Repository: microsoft/DeepSpeed

Issue:
[BUG] FP16 used for all reduce even if BFLOAT16 is enabled
Even if BFLOAT16 is enabled, unless the communication datatype is also explicitly set to bfp16, fp16 is still used: https://github.com/microsoft/DeepSpeed/blob/9b70ce56e7af89d5226f9b06ebe1137407f371dc/deepspeed/runtime/engine.py#L702-L711

Moreover in zero, even if the communication datatype is set to bfp16 at some place the default value is fp16
https://github.com/microsoft/DeepSpeed/blob/9b70ce56e7af89d5226f9b06ebe1137407f371dc/deepspeed/runtime/zero/stage_1_and_2.py#L1334-L1338




Discussion:

Resolving patch:
diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py
--- a/deepspeed/runtime/zero/stage3.py
+++ b/deepspeed/runtime/zero/stage3.py
@@ -1384,11 +1384,7 @@ def set_none_gradients_to_zero(self, i, partition_id):
 
     ######################Reduction Related Methods##############################
 
-    def allreduce_bucket(self,
-                         bucket,
-                         communication_data_type=torch.float16,
-                         rank=None,
-                         log=None):
+    def allreduce_bucket(self, bucket, rank=None, log=None):
         rank = None
         tensor = self.flatten(bucket)
 
@@ -1396,6 +1392,8 @@ def allreduce_bucket(self,
 
         if pg_correctness_test:
             communication_data_type = torch.float32
+        else:
+            communication_data_type = self.communication_data_type
 
         if communication_data_type != tensor.dtype:
             tensor_to_allreduce = tensor.to(communication_data_type)
diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py
--- a/deepspeed/runtime/zero/stage_1_and_2.py
+++ b/deepspeed/runtime/zero/stage_1_and_2.py
@@ -966,9 +966,13 @@ def average_tensor(self, tensor):
             if not self.ipg_bucket_has_moe_params:
                 tensor.div_(dist.get_world_size(group=self.dp_process_group))
 
+            tensor_to_reduce = tensor
+            if self.communication_data_type != tensor.dtype:
+                tensor_to_reduce = tensor.to(self.communication_data_type)
+
             async_handles = []
             for i, (dst, bucket_offset, numel) in enumerate(rank_and_offsets):
-                grad_slice = tensor.narrow(0, int(bucket_offset), int(numel))
+                grad_slice = tensor_to_reduce.narrow(0, int(bucket_offset), int(numel))
                 # if dist.get_rank() == 0:
                 #     print(f"Rank {dist.get_rank()} rank offset id {i} real dp size {dist.get_world_size(group=real_dp_process_group[i])} and dst: {dst}")
                 # dist.barrier()
@@ -983,6 +987,9 @@ def average_tensor(self, tensor):
             for handle in async_handles:
                 handle.wait()
 
+            if self.communication_data_type != tensor.dtype:
+                tensor.copy_(tensor_to_reduce)
+
     ##############################################################################
     ############################# CPU Offload Methods#############################
     ##############################################################################
@@ -1337,11 +1344,7 @@ def set_none_gradients_to_zero(self, i, partition_id):
                 param.grad = torch.zero_like(param)
 
     ######################Reduction Related Methods##############################
-    def allreduce_bucket(self,
-                         bucket,
-                         communication_data_type=torch.float16,
-                         rank=None,
-                         log=None):
+    def allreduce_bucket(self, bucket, rank=None, log=None):
         rank = None
         tensor = self.flatten(bucket)
 
@@ -1349,6 +1352,8 @@ def allreduce_bucket(self,
 
         if pg_correctness_test:
             communication_data_type = torch.float32
+        else:
+            communication_data_type = self.communication_data_type
 
         if communication_data_type != tensor.dtype:
             tensor_to_allreduce = tensor.to(communication_data_type)


Test patch:
diff --git a/tests/unit/test_bf16.py b/tests/unit/test_bf16.py
--- a/tests/unit/test_bf16.py
+++ b/tests/unit/test_bf16.py
@@ -6,6 +6,7 @@
 from deepspeed.ops.op_builder import CPUAdamBuilder
 from .simple_model import SimpleModel, SimpleOptimizer, random_dataloader, args_from_dict
 from .util import bf16_required_version_check
+from deepspeed import comm as dist
 
 
 @pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)])
@@ -318,3 +319,78 @@ def _go(args, model, hidden_dim):
             model.step()
 
     _go(args=args, model=model, hidden_dim=hidden_dim)
+
+
+@pytest.mark.parametrize('comp_type_str, comm_type_str',
+                         [("fp16",
+                           "fp16"),
+                          ("bfp16",
+                           "bfp16"),
+                          ("fp16",
+                           "bfp16"),
+                          ("bfp16",
+                           "fp16"),
+                          ("fp32",
+                           "fp16"),
+                          ("fp32",
+                           "bfp16")])
+def test_zero_dtype_cocktail(tmpdir, comp_type_str, comm_type_str):
+    torch_dtype_dict = {
+        "fp16": torch.float16,
+        "bfp16": torch.bfloat16,
+        "fp32": torch.float
+    }
+    comp_torch_dtype = torch_dtype_dict[comp_type_str]
+    comm_torch_dtype = torch_dtype_dict[comm_type_str]
+
+    if comp_torch_dtype == torch.bfloat16 or comm_torch_dtype == torch.bfloat16:
+        if not bf16_required_version_check():
+            pytest.skip(
+                " DeepSpeed BFloat16 tests need torch >= 1.10, NCCL >= 2.10.3, CUDA > =11.0 and HW support for BFloat16 to run correctly"
+            )
+
+    config_dict = {
+        "train_batch_size": 2,
+        "steps_per_print": 1,
+        "fp16": {
+            "enabled": torch_dtype_dict[comp_type_str] == torch.float16
+        },
+        "bf16": {
+            "enabled": torch_dtype_dict[comp_type_str] == torch.bfloat16
+        },
+        "zero_optimization": {
+            "stage": 2
+        },
+        "communication_data_type": comm_type_str
+    }
+
+    args = args_from_dict(tmpdir, config_dict)
+    hidden_dim = 10
+
+    model = SimpleModel(hidden_dim)
+
+    orig_torch_reduce = dist.reduce
+
+    def custom_reduce(tensor, dst, op=dist.ReduceOp.SUM, group=None, async_op=False):
+        assert tensor.dtype == comm_torch_dtype
+        return orig_torch_reduce(tensor, dst, op, group, async_op)
+
+    @distributed_test(world_size=[2])
+    def _go(args, model, hidden_dim, comp_dtype):
+        optimizer = torch.optim.Adam(model.parameters())
+        model, _, _, _ = deepspeed.initialize(args=args,
+                                              model=model,
+                                              optimizer=optimizer)
+        data_loader = random_dataloader(model=model,
+                                        total_samples=2,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=comp_dtype)
+        dist.reduce = custom_reduce
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+        dist.reduce = orig_torch_reduce
+
+    _go(args=args, model=model, hidden_dim=hidden_dim, comp_dtype=comp_torch_dtype)