[AutoParallel] Support unbalanced split save load. (#66103)

zhangbo9674 · web-flow · commit 59eae052c3ae · 2024-07-24T14:33:55.000+08:00
* add

* add

* add

* add

* add

* add

* add

* add comment
diff --git a/python/paddle/distributed/checkpoint/load_state_dict.py b/python/paddle/distributed/checkpoint/load_state_dict.py
@@ -522,6 +522,7 @@ def load_state_dict(
                 assert (
                     item.local_tensor_index.tensor_key in flat_state_dict
                 ), f"item:{item}, state_dict:{flat_state_dict}"
+
                 cur_local_tensor = (
                     flat_state_dict[
                         item.local_tensor_index.tensor_key
@@ -532,6 +533,7 @@ def load_state_dict(
                     ].is_dist()
                     else flat_state_dict[item.local_tensor_index.tensor_key]
                 )
+
                 cur_offsets = item.cur_offset
                 cur_lengths = item.lengths
                 cur_ends = [
@@ -549,23 +551,26 @@ def load_state_dict(
                 else:
                     cur_chunk_tensor = cur_local_tensor
             else:
+                # Why we use item.dtype: In static mode, the state_dict maybe incomplete in pp, the dtype is stored in advance.
                 cur_chunk_tensor = paddle.zeros(
                     item.lengths,
                     item.dtype,
                 )
 
+            # Src_rank represents the rank of data read from ckpt, item_rank is the rank of the parameter of the data to be loaded.
             if src_rank == item.rank:
                 if src_rank == paddle.distributed.get_rank():
-                    # assign value locally
+                    # Assign value locally: in the case of src_rank is cur_rank, it means that the ckpt and the parameters to be loaded are both in the current node.
                     paddle.assign(storage_chunk_tensor, cur_chunk_tensor)
             else:
-                # assign value remotely
+                # Assign value remotely: src_rank broadcasts the ckpt, and the parameters to be loaded receive the data broadcast by src_rank.
                 if src_rank == paddle.distributed.get_rank():
                     storage_chunk_tensor = storage_chunk_tensor.contiguous()
                     paddle.distributed.broadcast(
                         storage_chunk_tensor, src=src_rank, group=process_group
                     )
                 else:
+                    # The memory hold by cur_chunk_tensor may be non-contiguous, and the broadcast API does not support this type of tensor.
                     tmp_tensor = paddle.assign(cur_chunk_tensor)
                     paddle.distributed.broadcast(
                         tmp_tensor, src=src_rank, group=process_group
@@ -574,4 +579,7 @@ def load_state_dict(
 
         for k, v in flat_state_dict.items():
             if k in state_dict_in_cpu:
-                state_dict[k] = v.cpu()
+                value = state_dict
+                for key in mapping[k]:
+                    value = value[key]
+                paddle.assign(v.cpu(), value)
diff --git a/python/paddle/distributed/checkpoint/save_state_dict.py b/python/paddle/distributed/checkpoint/save_state_dict.py
@@ -160,6 +160,9 @@ def save_state_dict(
                 if not val._is_initialized():
                     continue
                 if val.is_dist():
+                    local_tensor = val._local_value()
+                    # Note: The local_tensor must keep the same name with the original tensor. Otherwise, the StructuredToParameterName@@ mapping will be wrong.
+                    local_tensor.name = val.name
                     # when val is scalar, the shape is []
                     (
                         local_shape,
@@ -175,9 +178,11 @@ def save_state_dict(
                     )
                     if local_shape is None or global_offset is None:
                         continue
-                    local_tensor = val._local_value()
-                    # Note: The local_tensor must keep the same name with the original tensor. Otherwise, the StructuredToParameterName@@ mapping will be wrong.
-                    local_tensor.name = val.name
+                    if (
+                        paddle.distributed.get_rank()
+                        not in val.process_mesh.process_ids
+                    ):
+                        continue
                 else:
                     local_shape = tuple(val.shape)
                     global_offset = (
diff --git a/python/paddle/distributed/checkpoint/utils.py b/python/paddle/distributed/checkpoint/utils.py
@@ -35,6 +35,16 @@ def get_coordinator(mesh: Union[np.array, List[List[int]]], rank: int):
     )
 
 
+# NOTE(zhangbo): Refer to the BalancedSplit function in the reshard_utils.cc file.
+def balanced_split(total_nums, num_of_pieces):
+    has_remainder = total_nums % num_of_pieces != 0
+    result = [(total_nums + num_of_pieces - 1) // num_of_pieces] * num_of_pieces
+    if has_remainder:
+        last_value = result[-1]
+        result[-1] = last_value - (last_value * num_of_pieces - total_nums)
+    return result
+
+
 def compute_local_shape_and_global_offset(
     global_shape: List[int],
     process_mesh: core.ProcessMesh,
@@ -52,12 +62,10 @@ def compute_local_shape_and_global_offset(
             continue
         else:
             i = placement.get_dim()
-            assert (
-                global_shape[i] % process_mesh.shape[dim] == 0
-            ), f"i:{i}, global_shape[i]:{global_shape[i]}, process_mesh.shape[dim]:{process_mesh.shape[dim]}"
-            local_shape[i] = global_shape[i] // process_mesh.shape[dim]
             chunk_idx = rank_coordinator[dim]
-            global_offset[i] = chunk_idx * local_shape[i]
+            chunks = balanced_split(global_shape[i], process_mesh.shape[dim])
+            local_shape[i] = chunks[chunk_idx]
+            global_offset[i] = sum(chunks[:chunk_idx])
 
     return tuple(local_shape), tuple(global_offset)
 
diff --git a/test/auto_parallel/hybrid_strategy/semi_auto_load_state_dict.py b/test/auto_parallel/hybrid_strategy/semi_auto_load_state_dict.py
@@ -67,7 +67,9 @@ def test_load_state_dict_with_four_devices(self):
         load_state_dict(state_dict, self._ckpt_path)
         # check
         cur_rank = paddle.distributed.get_rank()
-        expect_w1 = saved_w1.split(4, axis=0)[cur_rank]
+        expect_w1 = saved_w1.split(num_or_sections=[4, 4, 4, 1], axis=0)[
+            cur_rank
+        ]
         expect_w2 = sharded_w2
         expect_state_dict = dict(
             zip(list(global_state_dict.keys()), [expect_w1, expect_w2])
@@ -90,7 +92,7 @@ def test_load_state_dict_with_two_devices(self):
         load_state_dict(state_dict, self._ckpt_path)
         # check
         cur_rank = paddle.distributed.get_rank()
-        expect_w1 = saved_w1.split(2, axis=0)[cur_rank]
+        expect_w1 = saved_w1.split(num_or_sections=[7, 6], axis=0)[cur_rank]
         expect_w2 = saved_w2.split(2, axis=1)[cur_rank]
         expect_state_dict = dict(
             zip(list(global_state_dict.keys()), [expect_w1, expect_w2])
diff --git a/test/auto_parallel/hybrid_strategy/semi_auto_save_state_dict.py b/test/auto_parallel/hybrid_strategy/semi_auto_save_state_dict.py
@@ -20,7 +20,7 @@
 
 
 def get_global_state_dict():
-    w1 = paddle.arange(32).reshape([4, 8])
+    w1 = paddle.arange(104).reshape([13, 8])
     w2 = paddle.arange(32, 36).reshape([2, 2])
     return {"w1": w1, "w2": w2}
 
