We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG] Order of partitioning in universal checkpoint
**Describe the bug**
A clear and concise description of what the bug is.

Background:
Automatic adjustment of ZeRO's optimizer state partitioning with a new world size is not currently supported.

Issues:
I use the code from PR #4516  to convert a DeepSpeed checkpoint without TP and PP (128 ranks) to another DeepSpeed checkpoint (32 ranks). In my testing (ZeRO2), I've found that merging in order isn't particularly correct.

```
def check_mp_equal_to_fp32(args):
    output_folder = "./output"

    mp_sd = torch.load(
        os.path.join(output_folder, "output", "mp_rank_00_model_states.pt"),
        map_location=torch.device("cpu"),
    )
    zero_output_folder = os.path.join(output_folder, "zero")
    tensor_name_paths = sorted(glob.glob(f"{zero_output_folder}/*"))
    for tensor_name_path in tensor_name_paths:
        if "model" not in tensor_name_path:
            continue
        tensor_name = os.path.basename(tensor_name_path)
        fp32 = torch.load(os.path.join(tensor_name_path, "fp32.pt"))["param"].to(mp_sd["module"][tensor_name])
        torch.testing.assert_allclose(fp32, mp_sd["module"][tensor_name], msg=f"{tensor_name}, fp32: \n{fp32}, mp_sd: \n{mp_sd['module'][tensor_name]}")
```

In this example, I found that it wasn't the alphabetical order or number order.


**Expected behavior**
We need to get the linked matrix for partitioning, so that we can get the correct order. Maybe something look like this:


```python
def _build_linked_matrix(args: argparse.Namespace, temp_dir: str, mp_sd: dict):
    linked_matrix = defaultdict(list)
    reversed_linked_matrix = defaultdict(list)

    for tensor_name, mp_tensor in mp_sd["module"].items():
        fp32_paths = sorted(glob.glob(f"{temp_dir}/{tensor_name}/0/fp32.*"))
        fp32_tensors = []
        for fp32_path in fp32_paths:
            partition_idx = fp32_path.split("/")[-1].split(".")[-1]
            fp32_tensor = torch.load(os.path.join(fp32_path))
            fp32_tensors.append((partition_idx, fp32_tensor))

        flat_tensor = mp_tensor.flatten()
        visited = [False for _ in range(len(fp32_tensors))]
        offset = 0
        for merge_cnt in range(len(fp32_tensors)):
            for index, (partition_idx, fp32_tensor) in enumerate(fp32_tensors):
                numel = fp32_tensor.numel()
                if visited[index] is False:
                    cloned_flat_tensor = copy.deepcopy(flat_tensor)
                    cloned_flat_tensor[offset : offset + numel] = fp32_tensor.to(
                        cloned_flat_tensor
                    )
                    if torch.allclose(flat_tensor, cloned_flat_tensor):
                        visited[index] = True
                        # tensor_name -> [(partition_idx, numel)] in temp
                        linked_matrix[tensor_name].append((partition_idx, numel))
                        # fp32.idx -> [(tensor_name, offset, numel)]
                        # it is ordered because each time we merge one partition.
                        reversed_linked_matrix[partition_idx].append(
                            (tensor_name, offset, numel)
                        )
                        offset += numel
        print(tensor_name, linked_matrix[tensor_name], mp_tensor.shape)

    # linked_matrix["model.transformer_encoder.layers.31.linear1.weight"]
    # Out:
    # [('018' 54715476), ('019', 12393388)]

    # reverse_linked_matrix['019']
    # Out:
    # [('model.transformer_encoder.layers.31.linear1.weight', 54715476, 12393388),
    #  ('model.transformer_encoder.layers.31.linear1.bias', 0, 16384),
    #  ('model.transformer_encoder.layers.31.linear2.weight', 0, 46882008)]
    # => 54715476 is the offset of 12393388 elements in layers.31.linear1.weight
    print(linked_matrix, reversed_linked_matrix)
    torch.save(
        (linked_matrix, reversed_linked_matrix),
        os.path.join(args.input_folder, "linked_matrix.pt"),
    )
    return linked_matrix, reversed_linked_matrix
```

**System info (please complete the following information):**
 - OS: Ubuntu 22.04
 - GPU count and types: A100
 - Python version: 3.10
 - Any other relevant info about your setup




--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `deepspeed/checkpoint/reshape_utils.py`
Content:
```
1 # Copyright (c) Microsoft Corporation.
2 # SPDX-License-Identifier: Apache-2.0
3 
4 # DeepSpeed Team
5 
6 import os
7 import torch
8 from collections import OrderedDict
9 from .constants import (ZERO_FILE_PREFIX, FP16_ZERO_FILE_PREFIX, BF16_ZERO_FILE_PREFIX)
10 
11 
12 def basic_folder_validation(dir):
13     assert os.path.exists(dir), f'{dir} path does not exist'
14     assert os.path.isdir(dir), f'{dir} is not a folder'
15 
16 
17 def get_files_with_prefix(all_files, prefix):
18     file_list = []
19     for file_path in all_files:
20         _, fname = os.path.split(file_path)
21         if fname.startswith(prefix):
22             file_list.append(file_path)
23 
24     return sorted(file_list)
25 
26 
27 def validate_files(file_list):
28     for file in file_list:
29         if not os.path.isfile(file):
30             print(f'Error: {file} is not existent')
31 
32 
33 def get_files(dir):
34     file_list = []
35     for root, _, files in os.walk(dir):
36         for file in files:
37             file_list.append(os.path.join(root, file))
38     return file_list
39 
40 
41 def get_zero_files(dir):
42     file_list = get_files(dir)
43     for prefix in [ZERO_FILE_PREFIX, FP16_ZERO_FILE_PREFIX, BF16_ZERO_FILE_PREFIX]:
44         zero_files = get_files_with_prefix(file_list, prefix)
45         if len(zero_files) > 0:
46             return zero_files
47 
48     return []
49 
50 
51 def partition_data(data_list, num_partitions):
52     num_elems = len(data_list)
53     assert num_elems % num_partitions == 0
54     partition_size = num_elems // num_partitions
55     partitions_list = [data_list[i:i + partition_size] for i in range(0, num_elems, partition_size)]
56     return partitions_list
57 
58 
59 def _key_list_to_string(key_list):
60     return '.'.join(key_list)
61 
62 
63 def merge_state_dict(dict_a, dict_b, key_list):
64     merged_dict = type(dict_a)({})
65 
66     for key, value in dict_b.items():
67         if key in dict_a.keys():
68             merged_dict[key] = merge_state(dict_a[key], dict_b[key], [str(key)])
69         else:
70             merged_dict[key] = value
71 
72     return merged_dict
73 
74 
75 def merge_state_list(list_a, list_b, key_list):
76     if len(list_a) != len(list_b):
77         print(f'{_key_list_to_string(key_list)}')
78         raise ValueError(f'Cannot merge lists of different lengths, a = {len(list_a)} b = {len(list_b)}')
79 
80     return [merge_state(a, b, key_list) for a, b in zip(list_a, list_b)]
81 
82 
83 def merge_state(state_a, state_b, key_list=[]):
84     if type(state_a) != type(state_b):
85         key_list_string = _key_list_to_string(key_list)
86         print(f'key_list = {key_list_string}')
87         raise ValueError(f'Cannot merge two states of types {type(state_a)} and type {type(state_b)}')
88 
89     if type(state_a) in (dict, OrderedDict):
90         return merge_state_dict(state_a, state_b, key_list)
91     elif type(state_a) in (list, tuple):
92         return type(state_a)(merge_state_list(state_a, state_b, key_list))
93     elif torch.is_tensor(state_a):
94         return torch.cat([state_a, state_b], 0)
95     else:
96         return state_a
97
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/deepspeed/checkpoint/reshape_utils.py b/deepspeed/checkpoint/reshape_utils.py
--- a/deepspeed/checkpoint/reshape_utils.py
+++ b/deepspeed/checkpoint/reshape_utils.py
@@ -4,9 +4,10 @@
 # DeepSpeed Team
 
 import os
+import re
 import torch
 from collections import OrderedDict
-from .constants import (ZERO_FILE_PREFIX, FP16_ZERO_FILE_PREFIX, BF16_ZERO_FILE_PREFIX)
+from .constants import (ZERO_FILE_PREFIX, FP16_ZERO_FILE_PREFIX, BF16_ZERO_FILE_PREFIX, MODEL_FILE_PREFIX)
 
 
 def basic_folder_validation(dir):
@@ -38,12 +39,28 @@
     return file_list
 
 
+def sort_zero_files(files, prefix):
+    pattern = f"{prefix}([0-9]+)_{MODEL_FILE_PREFIX}([0-9]+)"
+    rank_pairs = []
+    for f in files:
+        m = re.search(pattern, f)
+        if m:
+            dp_rank = int(m.group(1))
+            mp_rank = int(m.group(2))
+            rank_pairs.append((dp_rank, mp_rank, f))
+        else:
+            raise ValueError(f"Cannot parse dp_rank and mp_rank from {f}")
+
+    sorted_files = sorted(rank_pairs, key=lambda x: (x[0], x[1]))
+    return [f for _, _, f in sorted_files]
+
+
 def get_zero_files(dir):
     file_list = get_files(dir)
     for prefix in [ZERO_FILE_PREFIX, FP16_ZERO_FILE_PREFIX, BF16_ZERO_FILE_PREFIX]:
         zero_files = get_files_with_prefix(file_list, prefix)
         if len(zero_files) > 0:
-            return zero_files
+            return sort_zero_files(zero_files, prefix)
 
     return []

{"golden_diff": "diff --git a/deepspeed/checkpoint/reshape_utils.py b/deepspeed/checkpoint/reshape_utils.py\n--- a/deepspeed/checkpoint/reshape_utils.py\n+++ b/deepspeed/checkpoint/reshape_utils.py\n@@ -4,9 +4,10 @@\n # DeepSpeed Team\n \n import os\n+import re\n import torch\n from collections import OrderedDict\n-from .constants import (ZERO_FILE_PREFIX, FP16_ZERO_FILE_PREFIX, BF16_ZERO_FILE_PREFIX)\n+from .constants import (ZERO_FILE_PREFIX, FP16_ZERO_FILE_PREFIX, BF16_ZERO_FILE_PREFIX, MODEL_FILE_PREFIX)\n \n \n def basic_folder_validation(dir):\n@@ -38,12 +39,28 @@\n     return file_list\n \n \n+def sort_zero_files(files, prefix):\n+    pattern = f\"{prefix}([0-9]+)_{MODEL_FILE_PREFIX}([0-9]+)\"\n+    rank_pairs = []\n+    for f in files:\n+        m = re.search(pattern, f)\n+        if m:\n+            dp_rank = int(m.group(1))\n+            mp_rank = int(m.group(2))\n+            rank_pairs.append((dp_rank, mp_rank, f))\n+        else:\n+            raise ValueError(f\"Cannot parse dp_rank and mp_rank from {f}\")\n+\n+    sorted_files = sorted(rank_pairs, key=lambda x: (x[0], x[1]))\n+    return [f for _, _, f in sorted_files]\n+\n+\n def get_zero_files(dir):\n     file_list = get_files(dir)\n     for prefix in [ZERO_FILE_PREFIX, FP16_ZERO_FILE_PREFIX, BF16_ZERO_FILE_PREFIX]:\n         zero_files = get_files_with_prefix(file_list, prefix)\n         if len(zero_files) > 0:\n-            return zero_files\n+            return sort_zero_files(zero_files, prefix)\n \n     return []\n", "issue": "[BUG] Order of partitioning in universal checkpoint\n**Describe the bug**\r\nA clear and concise description of what the bug is.\r\n\r\nBackground:\r\nAutomatic adjustment of ZeRO's optimizer state partitioning with a new world size is not currently supported.\r\n\r\nIssues:\r\nI use the code from PR #4516  to convert a DeepSpeed checkpoint without TP and PP (128 ranks) to another DeepSpeed checkpoint (32 ranks). In my testing (ZeRO2), I've found that merging in order isn't particularly correct.\r\n\r\n```\r\ndef check_mp_equal_to_fp32(args):\r\n    output_folder = \"./output\"\r\n\r\n    mp_sd = torch.load(\r\n        os.path.join(output_folder, \"output\", \"mp_rank_00_model_states.pt\"),\r\n        map_location=torch.device(\"cpu\"),\r\n    )\r\n    zero_output_folder = os.path.join(output_folder, \"zero\")\r\n    tensor_name_paths = sorted(glob.glob(f\"{zero_output_folder}/*\"))\r\n    for tensor_name_path in tensor_name_paths:\r\n        if \"model\" not in tensor_name_path:\r\n            continue\r\n        tensor_name = os.path.basename(tensor_name_path)\r\n        fp32 = torch.load(os.path.join(tensor_name_path, \"fp32.pt\"))[\"param\"].to(mp_sd[\"module\"][tensor_name])\r\n        torch.testing.assert_allclose(fp32, mp_sd[\"module\"][tensor_name], msg=f\"{tensor_name}, fp32: \\n{fp32}, mp_sd: \\n{mp_sd['module'][tensor_name]}\")\r\n```\r\n\r\nIn this example, I found that it wasn't the alphabetical order or number order.\r\n\r\n\r\n**Expected behavior**\r\nWe need to get the linked matrix for partitioning, so that we can get the correct order. Maybe something look like this:\r\n\r\n\r\n```python\r\ndef _build_linked_matrix(args: argparse.Namespace, temp_dir: str, mp_sd: dict):\r\n    linked_matrix = defaultdict(list)\r\n    reversed_linked_matrix = defaultdict(list)\r\n\r\n    for tensor_name, mp_tensor in mp_sd[\"module\"].items():\r\n        fp32_paths = sorted(glob.glob(f\"{temp_dir}/{tensor_name}/0/fp32.*\"))\r\n        fp32_tensors = []\r\n        for fp32_path in fp32_paths:\r\n            partition_idx = fp32_path.split(\"/\")[-1].split(\".\")[-1]\r\n            fp32_tensor = torch.load(os.path.join(fp32_path))\r\n            fp32_tensors.append((partition_idx, fp32_tensor))\r\n\r\n        flat_tensor = mp_tensor.flatten()\r\n        visited = [False for _ in range(len(fp32_tensors))]\r\n        offset = 0\r\n        for merge_cnt in range(len(fp32_tensors)):\r\n            for index, (partition_idx, fp32_tensor) in enumerate(fp32_tensors):\r\n                numel = fp32_tensor.numel()\r\n                if visited[index] is False:\r\n                    cloned_flat_tensor = copy.deepcopy(flat_tensor)\r\n                    cloned_flat_tensor[offset : offset + numel] = fp32_tensor.to(\r\n                        cloned_flat_tensor\r\n                    )\r\n                    if torch.allclose(flat_tensor, cloned_flat_tensor):\r\n                        visited[index] = True\r\n                        # tensor_name -> [(partition_idx, numel)] in temp\r\n                        linked_matrix[tensor_name].append((partition_idx, numel))\r\n                        # fp32.idx -> [(tensor_name, offset, numel)]\r\n                        # it is ordered because each time we merge one partition.\r\n                        reversed_linked_matrix[partition_idx].append(\r\n                            (tensor_name, offset, numel)\r\n                        )\r\n                        offset += numel\r\n        print(tensor_name, linked_matrix[tensor_name], mp_tensor.shape)\r\n\r\n    # linked_matrix[\"model.transformer_encoder.layers.31.linear1.weight\"]\r\n    # Out:\r\n    # [('018' 54715476), ('019', 12393388)]\r\n\r\n    # reverse_linked_matrix['019']\r\n    # Out:\r\n    # [('model.transformer_encoder.layers.31.linear1.weight', 54715476, 12393388),\r\n    #  ('model.transformer_encoder.layers.31.linear1.bias', 0, 16384),\r\n    #  ('model.transformer_encoder.layers.31.linear2.weight', 0, 46882008)]\r\n    # => 54715476 is the offset of 12393388 elements in layers.31.linear1.weight\r\n    print(linked_matrix, reversed_linked_matrix)\r\n    torch.save(\r\n        (linked_matrix, reversed_linked_matrix),\r\n        os.path.join(args.input_folder, \"linked_matrix.pt\"),\r\n    )\r\n    return linked_matrix, reversed_linked_matrix\r\n```\r\n\r\n**System info (please complete the following information):**\r\n - OS: Ubuntu 22.04\r\n - GPU count and types: A100\r\n - Python version: 3.10\r\n - Any other relevant info about your setup\r\n\r\n\r\n\n", "before_files": [{"content": "# Copyright (c) Microsoft Corporation.\n# SPDX-License-Identifier: Apache-2.0\n\n# DeepSpeed Team\n\nimport os\nimport torch\nfrom collections import OrderedDict\nfrom .constants import (ZERO_FILE_PREFIX, FP16_ZERO_FILE_PREFIX, BF16_ZERO_FILE_PREFIX)\n\n\ndef basic_folder_validation(dir):\n    assert os.path.exists(dir), f'{dir} path does not exist'\n    assert os.path.isdir(dir), f'{dir} is not a folder'\n\n\ndef get_files_with_prefix(all_files, prefix):\n    file_list = []\n    for file_path in all_files:\n        _, fname = os.path.split(file_path)\n        if fname.startswith(prefix):\n            file_list.append(file_path)\n\n    return sorted(file_list)\n\n\ndef validate_files(file_list):\n    for file in file_list:\n        if not os.path.isfile(file):\n            print(f'Error: {file} is not existent')\n\n\ndef get_files(dir):\n    file_list = []\n    for root, _, files in os.walk(dir):\n        for file in files:\n            file_list.append(os.path.join(root, file))\n    return file_list\n\n\ndef get_zero_files(dir):\n    file_list = get_files(dir)\n    for prefix in [ZERO_FILE_PREFIX, FP16_ZERO_FILE_PREFIX, BF16_ZERO_FILE_PREFIX]:\n        zero_files = get_files_with_prefix(file_list, prefix)\n        if len(zero_files) > 0:\n            return zero_files\n\n    return []\n\n\ndef partition_data(data_list, num_partitions):\n    num_elems = len(data_list)\n    assert num_elems % num_partitions == 0\n    partition_size = num_elems // num_partitions\n    partitions_list = [data_list[i:i + partition_size] for i in range(0, num_elems, partition_size)]\n    return partitions_list\n\n\ndef _key_list_to_string(key_list):\n    return '.'.join(key_list)\n\n\ndef merge_state_dict(dict_a, dict_b, key_list):\n    merged_dict = type(dict_a)({})\n\n    for key, value in dict_b.items():\n        if key in dict_a.keys():\n            merged_dict[key] = merge_state(dict_a[key], dict_b[key], [str(key)])\n        else:\n            merged_dict[key] = value\n\n    return merged_dict\n\n\ndef merge_state_list(list_a, list_b, key_list):\n    if len(list_a) != len(list_b):\n        print(f'{_key_list_to_string(key_list)}')\n        raise ValueError(f'Cannot merge lists of different lengths, a = {len(list_a)} b = {len(list_b)}')\n\n    return [merge_state(a, b, key_list) for a, b in zip(list_a, list_b)]\n\n\ndef merge_state(state_a, state_b, key_list=[]):\n    if type(state_a) != type(state_b):\n        key_list_string = _key_list_to_string(key_list)\n        print(f'key_list = {key_list_string}')\n        raise ValueError(f'Cannot merge two states of types {type(state_a)} and type {type(state_b)}')\n\n    if type(state_a) in (dict, OrderedDict):\n        return merge_state_dict(state_a, state_b, key_list)\n    elif type(state_a) in (list, tuple):\n        return type(state_a)(merge_state_list(state_a, state_b, key_list))\n    elif torch.is_tensor(state_a):\n        return torch.cat([state_a, state_b], 0)\n    else:\n        return state_a\n", "path": "deepspeed/checkpoint/reshape_utils.py"}], "after_files": [{"content": "# Copyright (c) Microsoft Corporation.\n# SPDX-License-Identifier: Apache-2.0\n\n# DeepSpeed Team\n\nimport os\nimport re\nimport torch\nfrom collections import OrderedDict\nfrom .constants import (ZERO_FILE_PREFIX, FP16_ZERO_FILE_PREFIX, BF16_ZERO_FILE_PREFIX, MODEL_FILE_PREFIX)\n\n\ndef basic_folder_validation(dir):\n    assert os.path.exists(dir), f'{dir} path does not exist'\n    assert os.path.isdir(dir), f'{dir} is not a folder'\n\n\ndef get_files_with_prefix(all_files, prefix):\n    file_list = []\n    for file_path in all_files:\n        _, fname = os.path.split(file_path)\n        if fname.startswith(prefix):\n            file_list.append(file_path)\n\n    return sorted(file_list)\n\n\ndef validate_files(file_list):\n    for file in file_list:\n        if not os.path.isfile(file):\n            print(f'Error: {file} is not existent')\n\n\ndef get_files(dir):\n    file_list = []\n    for root, _, files in os.walk(dir):\n        for file in files:\n            file_list.append(os.path.join(root, file))\n    return file_list\n\n\ndef sort_zero_files(files, prefix):\n    pattern = f\"{prefix}([0-9]+)_{MODEL_FILE_PREFIX}([0-9]+)\"\n    rank_pairs = []\n    for f in files:\n        m = re.search(pattern, f)\n        if m:\n            dp_rank = int(m.group(1))\n            mp_rank = int(m.group(2))\n            rank_pairs.append((dp_rank, mp_rank, f))\n        else:\n            raise ValueError(f\"Cannot parse dp_rank and mp_rank from {f}\")\n\n    sorted_files = sorted(rank_pairs, key=lambda x: (x[0], x[1]))\n    return [f for _, _, f in sorted_files]\n\n\ndef get_zero_files(dir):\n    file_list = get_files(dir)\n    for prefix in [ZERO_FILE_PREFIX, FP16_ZERO_FILE_PREFIX, BF16_ZERO_FILE_PREFIX]:\n        zero_files = get_files_with_prefix(file_list, prefix)\n        if len(zero_files) > 0:\n            return sort_zero_files(zero_files, prefix)\n\n    return []\n\n\ndef partition_data(data_list, num_partitions):\n    num_elems = len(data_list)\n    assert num_elems % num_partitions == 0\n    partition_size = num_elems // num_partitions\n    partitions_list = [data_list[i:i + partition_size] for i in range(0, num_elems, partition_size)]\n    return partitions_list\n\n\ndef _key_list_to_string(key_list):\n    return '.'.join(key_list)\n\n\ndef merge_state_dict(dict_a, dict_b, key_list):\n    merged_dict = type(dict_a)({})\n\n    for key, value in dict_b.items():\n        if key in dict_a.keys():\n            merged_dict[key] = merge_state(dict_a[key], dict_b[key], [str(key)])\n        else:\n            merged_dict[key] = value\n\n    return merged_dict\n\n\ndef merge_state_list(list_a, list_b, key_list):\n    if len(list_a) != len(list_b):\n        print(f'{_key_list_to_string(key_list)}')\n        raise ValueError(f'Cannot merge lists of different lengths, a = {len(list_a)} b = {len(list_b)}')\n\n    return [merge_state(a, b, key_list) for a, b in zip(list_a, list_b)]\n\n\ndef merge_state(state_a, state_b, key_list=[]):\n    if type(state_a) != type(state_b):\n        key_list_string = _key_list_to_string(key_list)\n        print(f'key_list = {key_list_string}')\n        raise ValueError(f'Cannot merge two states of types {type(state_a)} and type {type(state_b)}')\n\n    if type(state_a) in (dict, OrderedDict):\n        return merge_state_dict(state_a, state_b, key_list)\n    elif type(state_a) in (list, tuple):\n        return type(state_a)(merge_state_list(state_a, state_b, key_list))\n    elif torch.is_tensor(state_a):\n        return torch.cat([state_a, state_b], 0)\n    else:\n        return state_a\n", "path": "deepspeed/checkpoint/reshape_utils.py"}]}