Repository: microsoft/DeepSpeed

Issue:
[BUG] Unable to reconstruct full weights from ZeRO-2 checkpoint when layers share weights
**Describe the bug**
There is an edge case in which a ZeRO-2 checkpoint is missing weights after running it through the `deepspeed.utils.zero_to_fp32.get_fp32_state_dict_from_zero_checkpoint` utility function. Layers that share weights can't be recovered, leading to a missing keys error when calling `model.load_state_dict()`.

**To Reproduce**
Save the below script into a file `repro.py`.
```py
import os
import argparse
import torch
import torch.nn as nn

import deepspeed
from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint


class ModelWithSharedWeights(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer0 = nn.Linear(100, 100)
        self.layer1 = nn.Linear(200, 200)
        self.layer2 = nn.Linear(300, 300)
        # tie layer 1 and layer 2
        self.layer1.weight = self.layer2.weight


config = {
    "train_micro_batch_size_per_gpu": 2,
    "zero_allow_untested_optimizer": True,
    "zero_optimization": {"stage": 2},
}


rank = int(os.environ["LOCAL_RANK"])

torch.cuda.set_device(rank)
model = ModelWithSharedWeights()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

deepspeed_engine, _, _, _ = deepspeed.initialize(
    args=argparse.Namespace(device_rank=rank),
    config=config,
    model=model,
    optimizer=optimizer,
)
deepspeed_engine.save_checkpoint("checkpoints_ds/checkpoint0.pt", tag="checkpoint")

if rank == 0:
    model = ModelWithSharedWeights()
    state_dict = get_fp32_state_dict_from_zero_checkpoint("checkpoints_ds/checkpoint0.pt", tag="checkpoint")

    # Fails with
    # RuntimeError: Error(s) in loading state_dict for ModelWithSharedWeights:
    # Missing key(s) in state_dict: "layer2.weight".
    model.load_state_dict(state_dict)

torch.distributed.barrier()

```
Run this command:
```bash
torchrun --nproc_per_node 2 repro.py
```
Error:
```
Traceback (most recent call last):
  File "/home/adrian/repositories/lightning/repro_ds.py", line 52, in <module>
    model.load_state_dict(state_dict)
  File "/home/adrian/.conda/envs/lightning/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2056, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for ModelWithSharedWeights:
        Missing key(s) in state_dict: "layer2.weight". 
```

**Expected behavior**
The model should load the weights without error.

**ds_report output**

```
Setting ds_accelerator to cuda (auto detect)
Setting ds_accelerator to cuda (auto detect)
--------------------------------------------------
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]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.1
 [WARNING]  using untested triton version (2.1.0+7d1a95b046), 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 ............... ['/home/adrian/.conda/envs/lightning/lib/python3.10/site-packages/torch']
torch version .................... 2.1.0.dev20230529+cu118
deepspeed install path ........... ['/home/adrian/repositories/DeepSpeed/deepspeed']
deepspeed info ................... 0.9.4+f483c034, f483c034, master
torch cuda version ............... 11.8
torch hip version ................ None
nvcc version ..................... 12.1
deepspeed wheel compiled w. ...... torch 2.1, cuda 11.8
```


**System info (please complete the following information):**
 - OS: [e.g. Ubuntu 18.04] Ubuntu 18.04
 - GPU count and types [e.g. two machines with x8 A100s each] 8x A100
 - Interconnects (if applicable) [e.g., two machines connected with 100 Gbps IB]
 - Python version 3.10
 - Any other relevant info about your setup

**Launcher context**
See command above, it is using torchrun. 

**Docker context**
N/A

**Additional context**

Originally reported in Lightning: https://github.com/Lightning-AI/lightning/issues/15694



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
@@ -3236,24 +3236,30 @@ def _get_shared_params(self):
         e.g. in `zero_to_fp32`. Each dict entry is a pair of param names, where the key is the name
         of the variable that isn't stored and the value is the actual param holding data.
         """
-        shared_ds_ids = {}
+        shared_index = {}
         shared_params_by_full_name = {}
 
+        is_zero3_model = (self.zero_optimization_partition_weights()
+                          and any(hasattr(param, "ds_id") for param in self.module.parameters()))
+
         def get_layer_state_dict(module, prefix=""):
             # handle params
             for name, param in module.named_parameters(recurse=False):
-                if param is None or not hasattr(param, "ds_id"):
+                if param is None or (is_zero3_model and not hasattr(param, "ds_id")):
                     continue
                 key = prefix + name
-                # can't rely on param.data_ptr() as it will be reused as weights gets
-                # gathered and reduced, but param.ds_id is unique across all zero weights
+
+                # When weights are manged by stage 3, we can't rely on param.data_ptr() as it will be reused
+                # as weights get gathered and reduced, but param.ds_id is unique across all zero weights
                 # (and shared params will have the same param.ds_id)
-                if param.ds_id in shared_ds_ids:
+                param_id = param.ds_id if is_zero3_model else param.data_ptr()
+
+                if param_id in shared_index:
                     # shared weights
-                    #print(f"`{key}` is shared with `{shared_ds_ids[param.ds_id]}`")
-                    shared_params_by_full_name[key] = shared_ds_ids[param.ds_id]
+                    #print(f"`{key}` is shared with `{shared_index[param_id]}`")
+                    shared_params_by_full_name[key] = shared_index[param_id]
                 else:
-                    shared_ds_ids[param.ds_id] = key
+                    shared_index[param_id] = key
 
             for name, child in module.named_children():
                 if child is not None:


Test patch:
diff --git a/tests/unit/checkpoint/test_shared_weights.py b/tests/unit/checkpoint/test_shared_weights.py
new file mode 100644
--- /dev/null
+++ b/tests/unit/checkpoint/test_shared_weights.py
@@ -0,0 +1,49 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import torch.nn as nn
+
+import deepspeed
+from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+from unit.common import DistributedTest
+
+
+class ModelWithSharedWeights(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.layer0 = nn.Linear(100, 100)
+        self.layer1 = nn.Linear(200, 200)
+        self.layer2 = nn.Linear(300, 300)
+        # tie layer 1 and layer 2
+        self.layer1.weight = self.layer2.weight
+
+
+class TestCheckpointSharedWeights(DistributedTest):
+    world_size = 2
+
+    def test_checkpoint_shared_weights(self, tmp_path):
+        config = {
+            "train_micro_batch_size_per_gpu": 2,
+            "zero_allow_untested_optimizer": True,
+            "zero_optimization": {
+                "stage": 2
+            },
+        }
+        model = ModelWithSharedWeights()
+        optimizer = torch.optim.Adam(model.parameters())
+
+        deepspeed_engine, _, _, _ = deepspeed.initialize(
+            config=config,
+            model=model,
+            optimizer=optimizer,
+        )
+        filename = tmp_path / "checkpoint.pt"
+        deepspeed_engine.save_checkpoint(filename, tag="checkpoint")
+
+        model = ModelWithSharedWeights()
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(filename, tag="checkpoint")
+        model.load_state_dict(state_dict, strict=True)