Repository: microsoft/DeepSpeed

Issue:
[BUG] Loading MoE Checkpint with load_optimizer_states=False should tolerate missing optim_state files
**Describe the bug**
Even though not loaded into optimizer state dict for MoE model with `load_optimizer_states=False` the current code attempts to load the files at https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/engine.py#L2843, but they could be (intentionally or not) missing.
The only workaround right now is to use `load_module_only` which is not ideal if we want to load LR scheduler states for example.

**To Reproduce**
Steps to reproduce the behavior:
1. Delete `optim_state` files from the checkpoint
2. Try loading DeepSpeed checkpoint with `load_optimizer_states=False`
3. Observe `FileNotFoundError`:
```
Traceback (most recent call last):
--
...
_, checkpoint_state_dict = model.load_checkpoint(
File "/usr/local/lib/python3.8/dist-packages/deepspeed/runtime/engine.py", line 2415, in load_checkpoint
load_path, client_states = self._load_checkpoint(load_dir,
File "/usr/local/lib/python3.8/dist-packages/deepspeed/runtime/engine.py", line 2489, in _load_checkpoint
optim_checkpoint = torch.load(optim_load_path,
File "/usr/local/lib/python3.8/dist-packages/torch/serialization.py", line 594, in load
with _open_file_like(f, 'rb') as opened_file:
File "/usr/local/lib/python3.8/dist-packages/torch/serialization.py", line 230, in _open_file_like
return _open_file(name_or_buffer, mode)
File "/usr/local/lib/python3.8/dist-packages/torch/serialization.py", line 211, in __init__
super(_open_file, self).__init__(open(name, mode))
FileNotFoundError: [Errno 2] No such file or directory: '<REDACTED>/expp_rank_0_mp_rank_00_optim_states.pt'

```

**Expected behavior**
Loading checkpoints loads all specified state dicts, but for optimizer - it's initialized from model weights. `client_state` should reflect this.

**ds_report output**
```
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
async_io ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  using untested triton version (1.0.1), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/usr/local/lib64/python3.7/site-packages/torch']
torch version .................... 1.12.1+cu113
deepspeed install path ........... ['/home/$USER/DeepSpeed/deepspeed']
deepspeed info ................... 0.8.1+8df50a26, 8df50a26, master
torch cuda version ............... 11.3
torch hip version ................ None
nvcc version ..................... 11.0
deepspeed wheel compiled w. ...... torch 1.12, cuda 11.3
```

**Screenshots**
N/A

**System info (please complete the following information):**
 - OS: `Amazon Linux 2 4.14.200-155.322.amzn2.x86_64`
 - GPU count and types `8 x A100`
 - Interconnects (if applicable) `N/A`
 - Python version: `3.8.10`
 - Any other relevant info about your setup

**Launcher context**
DeepSpeed launcher / pytest

**Docker context**
`N/A`

**Additional context**
`N/A`



Discussion:
Ok, this could be a misunderstanding on my part. Maybe @tjruwase you can help me understand.

I thought `load_optimizer_states` effectively does not load/apply `optim_states.pt` files (especially for ZeRO 1-2), e.g. if the user wants to try a different world_size.

But as far as I can see it is only controlling loading optimizer `state_dict`, but other content from `optim_states.pt` are still loaded, isn't it? E.g. `dynamic_loss_scale` etc.

If that's the case I'm not sure what are the recommendations for loading a checkpoint with different world_size. Is the user expected to create dummy `optim_states.pt` for all the new ranks even if what I want is to, say, load LR schedules in addition to all model parameters.
I also meet this problem.
Besides, if I increases some nodes during fine-tuning, it will also raise an FileNotFoundError.

Resolving patch:
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -2759,26 +2759,29 @@ def _load_checkpoint(self,
 
         self.loaded_checkpoint_dp_world_size = checkpoint['dp_world_size']
 
+        optim_checkpoint = None
         if load_module_only:
             deepspeed_states = ['module']
             if self.optimizer is not None and self.fp16_enabled():
                 self.optimizer.refresh_fp32_params()
         else:
-            if self.has_moe_layers:
-                largest_group_name = groups._get_max_expert_size_name()
-                expp_rank = groups._get_expert_parallel_rank(largest_group_name)
-                optim_load_path = self._get_optimizer_ckpt_name(load_dir, tag, expp_rank)
-                optim_checkpoint = self.checkpoint_engine.load(optim_load_path, map_location=torch.device('cpu'))
-            else:
-                optim_checkpoint = checkpoint
-
             has_zero_optimizer_state = self.zero_optimization() or self.bfloat16_enabled()
             if load_optimizer_states and self.optimizer is not None and not has_zero_optimizer_state:
-                if self.fp16_enabled():
+                if self.has_moe_layers:
+                    largest_group_name = groups._get_max_expert_size_name()
+                    expp_rank = groups._get_expert_parallel_rank(largest_group_name)
+                    optim_load_path = self._get_optimizer_ckpt_name(load_dir, tag, expp_rank)
+                    optim_checkpoint = self.checkpoint_engine.load(optim_load_path, map_location=torch.device('cpu'))
+                else:
+                    optim_checkpoint = checkpoint
+
+                if self.fp16_enabled() or self.bfloat16_enabled():
                     self.optimizer.load_state_dict(optim_checkpoint['optimizer'],
                                                    load_optimizer_states=load_optimizer_states)
                 else:
-                    self.optimizer.load_state_dict(optim_checkpoint['optimizer'])
+                    optim_checkpoint = checkpoint
+
+                self.optimizer.load_state_dict(optim_checkpoint['optimizer'])
 
             if load_lr_scheduler_states and self.lr_scheduler is not None:
                 self.lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
@@ -2835,7 +2838,7 @@ def get_sparse_tensor_module_names(original_set, loaded_set, original_parameters
 
         client_state = {key: value for key, value in checkpoint.items() if not key in deepspeed_states}
 
-        if not load_optimizer_states and not load_module_only:
+        if optim_checkpoint is not None:
             client_state['optimizer'] = optim_checkpoint['optimizer']
 
         return load_path, client_state


Test patch:
diff --git a/tests/unit/checkpoint/common.py b/tests/unit/checkpoint/common.py
--- a/tests/unit/checkpoint/common.py
+++ b/tests/unit/checkpoint/common.py
@@ -15,6 +15,7 @@
 from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
 
 from unit.simple_model import *
+from unittest.mock import MagicMock, patch
 
 
 def compare_deepspeed_states(saved_model, loaded_model):
@@ -209,11 +210,17 @@ def checkpoint_correctness_verification(config_dict,
     loaded_model = create_deepspeed_model(config_dict=config_dict, model=models[1], base_optimizer=base_optimizers[1])
     assert list(trained_model.parameters())[0].dtype == list(loaded_model.parameters())[0].dtype
 
-    loaded_model.load_checkpoint(save_folder,
-                                 tag=save_tag,
-                                 load_optimizer_states=load_optimizer_states,
-                                 load_lr_scheduler_states=load_lr_scheduler_states,
-                                 load_module_only=load_module_only)
+    context = patch.object(loaded_model, "_get_optimizer_ckpt_name",
+                           wraps=loaded_model._get_optimizer_ckpt_name) if not load_optimizer_states else MagicMock()
+    with context as optim_load_state_dict_mock:
+        loaded_model.load_checkpoint(save_folder,
+                                     tag=save_tag,
+                                     load_optimizer_states=load_optimizer_states,
+                                     load_lr_scheduler_states=load_lr_scheduler_states,
+                                     load_module_only=load_module_only)
+        if not load_optimizer_states:
+            # should not attempt to get the file name to load it
+            optim_load_state_dict_mock.assert_not_called()
 
     compare_model_states(trained_model,
                          loaded_model,