Repository: microsoft/DeepSpeed

Issue:
[BUG] ZeRO 3: save_checkpoint immediately followed by load_checkpoint breaks
**Describe the bug**

This is the reverse of https://github.com/microsoft/DeepSpeed/issues/1748 and is also ZeRO 3-specific.

Here the opposite fails: 

```
deepspeed.load_checkpoint()
# no train in between:
deepspeed.save_checkpoint()
```

**To Reproduce**


```
# add to tests/unit/test_checkpointing.py

def test_immediate_save_load(tmpdir):

    config_dict = {
        "train_batch_size": 4,
        "optimizer": {
            "type": 'Adam'
        },
        "fp16": {
            "enabled": True,
            "initial_scale_power": 8
        },
        "zero_optimization": {
            "stage": 3,
        }
    }
    hidden_dim = 10
    model = SimpleModel(hidden_dim)
    args = args_from_dict(tmpdir, config_dict)

    @distributed_test(world_size=[1])
    def _test_immediate_save_load(args, model, tmpdir):

        ds_model = create_deepspeed_model(args=args, model=model, base_optimizer=None)
        ds_model.save_checkpoint(tmpdir)
        ds_model.load_checkpoint(tmpdir, load_optimizer_states=False, load_lr_scheduler_states=False, load_module_only=False)

    _test_immediate_save_load(args, model, tmpdir)
```

```
$ pytest tests/unit/test_checkpointing.py -k test_immediate_save_load
[...]
Traceback (most recent call last):
  File "/home/stas/anaconda3/envs/py38-pt110/lib/python3.8/multiprocessing/process.py", line 315, in _bootstrap
    self.run()
  File "/home/stas/anaconda3/envs/py38-pt110/lib/python3.8/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/mnt/nvme1/code/github/00optimize/DeepSpeed-save_16bit_model-save_checkpoint_prologue/tests/unit/common.py", line 93, in dist_init
    run_func(*func_args, **func_kwargs)
  File "/mnt/nvme1/code/github/00optimize/DeepSpeed-save_16bit_model-save_checkpoint_prologue/tests/unit/test_checkpointing.py", line 1298, in _test_immediate_save_load
    ds_model.load_checkpoint(tmpdir, load_optimizer_states=False, load_lr_scheduler_states=False, load_module_only=False)
  File "/mnt/nvme1/code/github/00optimize/DeepSpeed-save_16bit_model-save_checkpoint_prologue/deepspeed/runtime/engine.py", line 2452, in load_checkpoint
    load_path, client_states = self._load_checkpoint(load_dir,
  File "/mnt/nvme1/code/github/00optimize/DeepSpeed-save_16bit_model-save_checkpoint_prologue/deepspeed/runtime/engine.py", line 2506, in _load_checkpoint
    self.load_module_state_dict(state_dict=checkpoint['module'],
  File "/mnt/nvme1/code/github/00optimize/DeepSpeed-save_16bit_model-save_checkpoint_prologue/deepspeed/runtime/engine.py", line 2336, in load_module_state_dict
    self.module.load_state_dict(state_dict, strict=strict)
  File "/home/stas/anaconda3/envs/py38-pt110/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1482, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for SimpleModel:
        size mismatch for linears.0.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10, 10]).
        size mismatch for linears.0.bias: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([10]).
```

@tjruwase, @jeffra 



Discussion:
actually after filing this one I recalled that we have diagnosed this condition before. That was the whole point of doing a reinit of the deepspeed engine between save and load that I had to add to HF Transformers:

https://github.com/huggingface/transformers/blob/0fe17f375a4f0fdd9aea260d0645ccfd4896e958/src/transformers/trainer.py#L1521-L1530

So 2 things about this Issue that are new:

1. I added a failing test that you can add to the test suite - saves you time
2. until this is resolved could this situation be somehow detected at least and asserted instead of fail in the very not-obvious way?
e.g.  `load_checkpoint` could check if some state that `save_checkpoint` creates is there?

and then assert with something like:

> currently it's not possible to `load_checkpoint` after `save_checkpoint` was called and requires a new deepspeed engine to use `load_checkpoint` in this situation.



Resolving patch:
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -2457,6 +2457,10 @@ def load_checkpoint(self,
                 )
                 return None, None
 
+        if self.zero_optimization_partition_weights():
+            # Prepare for checkpoint load by ensuring all parameters are partitioned
+            self.optimizer.checkpoint_event_prologue()
+
         load_path, client_states = self._load_checkpoint(load_dir,
                                                          tag,
                                                          load_module_strict=load_module_strict,
@@ -2472,6 +2476,9 @@ def load_checkpoint(self,
             if not success:
                 self.optimizer._restore_from_bit16_weights()
 
+        if self.zero_optimization_partition_weights():
+            self.optimizer.checkpoint_event_epilogue()
+
         return load_path, client_states
 
     def _load_checkpoint(self,
@@ -2737,8 +2744,8 @@ def save_checkpoint(self, save_dir, tag=None, client_state={}, save_latest=True)
         process with rank 0.
         """
         if self.zero_optimization_partition_weights():
-            # Prepare for state_dict() by ensuring all parameters are partitioned
-            self.optimizer.save_checkpoint_prologue()
+            # Prepare for checkpoint save by ensuring all parameters are partitioned
+            self.optimizer.checkpoint_event_prologue()
 
         # This is to make sure the checkpoint names are created without collision
         # There seems to be issue creating them in parallel
@@ -2770,7 +2777,7 @@ def save_checkpoint(self, save_dir, tag=None, client_state={}, save_latest=True)
             self._save_zero_checkpoint(save_dir, tag)
 
         if self.zero_optimization_partition_weights():
-            self.optimizer.save_checkpoint_epilogue()
+            self.optimizer.checkpoint_event_epilogue()
 
         # Save latest checkpoint tag
         torch.distributed.barrier()
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
@@ -401,9 +401,9 @@ def release_sub_module(self, submodule: Module) -> None:
 
     @instrument_w_nvtx
     @torch.no_grad()
-    def release_and_reset_all(self) -> None:
+    def release_and_reset_all(self, module: Module) -> None:
         """release all module parameters"""
-        for param in map(lambda p: p.param, self.__param_order):
+        for param in iter_params(module, recurse=True):
             if param in self.__inflight_param_registry:
                 raise RuntimeError(f"param {param.ds_summary()} still in flight")
 
@@ -412,10 +412,9 @@ def release_and_reset_all(self) -> None:
             param.ds_active_sub_modules.clear()
             self.__release_param(param)
 
-        for param_in_trace in self.__param_order:
-            if param_in_trace.param.ds_status != ZeroParamStatus.NOT_AVAILABLE:
-                raise RuntimeError(
-                    f"{param_in_trace.param.ds_summary()} expected to be released")
+        for param in iter_params(module, recurse=True):
+            if param.ds_status != ZeroParamStatus.NOT_AVAILABLE:
+                raise RuntimeError(f"{param.ds_summary()} expected to be released")
 
     @instrument_w_nvtx
     def __all_gather_params(self, params: Set[Parameter]) -> None:
@@ -2825,7 +2824,7 @@ def _partition_all_parameters(self):
         """Partitioning Parameters that were not partitioned usually if parameters
         of modules whose input parameters do not require grad computation do not
         trigger post call and will therefore will remain unpartitioned"""
-        self.param_coordinator.release_and_reset_all()
+        self.param_coordinator.release_and_reset_all(self.module)
         for param in iter_params(self.module, recurse=True):
             if param.ds_status != ZeroParamStatus.NOT_AVAILABLE:
                 raise RuntimeError(f"{param.ds_summary()} expected to be released")
@@ -3110,10 +3109,10 @@ def load_state_dict(self,
             self.persistent_parameters[0].partition(self.persistent_parameters)
             self.persistent_parameters[0].all_gather(self.persistent_parameters)
 
-    def save_checkpoint_prologue(self):
+    def checkpoint_event_prologue(self):
         self._partition_all_parameters()
 
-    def save_checkpoint_epilogue(self):
+    def checkpoint_event_epilogue(self):
         if len(self.persistent_parameters) > 0:
             self.persistent_parameters[0].all_gather(self.persistent_parameters)
 


Test patch:
diff --git a/tests/unit/test_checkpointing.py b/tests/unit/test_checkpointing.py
--- a/tests/unit/test_checkpointing.py
+++ b/tests/unit/test_checkpointing.py
@@ -1266,3 +1266,82 @@ def _go1(models):
             model.load_checkpoint(tmpdir, load_optimizer_states=load_optim)
 
     _go1(models)
+
+
+@pytest.mark.parametrize('zero_stage', [0, 1, 2, 3])
+def test_immediate_save_load(tmpdir, zero_stage):
+    config_dict = {
+        "train_batch_size": 4,
+        "optimizer": {
+            "type": 'Adam'
+        },
+        "fp16": {
+            "enabled": True,
+            "initial_scale_power": 8
+        },
+        "zero_optimization": {
+            "stage": zero_stage,
+        }
+    }
+    hidden_dim = 10
+    model = SimpleModel(hidden_dim)
+    args = args_from_dict(tmpdir, config_dict)
+
+    @distributed_test(world_size=[1])
+    def _test_immediate_save_load(args, model, tmpdir):
+
+        ds_model = create_deepspeed_model(args=args, model=model, base_optimizer=None)
+        ds_model.save_checkpoint(tmpdir)
+        ds_model.load_checkpoint(tmpdir,
+                                 load_optimizer_states=False,
+                                 load_lr_scheduler_states=False,
+                                 load_module_only=False)
+
+    _test_immediate_save_load(args, model, tmpdir)
+
+
+@pytest.mark.parametrize('zero_stage', [0, 1, 2, 3])
+def test_load_immediate_save(tmpdir, zero_stage):
+    config_dict = {
+        "train_batch_size": 4,
+        "optimizer": {
+            "type": 'Adam'
+        },
+        "fp16": {
+            "enabled": True,
+            "initial_scale_power": 8
+        },
+        "zero_optimization": {
+            "stage": zero_stage,
+        }
+    }
+    hidden_dim = 10
+    model = SimpleModel(hidden_dim)
+    args = args_from_dict(tmpdir, config_dict)
+
+    @distributed_test(world_size=[1])
+    def _test_load_immediate_save(args, model, tmpdir):
+
+        # 1. pretrain a model and save it
+        dtype = torch.half
+        ds_model = create_deepspeed_model(args=args, model=model, base_optimizer=None)
+        data_loader = random_dataloader(model=ds_model,
+                                        total_samples=1,
+                                        hidden_dim=hidden_dim,
+                                        device=ds_model.device,
+                                        dtype=dtype)
+        for n, batch in enumerate(data_loader):
+            loss = ds_model(batch[0], batch[1])
+            ds_model.backward(loss)
+            ds_model.step()
+        ds_model.save_checkpoint(tmpdir)
+
+        # 2. load and immediately save a model with a fresh ds engine
+        ds_model = create_deepspeed_model(args=args, model=model, base_optimizer=None)
+        ds_model.load_checkpoint(tmpdir,
+                                 load_optimizer_states=False,
+                                 load_lr_scheduler_states=False,
+                                 load_module_only=False)
+        ds_model.save_checkpoint(tmpdir)
+
+    _test_load_immediate_save(args, model, tmpdir)