[FlexCheckPoint]adapt fc to sharding stage3 (#76538)

* adapt fc to sharding stage3

* add test and fix bug

* fix bug

* fix bug

* add test

Author: zty-king

python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_stage3.py

@@ -24,6 +24,11 @@
 from paddle.autograd import PyLayer
 from paddle.base.framework import EagerParamBase
 from paddle.distributed import collective
+from paddle.distributed.flex_checkpoint.dcp.sharded_weight import (
+    ShardedStateDict,
+    ShardedWeight,
+    create_sharded_weight_with_new_local,
+)
 from paddle.framework import core
 from paddle.nn import ClipGradByGlobalNorm
 
@@ -182,7 +187,11 @@ def __init__(
             "Multiple optimizers are not supported now."
         )
         self._optim = _OptimizerWrapper(
-            optimizer, self._offload, self._group, self._update_params_slice
+            optimizer,
+            self._offload,
+            self._group,
+            self._update_params_slice,
+            self._sharded_state_dict,
         )
         self._ori_parameter_list = self._optim._parameter_list
         self._ori_param_groups = self._optim._param_groups
@@ -850,6 +859,193 @@ def _opt_clear(self):
 
         self._optim.clear_grad = MethodType(_opt_clear, self._optim)
 
+    def init_slice_param(self):
+        for layer_id, params in self._trainable_params.items():
+            for param in params:
+                value = paddle.zeros(param.shape, dtype=param.dtype)
+                value._share_buffer_to(param)
+
+    def align_param_to_buffer_and_clear_slice_param(self):
+        for layer_id, params in self._trainable_params.items():
+            for param in params:
+                param_shape = param.shape
+                origin_state = param.stop_gradient
+                param.stop_gradient = True
+                start, end = self._param2buffer[param.name][self._rank]
+                param.flatten_()
+                param.stop_gradient = origin_state
+                param_numel = param.numel().item()
+                start = min(start, param_numel)
+                end = min(end, param_numel)
+                if end > start:
+                    tmp_tensor = param._slice(start, end).detach()
+                    buffer_slice = param.fw_storage._slice(
+                        0, end - start
+                    ).detach()
+                    buffer_slice.set_value(tmp_tensor)
+                    del buffer_slice
+                param.get_tensor()._set_dims(param_shape)
+                param._clear_data()
+
+    def init_optimizer_for_slice_param(self):
+        local_param_list = []
+        for param in self._optim._parameter_list:
+            if hasattr(param, "fw_storage"):
+                var = param.fw_storage
+                tmp_param = EagerParamBase(
+                    shape=var.shape, dtype=var.dtype, name="slice@" + param.name
+                )
+                local_param_list.append(tmp_param)
+            else:
+                local_param_list.append(param)
+        self._optim._parameter_list = local_param_list
+
+    def _sharded_state_dict(
+        self,
+        model_sharded_state_dict: ShardedStateDict,
+    ) -> ShardedStateDict:
+        """
+        Convert optimizer state dict to a sharded state dict based on model sharding information.
+
+        Args:
+            model_sharded_state_dict (dict): Sharded state dict of the model, containing tensor metadata.
+
+        Returns:
+            dict: A new optimizer state dict where weights are wrapped as ShardedWeight.
+        """
+
+        _FP32_MASTER = "fp32_master_0"
+        _MOMENT_NAME = "moment"
+        _optimizer_scalar_name = [
+            "beta1_pow_acc_0",
+            "beta2_pow_acc_0",
+        ]
+        _optimizer_non_scaler_name = [
+            "moment1_0",
+            "moment2_0",
+            "velocity_0",
+        ]
+
+        param_to_slice = {}
+        for param in self._ori_parameter_list:
+            if hasattr(param, "fw_storage"):
+                param_to_slice[param.name] = True
+            else:
+                param_to_slice[param.name] = False
+
+        def _create_sharded_weight(
+            unified_name, tensor, sharded_param, static_name
+        ):
+            if param_to_slice[static_name]:
+                padding_begin = sharded_param.local_tensor.numel().item()
+                slice_begin = min(padding_begin, self._rank * tensor.shape[0])
+                slice_end = min(
+                    padding_begin, (self._rank + 1) * tensor.shape[0]
+                )
+                if slice_begin == padding_begin or slice_end == padding_begin:
+                    local_tensor = paddle.slice(
+                        tensor,
+                        axes=[0],
+                        starts=[0],
+                        ends=[slice_end - slice_begin],
+                    )
+                else:
+                    local_tensor = tensor
+                return ShardedWeight(
+                    key=unified_name,
+                    local_tensor=local_tensor,
+                    local_shape=sharded_param.local_shape,
+                    global_shape=sharded_param.global_shape,
+                    global_offset=sharded_param.global_offset,
+                    is_flattened=True,
+                    flattened_range=slice(slice_begin, slice_end),
+                )
+            else:
+                return create_sharded_weight_with_new_local(
+                    unified_name, tensor, sharded_param
+                )
+
+        def _generate_base_static_name(vname):
+            if _FP32_MASTER in vname:
+                return tuple(vname.split("_" + _FP32_MASTER + "_", 1))
+            for name in _optimizer_scalar_name + _optimizer_non_scaler_name:
+                if vname.endswith(name):
+                    return vname[: -(len(name) + 1)], name
+            raise ValueError(f"Cannot split variable name: {vname}.")
+
+        optimizer_sharded_state_dict = {}
+        optimizer_state_dict = self._optim.state_dict()
+        # Build name mapping and remove non-tensor entries from optimizer state
+        static_to_struct_mapping = {}
+        model_sharded_state_dict = dict(
+            sorted(model_sharded_state_dict.items())
+        )
+        for k, v in model_sharded_state_dict.items():
+            # When shared weights exist, the v.local_tensor.name of shared parameters are identical, but only the first parameter has optimizer states. Therefore, only the key-value pairs of the first occurrence in the shared parameter group need to be retained.
+            if v.local_tensor.name not in static_to_struct_mapping:
+                static_to_struct_mapping[v.local_tensor.name] = k
+
+        master_weights = optimizer_state_dict.pop("master_weights", None)
+        optimizer_state_dict.pop("LR_Scheduler", None)
+        # Process main optimizer states
+        for key, tensor in optimizer_state_dict.items():
+            static_name, optim_state_type = _generate_base_static_name(key)
+            static_name = static_name.replace("slice@", "")
+            struct_name = static_to_struct_mapping[static_name]
+            sharded_weight = model_sharded_state_dict[struct_name]
+
+            unified_name = f"{struct_name}.{optim_state_type}"
+
+            # Determine tensor partitioning scheme
+            if _MOMENT_NAME in optim_state_type:
+                if tensor.is_dist():
+                    optimizer_sharded_state_dict[unified_name] = ShardedWeight(
+                        key=unified_name,
+                        local_tensor=tensor,
+                        local_shape=tensor.shape,
+                        global_shape=tensor.shape,
+                        global_offset=sharded_weight.global_offset,
+                    )
+                else:
+                    optimizer_sharded_state_dict[unified_name] = (
+                        _create_sharded_weight(
+                            unified_name, tensor, sharded_weight, static_name
+                        )
+                    )
+
+            else:  # Non-momentum parameters
+                optimizer_sharded_state_dict[unified_name] = ShardedWeight(
+                    key=unified_name,
+                    local_tensor=tensor,
+                    local_shape=(1,),
+                    global_shape=(1,),
+                    global_offset=(0,),
+                )
+
+        # Process master weights if using mixed precision
+        if master_weights is not None:
+            for key, tensor in master_weights.items():
+                key = key.replace("slice@", "")
+                struct_name = static_to_struct_mapping[key]
+                sharded_weight = model_sharded_state_dict[struct_name]
+                unified_name = f"{struct_name}.w_0"
+                if tensor.is_dist():
+                    optimizer_sharded_state_dict[unified_name] = ShardedWeight(
+                        key=unified_name,
+                        local_tensor=tensor,
+                        local_shape=tensor.shape,
+                        global_shape=tensor.shape,
+                        global_offset=sharded_weight.global_offset,
+                    )
+                else:
+                    optimizer_sharded_state_dict[unified_name] = (
+                        _create_sharded_weight(
+                            unified_name, tensor, sharded_weight, key
+                        )
+                    )
+
+        return optimizer_sharded_state_dict
+
 
 def ForwardPreHooks(
     layer,
@@ -1200,13 +1396,16 @@ def _TensorWrapper(param):
     return tmp_param
 
 
-def _OptimizerWrapper(optimizer, offload, group, update_params_slice):
+def _OptimizerWrapper(
+    optimizer, offload, group, update_params_slice, sharded_state_dict
+):
     if not hasattr(optimizer, "_optim"):
         optimizer._optim = optimizer
         optimizer.offload = offload
         optimizer._group = group
         optimizer.update_scaler = None
         optimizer.update_slice = update_params_slice
+        optimizer.sharded_state_dict = sharded_state_dict
     return optimizer
 
 

python/paddle/distributed/flex_checkpoint/dcp/load_state_dict.py

@@ -716,7 +716,7 @@ def _handle_aoa(
     src_desc_to_postprocess_list = {}
     force_gc = []
 
-    for param_name, tgt_shard in load_dict.items():
+    for param_name, tgt_shard in sorted(load_dict.items()):
         tgt_desc = build_shard_desc(tgt_shard)
         shard_mappings = aoa_engine.find_shard_sources(tgt_desc)
         for mapping in shard_mappings:

test/flex_checkpoint/sharded_state_dict_logic.py

@@ -30,6 +30,9 @@
 from paddle.distributed.fleet.meta_parallel.sharding.group_sharded_optimizer_stage2 import (
     GroupShardedOptimizerStage2,
 )
+from paddle.distributed.fleet.meta_parallel.sharding.group_sharded_stage3 import (
+    GroupShardedStage3,
+)
 from paddle.distributed.fleet.utils.sequence_parallel_utils import (
     ColumnSequenceParallelLinear,
     RowSequenceParallelLinear,
@@ -410,6 +413,116 @@ def run_optimizer_test(self):
                         assert tuple(
                             opt_sharded_state_dict[opt__var_name].global_offset
                         ) == tuple(value.global_offset)
+
+        elif self.layer_type == "GroupShardedStage3":
+            model = fleet.distributed_model(model)
+            wrapped_model = GroupShardedStage3(
+                model, opt, segment_size=2**12
+            )  # slice the linear1、linear2 weight
+            for param in opt._parameter_list:
+                if hasattr(param, "fw_storage"):
+                    assert len(param.shape) != 1
+            wrapped_model.init_optimizer_for_slice_param()
+            for param in opt._parameter_list:
+                if hasattr(param, "fw_storage"):
+                    assert len(param.shape) == 1
+            model_sharded_state_dict = model.sharded_state_dict()
+            for k, v in model_sharded_state_dict.items():
+                if (
+                    k == "_layers.linear1.weight"
+                    or k == "_layers.linear2.weight"
+                ):
+                    assert not v.local_tensor._is_initialized()
+            wrapped_model.init_slice_param()
+            for k, v in model_sharded_state_dict.items():
+                if (
+                    k == "_layers.linear1.weight"
+                    or k == "_layers.linear2.weight"
+                ):
+                    assert v.local_tensor._is_initialized()
+            wrapped_model.align_param_to_buffer_and_clear_slice_param()
+            for k, v in model_sharded_state_dict.items():
+                if (
+                    k == "_layers.linear1.weight"
+                    or k == "_layers.linear2.weight"
+                ):
+                    assert not v.local_tensor._is_initialized()
+            model.train()
+            x = paddle.randint(
+                low=0,
+                high=self.vocab_size,
+                shape=[self.batch_size, self.seq_len, self.hidden_size],
+                dtype='int64',
+            )
+            rank = paddle.distributed.get_rank()
+            sharidng_x = (
+                x[0 : self.batch_size // 2]
+                if rank == 0
+                else x[self.batch_size // 2 :]
+            )
+            y = model(sharidng_x).mean()
+            y.backward()
+            opt.step()
+            opt.clear_grad()
+            model_sharded_state_dict = model.sharded_state_dict()
+            for k, v in model_sharded_state_dict.items():
+                if (
+                    k == "_layers.linear1.weight"
+                    or k == "_layers.linear2.weight"
+                ):
+                    assert not v.local_tensor._is_initialized()
+            wrapped_model.get_all_parameters()
+            opt_sharded_state_dict = opt.sharded_state_dict(
+                model_sharded_state_dict
+            )
+
+            for k, v in model_sharded_state_dict.items():
+                if (
+                    k == "_layers.linear1.weight"
+                    or k == "_layers.linear2.weight"
+                ):
+                    assert v.local_tensor._is_initialized()
+
+            for key, value in model_sharded_state_dict.items():
+                for state_name in self.optimizer_var_suffix:
+                    opt__var_name = key + state_name
+                    if opt__var_name in opt_sharded_state_dict:
+                        if hasattr(
+                            value.local_tensor, "fw_storage"
+                        ):  # check the optimizer_var which isFragment
+                            opt_var_globle_flattened_range = []
+                            paddle.distributed.all_gather_object(
+                                opt_var_globle_flattened_range,
+                                opt_sharded_state_dict[
+                                    opt__var_name
+                                ].flattened_range,
+                            )
+
+                            first_fragment = opt_var_globle_flattened_range[0]
+                            second_fragment = opt_var_globle_flattened_range[1]
+                            assert (
+                                first_fragment.stop == second_fragment.start
+                            )  # the first_flattened_range_stop == the second_flattened_range_start
+                            opt_var_globle_size_flattened = (
+                                second_fragment.stop - first_fragment.start
+                            )
+                            model_var_globle_size_flattened = math.prod(
+                                value.local_shape
+                            )
+                            assert (
+                                opt_var_globle_size_flattened
+                                == model_var_globle_size_flattened
+                            )
+
+                        assert tuple(
+                            opt_sharded_state_dict[opt__var_name].local_shape
+                        ) == tuple(value.local_shape)
+                        assert tuple(
+                            opt_sharded_state_dict[opt__var_name].global_shape
+                        ) == tuple(value.global_shape)
+                        assert tuple(
+                            opt_sharded_state_dict[opt__var_name].global_offset
+                        ) == tuple(value.global_offset)
         else:
             raise ValueError(f"Unknown layer_type: {self.layer_type}")
 

test/flex_checkpoint/test_sharded_state_dict.py

@@ -141,6 +141,15 @@
             "has_bias": "True",
             "master_weight": "True",
         },
+        {
+            "test_type": "optimizer",
+            "layer_type": "GroupShardedStage3",
+            "world_size": 2,
+            "tp": 1,
+            "sharding_degree": 2,
+            "has_bias": "True",
+            "master_weight": "True",
+        },
     ],
     "4_card_tests": [
         {