ZeRO2-Offload: Load balance gradient copying to CPU
diff --git a/deepspeed/runtime/zero/stage2.py b/deepspeed/runtime/zero/stage2.py
index e41d848a8693..188c1a0bbe19 100755
--- a/deepspeed/runtime/zero/stage2.py
+++ b/deepspeed/runtime/zero/stage2.py
@@ -191,12 +191,17 @@ def __init__(self,
         self.all_reduce_print = False
         self.dtype = self.optimizer.param_groups[0]['params'][0].dtype
 
+        self.round_robin_fp16_groups = []
+        self.round_robin_fp6_indices = []
+
         # padding on each partition for alignment purposes
         self.groups_padding = []
         # loop to deal with groups
         for i, param_group in enumerate(self.optimizer.param_groups):
             # push this group to list before modify
+            # TODO: Explore simplification that avoids the extra book-keeping by pushing the reordered group
             self.fp16_groups.append(param_group['params'])
+
             # Record padding required to align group to world size
             if partition_id == dist.get_world_size(group=self.dp_process_group) - 1:
                 padding = get_alignment_padding(self.fp16_groups[i],
@@ -213,10 +218,21 @@ def __init__(self,
             move_to_cpu(self.fp16_groups[i])
             see_memory_usage(f"After moving param group {i} to CPU")
 
+            # Reorder group parameters for load balancing of gradient partitioning during backward among ranks.
+            # This ensures that gradients are reduced in a fashion such that ownership round robins among the ranks.
+            # For example, rather than 3 gradients (g_n+2, g_n+1, g_n) that are reduced consecutively belonging
+            # to the same rank, instead they will belong to 3 ranks (r_m+2, r_m+1, r_m).
+            round_robin_tensors, round_robin_indices = self._round_robin_reorder(
+                self.fp16_groups[i],
+                dist.get_world_size(group=self.dp_process_group)
+            )
+            self.round_robin_fp16_groups.append(round_robin_tensors)
+            self.round_robin_fp6_indices.append(round_robin_indices)
+
             #create flat buffer in CPU and move to GPU
             self.fp16_groups_flat.append(
                 self.flatten_dense_tensors_aligned(
-                    self.fp16_groups[i],
+                    self.round_robin_fp16_groups[i],
                     dist.get_world_size(group=self.dp_process_group)).cuda(
                         torch.cuda.current_device()))
             see_memory_usage(f"After flattening and moving param group {i} to GPU")
@@ -226,10 +242,7 @@ def __init__(self,
                     f"After Flattening and after emptying param group {i} cache")
 
             # set model fp16 weight to slices of flattened buffer
-            updated_params = self.unflatten(self.fp16_groups_flat[i],
-                                            self.fp16_groups[i])
-            for p, q in zip(self.fp16_groups[i], updated_params):
-                p.data = q.data
+            self._update_model_fp16_weights(i)
 
             #divide the flat weights into near equal partition equal to the data parallel degree
             #each process will compute on a different part of the partition
@@ -249,7 +262,10 @@ def __init__(self,
 
             partition_size = len(self.fp16_groups_flat[i]) / dist.get_world_size(
                 group=self.dp_process_group)
-            params_in_partition, params_not_in_partition, first_offset = self.get_partition_info(self.fp16_groups[i], partition_size, partition_id)
+            params_in_partition, params_not_in_partition, first_offset = self.get_partition_info(
+                self.round_robin_fp16_groups[i],
+                partition_size,
+                partition_id)
 
             self.partition_size.append(partition_size)
             self.params_in_partition.append(params_in_partition)
@@ -263,6 +279,7 @@ def __init__(self,
         self.reduction_stream = torch.cuda.Stream()
         self.cpu_computation_stream = torch.cuda.Stream()
         self.migration_stream = torch.cuda.Stream()
+        self.copy_grad_stream = torch.cuda.Stream()
         self.callback_queued = False
 
         self.param_dict = {}
@@ -383,6 +400,36 @@ def __init__(self,
         if dist.get_rank(group=self.dp_process_group) == 0:
             see_memory_usage(f"After initializing ZeRO optimizer")
 
+    def _update_model_fp16_weights(self, group_index):
+        updated_params = self.unflatten(self.fp16_groups_flat[group_index],
+                                        self.round_robin_fp16_groups[group_index])
+        for p, q in zip(self.round_robin_fp16_groups[group_index], updated_params):
+            p.data = q.data
+
+        # set model fp16 weight to slices of reordered flattened buffer
+        for param_index, param in enumerate(self.fp16_groups[group_index]):
+            new_index = self.round_robin_fp6_indices[group_index][param_index]
+            param.data = self.round_robin_fp16_groups[group_index][new_index].data
+
+    def _round_robin_reorder(self, tensor_list, num_partitions):
+        partition_tensors = {}
+
+        for i, tensor in enumerate(tensor_list):
+            j = i % num_partitions
+            if not j in partition_tensors:
+                partition_tensors[j] = []
+            partition_tensors[j].append((i, tensor))
+
+        reordered_tensors = []
+        reordered_indices = {}
+
+        for partition_index in partition_tensors.keys():
+            for i, (original_index, tensor) in enumerate(partition_tensors[partition_index]):
+                reordered_indices[original_index] = len(reordered_tensors)
+                reordered_tensors.append(tensor)
+
+        return reordered_tensors, reordered_indices
+
     def _release_ipg_buffers(self):
         if self.contiguous_gradients:
             self.ipg_buffer = None
@@ -966,6 +1013,8 @@ def copy_grads_in_partition(self, param):
     def reduce_ipg_grads(self):
         if self.overlap_comm:
             stream = self.reduction_stream
+        elif self.cpu_offload:
+            stream = self.copy_grad_stream
         else:
             stream = torch.cuda.current_stream()
 
@@ -1521,10 +1570,7 @@ def step(self, closure=None):
 
         # TODO: we probably don't need this? just to be safe
         for i in range(len(norm_groups)):
-            updated_params = self.unflatten(self.fp16_groups_flat[i],
-                                            self.fp16_groups[i])
-            for p, q in zip(self.fp16_groups[i], updated_params):
-                p.data = q.data
+            self._update_model_fp16_weights(i)
 
         self.log_timers(timer_names)
         see_memory_usage('After zero_optimizer step')

During backward in ZeRO-2 Offload, reduced gradients are accumulated into CPU memory for later optimizer step. Due to the previous model partitioning scheme, gradient copying to CPU occurs one rank at a time, which slows down backward and under-utilizes PCIe. This PR introduces a new model partitioning scheme that spreads gradient copying evenly among all (most) ranks at any point in time. This improves backward time and PCIe utilization.   