[squash] zero-ckpt-cpu-issue
diff --git a/deepspeed/runtime/constants.py b/deepspeed/runtime/constants.py
index daf9ffeaadff..7892dda1a76f 100755
--- a/deepspeed/runtime/constants.py
+++ b/deepspeed/runtime/constants.py
@@ -444,3 +444,9 @@ class ValidationMode:
 '''
 DATALOADER_DROP_LAST = "dataloader_drop_last"
 DATALOADER_DROP_LAST_DEFAULT = False
+
+#########################################
+# Optimizer checkpoint keys
+#########################################
+OPTIMIZER_STATE_DICT = "optimizer_state_dict"
+FP32_GROUPS = "fp32_groups"
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
old mode 100755
new mode 100644
index 4ba3a2be7627..e206444efde9
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -42,9 +42,10 @@
 from deepspeed.runtime.dataloader import DeepSpeedDataLoader
 from deepspeed.runtime.constants import \
     ROUTE_TRAIN, ROUTE_PREDICT, ROUTE_EVAL, \
-    PLD_THETA, PLD_GAMMA
+    PLD_THETA, PLD_GAMMA, OPTIMIZER_STATE_DICT
 from deepspeed.runtime.zero.constants import \
-    ZERO_OPTIMIZATION_OPTIMIZER_STATES, ZERO_OPTIMIZATION_GRADIENTS, ZERO_OPTIMIZATION_WEIGHTS
+    ZERO_OPTIMIZATION_OPTIMIZER_STATES, ZERO_OPTIMIZATION_GRADIENTS, ZERO_OPTIMIZATION_WEIGHTS, \
+    SINGLE_PARTITION_OF_FP32_GROUPS
 from deepspeed.runtime.sparse_tensor import SparseTensor
 
 import deepspeed.runtime.lr_schedules as lr_schedules
@@ -1385,7 +1386,8 @@ def _configure_zero_optimizer(self, optimizer):
                 round_robin_gradients=round_robin_gradients,
                 has_moe_layers=self.has_moe_layers,
                 fp16_master_weights_and_gradients=self.fp16_master_weights_and_gradients(
-                ))
+                ),
+                elastic_checkpoint=self.zero_elastic_checkpoint())
 
         elif zero_stage == ZERO_OPTIMIZATION_WEIGHTS:
             assert not self.has_moe_layers, "MoE not supported with Stage 3"
@@ -2577,10 +2579,27 @@ def _get_all_zero_checkpoints(self, load_dir, tag):
             return None
 
         zero_sd_list = []
-        for ckpt_name in zero_ckpt_names:
-            zero_sd_list.append(torch.load(ckpt_name, map_location="cpu"))
+        for i, ckpt_name in enumerate(zero_ckpt_names):
+            _state = None
+            # Fully load state for current rank
+            if self.zero_elastic_checkpoint() or dist.get_rank(
+                    group=self.optimizer.dp_process_group) == i:
+                _state = torch.load(ckpt_name, map_location='cpu')
+            elif self.zero_optimization_stage(
+            ) <= ZERO_OPTIMIZATION_GRADIENTS and self.zero_load_from_fp32_weights():
+                # Extract fp32 groups only, otherwise throw away to prevent unnecessary CPU memory overheads
+                _state = torch.load(ckpt_name, map_location='cpu')
+                _state = {
+                    OPTIMIZER_STATE_DICT: {
+                        SINGLE_PARTITION_OF_FP32_GROUPS:
+                        _state[OPTIMIZER_STATE_DICT][SINGLE_PARTITION_OF_FP32_GROUPS]
+                    }
+                }
+            else:
+                _state = {OPTIMIZER_STATE_DICT: None}
+            zero_sd_list.append(_state)
 
-        zero_optimizer_sd = [sd["optimizer_state_dict"] for sd in zero_sd_list]
+        zero_optimizer_sd = [sd[OPTIMIZER_STATE_DICT] for sd in zero_sd_list]
         print(
             f"successfully loaded {len(zero_optimizer_sd)} ZeRO state_dicts for rank {self.global_rank}"
         )
diff --git a/deepspeed/runtime/fp16/fused_optimizer.py b/deepspeed/runtime/fp16/fused_optimizer.py
index 370bd9cae5c1..8b96d511a8eb 100755
--- a/deepspeed/runtime/fp16/fused_optimizer.py
+++ b/deepspeed/runtime/fp16/fused_optimizer.py
@@ -6,12 +6,12 @@
 '''
 
 import torch
-import math
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 
 from deepspeed.runtime.utils import get_global_norm, get_grad_norm, CheckOverflow, get_weight_norm
 from deepspeed.runtime.fp16.loss_scaler import INITIAL_LOSS_SCALE, SCALE_WINDOW, MIN_LOSS_SCALE
 from deepspeed.utils import groups, logger, log_dist
+from deepspeed.runtime.constants import OPTIMIZER_STATE_DICT
 import torch.distributed as dist
 
 
@@ -396,7 +396,7 @@ def state_dict(self):
             state_dict['last_overflow_iter'] = self.last_overflow_iter
             state_dict['scale_factor'] = self.scale_factor
             state_dict['scale_window'] = self.scale_window
-        state_dict['optimizer_state_dict'] = self.optimizer.state_dict()
+        state_dict[OPTIMIZER_STATE_DICT] = self.optimizer.state_dict()
         state_dict['fp32_groups_flat'] = self.fp32_groups_flat
         state_dict['clip_grad'] = self.clip_grad
         return state_dict
@@ -431,7 +431,7 @@ def load_state_dict(self, state_dict, load_optimizer_states=True):
             self.scale_factor = state_dict['scale_factor']
             self.scale_window = state_dict['scale_window']
         if load_optimizer_states:
-            self.optimizer.load_state_dict(state_dict['optimizer_state_dict'])
+            self.optimizer.load_state_dict(state_dict[OPTIMIZER_STATE_DICT])
         self.clip_grad = state_dict['clip_grad']
         # At this point, the optimizer's references to the model's fp32 parameters are up to date.
         # The optimizer's hyperparameters and internal buffers are also up to date.
diff --git a/deepspeed/runtime/fp16/unfused_optimizer.py b/deepspeed/runtime/fp16/unfused_optimizer.py
index d0110d5f828c..cd1c20ef4235 100755
--- a/deepspeed/runtime/fp16/unfused_optimizer.py
+++ b/deepspeed/runtime/fp16/unfused_optimizer.py
@@ -7,12 +7,12 @@
 
 from deepspeed.moe.utils import split_params_grads_into_shared_and_expert_params
 import torch
-from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
-import math
+from torch._utils import _flatten_dense_tensors
 
-from deepspeed.runtime.utils import get_global_norm, get_grad_norm, CheckOverflow, get_weight_norm
+from deepspeed.runtime.utils import get_global_norm, CheckOverflow, get_weight_norm
 from deepspeed.runtime.fp16.loss_scaler import INITIAL_LOSS_SCALE, SCALE_WINDOW, MIN_LOSS_SCALE
 from deepspeed.utils import logger
+from deepspeed.runtime.constants import OPTIMIZER_STATE_DICT
 
 
 class FP16_UnfusedOptimizer(object):
@@ -332,7 +332,7 @@ def state_dict(self):
             state_dict['last_overflow_iter'] = self.last_overflow_iter
             state_dict['scale_factor'] = self.scale_factor
             state_dict['scale_window'] = self.scale_window
-        state_dict['optimizer_state_dict'] = self.optimizer.state_dict()
+        state_dict[OPTIMIZER_STATE_DICT] = self.optimizer.state_dict()
         state_dict['fp32_groups'] = self.fp32_groups
         return state_dict
 
@@ -368,7 +368,7 @@ def load_state_dict(self, state_dict, load_optimizer_states=True):
             self.scale_window = state_dict['scale_window']
 
         if load_optimizer_states:
-            self.optimizer.load_state_dict(state_dict['optimizer_state_dict'])
+            self.optimizer.load_state_dict(state_dict[OPTIMIZER_STATE_DICT])
         # At this point, the optimizer's references to the model's fp32 parameters are up to date.
         # The optimizer's hyperparameters and internal buffers are also up to date.
         # However, the fp32 master copies of the model's fp16 params stored by the optimizer are still
diff --git a/deepspeed/runtime/zero/constants.py b/deepspeed/runtime/zero/constants.py
index e3b2dfc0c68f..7ff99b25b10b 100755
--- a/deepspeed/runtime/zero/constants.py
+++ b/deepspeed/runtime/zero/constants.py
@@ -48,6 +48,8 @@
 ZERO_OPTIMIZATION_STAGE_2 = 'stage_2'
 ZERO_OPTIMIZATION_STAGE_3 = 'stage_3'
 
+SINGLE_PARTITION_OF_FP32_GROUPS = "single_partition_of_fp32_groups"
+
 ZERO_OPTIMIZATION_STAGE_DEFAULT = ZERO_OPTIMIZATION_DISABLED
 
 ZERO_OPTIMIZATION_ALLGATHER_PARTITIONS = 'allgather_partitions'
@@ -74,7 +76,7 @@
 ZERO_OPTIMIZATION_LOAD_FROM_FP32_WEIGHTS_DEFAULT = True
 
 ZERO_OPTIMIZATION_ELASTIC_CHECKPOINT = 'elastic_checkpoint'
-ZERO_OPTIMIZATION_ELASTIC_CHECKPOINT_DEFAULT = True
+ZERO_OPTIMIZATION_ELASTIC_CHECKPOINT_DEFAULT = False
 
 ZERO_OPTIMIZATION_CPU_OFFLOAD = 'cpu_offload'
 ZERO_OPTIMIZATION_CPU_OFFLOAD_DEFAULT = False
diff --git a/deepspeed/runtime/zero/stage2.py b/deepspeed/runtime/zero/stage2.py
index b995e4dd975c..604bac71f771 100755
--- a/deepspeed/runtime/zero/stage2.py
+++ b/deepspeed/runtime/zero/stage2.py
@@ -5,13 +5,9 @@
 import torch
 from torch.distributed.distributed_c10d import _get_global_rank
 import torch.distributed as dist
-import math
 from torch._six import inf
-from torch.autograd import Variable
 from packaging import version as pkg_version
 
-import collections
-
 from deepspeed.runtime.fp16.loss_scaler import LossScaler, DynamicLossScaler
 from deepspeed.runtime.utils import bwc_tensor_model_parallel_rank, get_global_norm, see_memory_usage, is_model_parallel_parameter
 from deepspeed.runtime.zero.config import ZERO_OPTIMIZATION_GRADIENTS
@@ -22,6 +18,8 @@
 from deepspeed.moe.utils import is_moe_param
 from deepspeed.git_version_info import version
 
+from .constants import SINGLE_PARTITION_OF_FP32_GROUPS
+
 # Toggle this to true to enable correctness test
 # with gradient partitioning and without
 pg_correctness_test = False
@@ -107,7 +105,8 @@ def __init__(self,
                  partition_grads=True,
                  round_robin_gradients=False,
                  has_moe_layers=False,
-                 fp16_master_weights_and_gradients=False):
+                 fp16_master_weights_and_gradients=False,
+                 elastic_checkpoint=False):
 
         if dist.get_rank() == 0:
             logger.info(f"Reduce bucket size {reduce_bucket_size}")
@@ -118,6 +117,8 @@ def __init__(self,
         # 1. maintain same user API from apex.fp16_utils
         # 2. keep common stuff here in case we need to add ne552w fused optimizer later
 
+        self.elastic_checkpoint = elastic_checkpoint
+
         # differences from apex.fp16_utils:
         # - assume all model params in fp16
         # - assume all params requires grad
@@ -1969,7 +1970,17 @@ def state_dict(self):
         state_dict['loss_scaler'] = self.loss_scaler
         state_dict['dynamic_loss_scale'] = self.dynamic_loss_scale
         state_dict['overflow'] = self.overflow
-        state_dict['base_optimizer_state'] = self._get_base_optimizer_state()
+
+        if self.elastic_checkpoint:
+            state_dict['base_optimizer_state'] = self._get_base_optimizer_state()
+            # Remove paddings for DP alignment to enable loading for other alignment values
+            fp32_groups_without_padding = self._get_groups_without_padding(
+                self.single_partition_of_fp32_groups)
+            state_dict[SINGLE_PARTITION_OF_FP32_GROUPS] = fp32_groups_without_padding
+        else:
+            state_dict['base_optimizer_state'] = self.optimizer.state_dict()
+            state_dict[
+                SINGLE_PARTITION_OF_FP32_GROUPS] = self.single_partition_of_fp32_groups
 
         state_dict['zero_stage'] = ZERO_OPTIMIZATION_GRADIENTS
         state_dict['partition_count'] = self.partition_count
@@ -1979,7 +1990,7 @@ def state_dict(self):
         # Remove paddings for DP alignment to enable loading for other alignment values
         fp32_groups_without_padding = self._get_groups_without_padding(
             self.single_partition_of_fp32_groups)
-        state_dict['single_partition_of_fp32_groups'] = fp32_groups_without_padding
+        state_dict[SINGLE_PARTITION_OF_FP32_GROUPS] = fp32_groups_without_padding
 
         #        if self.cpu_offload:
         #            state_dict_tmp = async_copy_to(state_dict,
@@ -1999,7 +2010,7 @@ def _restore_from_fp32_weights(self, all_state_dict):
         for i in range(len(self.single_partition_of_fp32_groups)):
             partition_id = dist.get_rank(group=self.real_dp_process_group[i])
             merged_partitions = [
-                sd['single_partition_of_fp32_groups'][i] for sd in all_state_dict
+                sd[SINGLE_PARTITION_OF_FP32_GROUPS][i] for sd in all_state_dict
             ]
             if self.is_moe_group(self.optimizer.param_groups[i]):
                 ranks = self.get_ep_ranks()
@@ -2112,25 +2123,36 @@ def load_state_dict(self,
             model.load_state_dict(checkpoint['model'])
             optimizer.load_state_dict(checkpoint['optimizer'])
         """
+        dp_rank = dist.get_rank(group=self.dp_process_group)
         # I think it should actually be ok to reload the optimizer before the model.
-        self.loss_scaler = state_dict_list[0]['loss_scaler']
-        self.dynamic_loss_scale = state_dict_list[0]['dynamic_loss_scale']
-        self.overflow = state_dict_list[0]['overflow']
+        self.loss_scaler = state_dict_list[dp_rank]['loss_scaler']
+        self.dynamic_loss_scale = state_dict_list[dp_rank]['dynamic_loss_scale']
+        self.overflow = state_dict_list[dp_rank]['overflow']
+
+        ckpt_version = state_dict_list[dp_rank].get("ds_version", False)
+        assert ckpt_version, f"Empty ds_version! {error_str}"
+        ckpt_version = pkg_version.parse(ckpt_version)
 
         # zero stage 1 mode
         if not self.partition_gradients:
             required_version = pkg_version.parse("0.3.17")
-            ckpt_version = state_dict_list[0].get("ds_version", False)
             error_str = f"ZeRO stage 1 changed in {required_version} and is not backwards compatible " \
                 "with older stage 1 checkpoints. If you'd like to load an old ZeRO-1 checkpoint " \
                 "please set 'legacy_stage1': true in your zero config json. This old version of " \
                 "stage 1 will be removed in v0.4.0."
+            assert required_version <= ckpt_version, f"Old version: {ckpt_version} {error_str}"
 
-            assert ckpt_version, f"Empty ds_version! {error_str}"
-            assert required_version <= pkg_version.parse(ckpt_version), f"Old version: {ckpt_version} {error_str}"
+        if ckpt_version < pkg_version.parse("0.5.8"):
+            # zero checkpoints before 0.5.8 defaulted to elastic enabled, must
+            # load checkpoint state using elastic logic
+            self.elastic_checkpoint = True
 
         if load_optimizer_states:
-            self._restore_base_optimizer_state(state_dict_list)
+            if self.elastic_checkpoint:
+                self._restore_base_optimizer_state(state_dict_list)
+            else:
+                self.optimizer.load_state_dict(
+                    state_dict_list[dp_rank]['base_optimizer_state'])
 
         # At this point, the optimizer's references to the model's fp32 parameters are up to date.
         # The optimizer's hyperparameters and internal buffers are also up to date.
@@ -2148,8 +2170,10 @@ def load_state_dict(self,
         # are guaranteed to exist, so we can just copy_() from the saved master params.
 
         if load_from_fp32_weights:
+            # option 2 from above
             self._restore_from_fp32_weights(state_dict_list)
         else:
+            # option 1 from above
             self._restore_from_bit16_weights()
 
 
diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py
index 1387ced6fc8d..a86ed6003e32 100755
--- a/deepspeed/runtime/zero/stage3.py
+++ b/deepspeed/runtime/zero/stage3.py
@@ -26,6 +26,7 @@
 from deepspeed.runtime.swap_tensor.partitioned_param_swapper import PartitionedParamStatus
 from deepspeed.runtime.swap_tensor.partitioned_optimizer_swapper import PartitionedOptimizerSwapper
 from deepspeed.runtime.swap_tensor.pipelined_optimizer_swapper import PipelinedOptimizerSwapper
+from deepspeed.runtime.constants import OPTIMIZER_STATE_DICT
 
 # Toggle this to true to enable correctness test
 # with gradient partitioning and without
@@ -3114,7 +3115,7 @@ def _rigid_state_dict(self):
         state_dict['partition_count'] = self.partition_count
 
         self._set_fp32_optimizer_param_groups()
-        state_dict['optimizer_state_dict'] = self.optimizer.state_dict()
+        state_dict[OPTIMIZER_STATE_DICT] = self.optimizer.state_dict()
         state_dict['fp32_flat_groups'] = self.fp32_partitioned_groups_flat
         self._clear_fp32_optimizer_param_groups()
 
@@ -3227,7 +3228,7 @@ def _rigid_load_state_dict(self, state_dict, load_optimizer_states=True):
 
         if load_optimizer_states:
             self._set_fp32_optimizer_param_groups()
-            self.optimizer.load_state_dict(state_dict['optimizer_state_dict'])
+            self.optimizer.load_state_dict(state_dict[OPTIMIZER_STATE_DICT])
             self._clear_fp32_optimizer_param_groups()
 
         # restore fp32 partitions
