Load non-DeepSpeed checkpoints into ZeRO optimizer

deepspeed/pt/deepspeed_light.py

@@ -561,7 +561,6 @@ class DeepSpeedLight(Module):
 
         if zero_stage == ZERO_OPTIMIZATION_OPTIMIZER_STATES:
             assert self.zero_reduce_scatter(), 'Stage 1 only supports reduce scatter mode'
-            logger.info('Creating fp16 ZeRO Optimizer Stage 1')
             optimizer = FP16_DeepSpeedZeroOptimizer_Stage1(
                 optimizer,
                 static_loss_scale=self.loss_scale(),
@@ -593,7 +592,6 @@ class DeepSpeedLight(Module):
                 gradient_predivide_factor=self.gradient_predivide_factor())
         else:
             raise NotImplementedError("ZeRO stage {} not implemented".format(zero_stage))
-        logger.info('Creating fp16 zero stage {} optimizer'.format(zero_stage))
 
         return optimizer
 

deepspeed/pt/fp16_optimizer.py

@@ -353,6 +353,7 @@ class FP16_Optimizer(object):
         state_dict['clip_grad'] = self.clip_grad
         return state_dict
 
+    # Refresh fp32 master params from fp16 copies
     def refresh_fp32_params(self):
         for current, saved in zip(self.fp32_groups_flat, self.fp16_groups_flat):
             current.data.copy_(saved.data)

deepspeed/pt/zero_optimizer_stage1.py

@@ -10,6 +10,24 @@ from deepspeed.pt.loss_scaler import LossScaler, DynamicLossScaler
 from deepspeed.pt.deepspeed_utils import get_grad_norm, CheckOverflow
 
 
+def flatten_dense_tensors_sub_partition_aligned_(tensor_list,
+                                                dp,
+                                                max_elements_per_comm,
+                                                pg):
+    num_elements = sum(t.numel() for t in tensor_list)
+    log_dist("Total number of elements in model: {}, max elements per com: {}".format(
+        num_elements,
+        max_elements_per_comm),
+        ranks=[0])
+
+    # Compute aligned partition size based on parameter count
+    aligned_param_partition_size = math.ceil(num_elements / dp)
+
+    # Compute aligned partition size based on communication size
+    aligned_comm_partition_size = int(max_elements_per_comm // dp)
+
+
+
 def flatten_dense_tensors_sub_partition_aligned(tensor_list,
                                                 dp,
                                                 max_elements_per_comm,
@@ -780,6 +798,14 @@ class FP16_DeepSpeedZeroOptimizer_Stage1(object):
             'local_sub_partitions_of_fp32_groups'] = self.local_sub_partitions_of_fp32_groups
         return state_dict
 
+    # Refresh the fp32 master params from the fp16 copies.
+    def refresh_fp32_params(self):
+        partition_id = dist.get_rank(group=self.dp_process_group)
+        for fp16_all_sub_partitions, fp32_local_sub_partitions in zip(self.parallel_sub_partitioned_fp16_groups, self.local_sub_partitions_of_fp32_groups):
+            for local_sub_partition_param_fp16, local_sub_partition_param_fp32 in zip(fp16_all_sub_partitions[partition_id], fp32_local_sub_partitions):
+                local_sub_partition_param_fp32.data.copy_(
+                    local_sub_partition_param_fp16.data)
+
     def load_state_dict(self, state_dict, load_optimizer_states=True):
         """
         Loads a state_dict created by an earlier call to state_dict().