Repository: microsoft/DeepSpeed

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





Discussion:

Resolving patch:
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 @@ def get_files(dir):
     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 []
 


Test patch: