Repository: microsoft/DeepSpeed

Issue:
[BUG] Bloated MoE expert checkpoint files with ZeRO 2
**Describe the bug**
It appears that with ZeRO 2 MoE layer expert parameters are replicated across all ranks by sharing a single tensor storage (with different offsets).
Since all MoE layer experts are stored in individual files for all ranks the default `torch.save()` behavior stores the underlying tensor's storage (even though a given expert is only covering a fraction of the storage).

Thus we end up with the following file structure per layer N:
```
layer_N_expert_E_mp_rank_00_model_states.pt
...
layer_N_expert_2_mp_rank_00_model_states.pt
layer_N_expert_1_mp_rank_00_model_states.pt
layer_N_expert_0_mp_rank_00_model_states.pt
```

Where _each_ file has the storage for all E experts across layer N, but the tensors in the files are only using a fraction of it.
When `torch.load()`ed each end up being tensors with independent storage.

It seems like using `.clone().detach()` when saving MoE expert tensors to individual files should address the issue.

But I was wondering about the reason behind making each expert and individual file for all ranks. E.g. if we stored a single file for expert layer for all ranks then the storage would have been reused for the loaded tensors and this single file would not be bloated.

I'll provide a test shortly, willing to contribute the fix.

**To Reproduce**
Steps to reproduce the behavior:
1. Go to '...'
2. Click on '....'
3. Scroll down to '....'
4. See error

**Expected behavior**
Checkpoint files should be compact, no redundancy.

**ds_report output**
Please run `ds_report` to give us details about your setup.

**Screenshots**
If applicable, add screenshots to help explain your problem.

**System info (please complete the following information):**
 - OS: Amazon Linux 4.14.252-195.483.amzn2.x86_64
 - GPU count and types: 8x A100 on a single machine
 - Interconnects (if applicable): single machine
 - Python version: 3.8
 - Any other relevant info about your setup: DeepSpeed 0.6.0, 0.7.4

**Launcher context**
DeepSpeed runner/launcher

**Docker context**
Custom docker image

**Additional context**
Using MoE layers, ZeRO 2



Discussion:

Resolving patch:
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -2989,8 +2989,9 @@ def _save_moe_checkpoint(self, save_dir, tag, client_state={}):
                         num_local_experts + int(local_expert_id)
                     expert_key = key.replace(f'{moe_str_prefix}{local_expert_id}',
                                              f'{moe_str_prefix}{global_expert_id}')
-                    experts_state_dict[str(
-                        global_expert_id)][expert_key] = moe_state_dict.pop(key)
+                    # truncating extra tensor (shared) storage
+                    truncated = moe_state_dict.pop(key).clone().detach()
+                    experts_state_dict[str(global_expert_id)][expert_key] = truncated
 
                 # let save the moe parameters
                 for global_expert_id, expert_state_dict in experts_state_dict.items():


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
@@ -115,16 +115,19 @@ def compare_lr_scheduler_states(saved_model, loaded_model):
             assert state0 == state1
 
 
-def create_deepspeed_model(config_dict, model, base_optimizer):
-    if base_optimizer is None:
-        ds_model, _, _, _ = deepspeed.initialize(config=config_dict,
-                                                 model=model,
-                                                 model_parameters=model.parameters())
-    else:
-        ds_model, _, _, _ = deepspeed.initialize(config=config_dict,
-                                                model=model,
-                                                optimizer=base_optimizer)
+# following mixture-of-experts.md
+def create_moe_param_groups(model):
+    from deepspeed.moe.utils import split_params_into_different_moe_groups_for_optimizer
+
+    parameters = {'params': [p for p in model.parameters()], 'name': 'parameters'}
+    return split_params_into_different_moe_groups_for_optimizer(parameters)
 
+
+def create_deepspeed_model(config_dict, model, base_optimizer):
+    ds_model, _, _, _ = deepspeed.initialize(config=config_dict,
+                                             model=model,
+                                             model_parameters=create_moe_param_groups(model),
+                                             optimizer=base_optimizer)
     return ds_model
 
 
@@ -178,6 +181,19 @@ def checkpoint_correctness_verification(config_dict,
 
     dist.barrier()
 
+    for root, _, files in os.walk(save_folder):
+        for f in files:
+            if "_expert_" in f and "_model_states" in f:
+                expert = torch.load(os.path.join(root, f))
+                needed, storages = 0, {}
+                for name, tensor in expert.items():
+                    needed += tensor.size().numel()
+                    storage = tensor.storage()
+                    # some storage can be shared within an expert's checkpoint
+                    storages[storage.data_ptr()] = storage.size()
+                stored = sum(v for _, v in storages.items())
+                assert needed == stored, f"MoE expert checkpoint uses more storage than required: {f}"
+
     loaded_model = create_deepspeed_model(config_dict=config_dict,
                                           model=models[1],
                                           base_optimizer=base_optimizers[1])
diff --git a/tests/unit/simple_model.py b/tests/unit/simple_model.py
--- a/tests/unit/simple_model.py
+++ b/tests/unit/simple_model.py
@@ -43,9 +43,16 @@ class SimpleMoEModel(torch.nn.Module):
     def __init__(self, hidden_dim, num_experts=4, ep_size=1, use_residual=False):
         super(SimpleMoEModel, self).__init__()
         self.linear = torch.nn.Linear(hidden_dim, hidden_dim)
-        linear2 = torch.nn.Linear(hidden_dim, hidden_dim)
+        expert = torch.nn.Linear(hidden_dim, hidden_dim)
+        # using two MoE layers to check implications of sharing a single storage
         self.linear2 = MoE(hidden_size=hidden_dim,
-                           expert=linear2,
+                           expert=expert,
+                           ep_size=ep_size,
+                           use_residual=use_residual,
+                           num_experts=num_experts,
+                           k=1)
+        self.linear3 = MoE(hidden_size=hidden_dim,
+                           expert=expert,
                            ep_size=ep_size,
                            use_residual=use_residual,
                            num_experts=num_experts,
@@ -53,9 +60,9 @@ def __init__(self, hidden_dim, num_experts=4, ep_size=1, use_residual=False):
         self.cross_entropy_loss = torch.nn.CrossEntropyLoss()
 
     def forward(self, x, y):
-        hidden_dim = x
-        hidden_dim = self.linear(hidden_dim)
+        hidden_dim = self.linear(x)
         output, _, _ = self.linear2(hidden_dim)
+        output, _, _ = self.linear3(output)
         hidden_dim = hidden_dim + output
         sentence_embed = hidden_dim.mean(1)
         return self.cross_entropy_loss(sentence_embed, y)