# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import copy
import json
import os
from collections import OrderedDict

import numpy
import paddle
import paddle.distributed as dist
from paddle.distributed import fleet
from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer import (
    DygraphShardingOptimizer,
)

try:
    from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
        DygraphShardingOptimizerV2,
    )
except:
    DygraphShardingOptimizerV2 = None

from paddlenlp.transformers.model_utils import (
    _add_variant,
    get_parameter_dtype,
    unwrap_optimizer,
)
from paddlenlp.transformers.utils import paddlenlp_load
from paddlenlp.utils.env import MODEL_META_NAME, SHARDING_META_NAME
from paddlenlp.utils.log import logger
from paddlenlp.utils.tools import get_env_device

from . import reshard as reshard_util
from .reshard import (
    SHARDING_STRATEGY_V1,
    SHARDING_STRATEGY_V2,
    get_param_sharding_group,
    merge_model_state,
    merge_opt_state,
    pp_reshard,
    split_model_state,
    split_opt_state,
    split_structure_name_mapping,
)


def to_device(tensor, place=None):
    if place is None:
        place = get_env_device()

    if isinstance(place, str):
        place = paddle.device._convert_to_place(place)

    if not tensor.place._equals(place):
        new_t = tensor._copy_to(place, True)
        dst_tensor = tensor.value().get_tensor()
        src_tensor = new_t.value().get_tensor()
        dst_tensor._share_data_with(src_tensor)

    return tensor


def filter_sharded_params(state_dict, optimizer, sharding_group, include_freeze_params=False):

    sharding_rank = max(sharding_group.rank, 0)
    sharding_world_size = sharding_group.nranks
    from paddlenlp.trainer.utils import reshard as reshard_util

    logger.info(f"filter sharded_params not placed in sharding_rank {sharding_rank} .")
    if not reshard_util.is_sharding_opt(optimizer):
        return state_dict

    filtered_state_dict = OrderedDict()
    if reshard_util.get_sharding_strategy(optimizer) == reshard_util.SHARDING_STRATEGY_V1:
        optimizer = unwrap_optimizer(optimizer, DygraphShardingOptimizer)
        for (k, v) in state_dict.items():
            if v.name in optimizer._param2rank:
                sharded_rank = optimizer._param2rank[v.name]
                if sharded_rank != sharding_rank:
                    continue
                filtered_state_dict[k] = v
            elif include_freeze_params:
                if sharding_rank == 0:
                    filtered_state_dict[k] = v
    else:
        optimizer = unwrap_optimizer(optimizer, DygraphShardingOptimizerV2)
        parameters = optimizer._parameter_list
        filtered_parameters = [p.name for (i, p) in enumerate(parameters) if i % sharding_world_size == sharding_rank]
        filtered_parameters = set(filtered_parameters)
        for (k, v) in state_dict.items():
            if v.name in filtered_parameters:
                filtered_state_dict[k] = v
            elif include_freeze_params and (v.name not in [p.name for p in parameters]):
                if sharding_rank == 0:
                    filtered_state_dict[k] = v
    return filtered_state_dict


def exclude_parameters_in_state_dict(
    model_state_dict, param_names_in_master_weights, sharding_group, should_save_sharding_stage1_model=True
):
    assert sharding_group is not None
    assert isinstance(model_state_dict, dict) and isinstance(
        param_names_in_master_weights, (list, set)
    ), "param_names_in_master_weights type:{}".format(type(param_names_in_master_weights))
    state_param_names = [v.name for k, v in model_state_dict.items()]
    logger.debug(
        "param_names_in_master_weights:{}, state_param_names:{}".format(
            param_names_in_master_weights, state_param_names
        )
    )
    # allgather parameter names in sharding group
    tmp = []
    if sharding_group.nranks > 1:
        paddle.distributed.all_gather_object(tmp, param_names_in_master_weights, group=sharding_group)
    else:
        tmp = [param_names_in_master_weights]
    param_names_in_master_weights = set([v for item in tmp for v in item])
    logger.info("sharding_group_param_names:{}".format(param_names_in_master_weights))
    non_parameters_state_dict = copy.copy(model_state_dict)
    for k, v in model_state_dict.items():
        if v.name in param_names_in_master_weights:
            non_parameters_state_dict.pop(k)

    return non_parameters_state_dict


class ParameterNameRemapper:
    def __init__(self, old_mapping, new_mapping, checkpoint):
        self.checkpoint = checkpoint
        self.p_name_map = {}
        for k, v in old_mapping.items():
            assert k in new_mapping, f"structure name not found: {k} {new_mapping.keys()}"
            new_v = new_mapping[k]
            if v not in self.p_name_map:
                self.p_name_map[v] = new_v
            else:
                old_v = self.p_name_map[v]
                assert old_v == new_v, f"structure name {k} has different parameter name {new_v} {self.p_name_map[v]}"
        self.old_p_names = list(self.p_name_map.keys())
        self.new_mapping = dict([(k, new_mapping[k]) for k in old_mapping.keys()])

    def _map_tensor(self, tensor, old_p_name=None):
        if old_p_name is not None:
            new_p_name = self.p_name_map.get(old_p_name)
            assert new_p_name is not None, f"parameter name {old_p_name} not found"
        else:
            new_p_name = None

        def _map_name(old_name):
            if new_p_name is not None:
                assert old_name.startswith(old_p_name)
                return new_p_name + old_name[len(old_p_name) :]
            else:
                new_name = self.p_name_map.get(old_name)
                assert new_name is not None, f"parameter name {old_name} not found"
                return new_name

        if isinstance(tensor, paddle.Tensor):
            new_name = _map_name(tensor.name)
            tensor.name = new_name
            return new_name, tensor
        else:
            assert isinstance(tensor, (list, tuple)), type(tensor)
            old_name, value = tensor
            new_name = self._map_name(old_name)
            return new_name, (new_name, value)

    def remap_model_state(self, model_state):
        for k, v in model_state.items():
            if not isinstance(v, numpy.ndarray):
                model_state[k] = self._map_tensor(v)[1]
        return model_state

    def remap_optimizer_state(self, opt_state):
        lr_scheduler_key = "LR_Scheduler"
        master_weight_key = "master_weights"

        new_opt_state = {}
        new_master_weights = None
        opt_names = []
        for k, v in opt_state.items():
            if k == lr_scheduler_key:
                new_opt_state[k] = v
            elif k == master_weight_key:
                for kk, vv in v.items():
                    new_kk = self.p_name_map[kk]
                    if new_master_weights is None:
                        new_opt_state[master_weight_key] = {}
                        new_master_weights = new_opt_state[master_weight_key]
                    new_master_weights[new_kk] = self._map_tensor(vv, kk)[1]
            else:
                assert isinstance(v, paddle.Tensor), type(v)
                opt_names.append(v.name)

        opt_to_pname = reshard_util.convert_opt_name_to_tname(self.old_p_names, opt_names)
        for opt_name in opt_names:
            v = opt_state[opt_name]
            new_opt_name, new_v = self._map_tensor(v, opt_to_pname[opt_name])
            new_opt_state[new_opt_name] = new_v

        opt_state.clear()
        opt_state.update(new_opt_state)
        return opt_state


class GroupGetter:
    def __init__(self, model, hcg=None):
        self.structure_name_mapping = {}
        self.structure_name_to_group = {}
        self.tensor_name_to_group = {}
        self.parameter_names = []
        self.group_map = OrderedDict()
        self.hcg = hcg or fleet.get_hybrid_communicate_group()
        for k, v in model.state_dict().items():
            self.structure_name_mapping[k] = v.name
            group = get_param_sharding_group(v, self.hcg)
            self.structure_name_to_group[k] = group
            self.tensor_name_to_group[v.name] = group
            self.group_map[group.id] = group
            self.parameter_names.append(v.name)

    def _get_parameter_name(self, name):
        if name in self.tensor_name_to_group:
            return name

        suffix = [
            "_fp32_master_0_beta1_pow_acc_0",
            "_fp32_master_0_beta2_pow_acc_0",
            "_fp32_master_0_moment1_0",
            "_fp32_master_0_moment2_0",
            "_beta1_pow_acc_0",
            "_beta2_pow_acc_0",
            "_moment1_0",
            "_moment2_0",
        ]

        for s in suffix:
            if name.endswith(s):
                tmp = name[: -len(s)]
                assert tmp in self.tensor_name_to_group, f"cannot find {name}"
                return tmp

        raise ValueError(f"cannot find {name}")

    def get_group(self, name):
        if name in self.structure_name_to_group:
            assert name not in self.tensor_name_to_group, name
            return self.structure_name_to_group[name]
        else:
            return self.tensor_name_to_group[self._get_parameter_name(name)]

    def get_group_by_id(self, gid):
        return self.group_map[gid]

    def get_group_ids(self):
        return list(self.group_map.keys())


class ShardingIO:
    def __init__(self, args, model, optimizer=None, hcg=None, remap_parameter_name=False, is_ema=False):
        self.args = args
        self.model = model
        self.optimizer = optimizer
        self.hcg = hcg
        self.sharding_group = None
        if self.hcg is None and paddle.distributed.get_world_size() > 1 and self.args.use_hybrid_parallel:
            self.hcg = fleet.get_hybrid_communicate_group()
            self.sharding_group = self.hcg.get_sharding_parallel_group()

        self.remap_parameter_name = remap_parameter_name
        self.remapper = None
        self.is_ema = is_ema

    def _get_remapper(self, checkpoint):
        if not self.remap_parameter_name:
            return None

        if self.remapper is None or self.remapper.checkpoint != checkpoint:
            new_mapping = {}
            for k, v in self.model.state_dict().items():
                new_mapping[k] = v.name

            suffix = self._sharding_meta_suffix()
            model_meta = self._load_model_meta_impl(checkpoint)
            old_mapping = model_meta["sharding_metas"][suffix]["structure_name_mapping"]
            self.remapper = ParameterNameRemapper(old_mapping, new_mapping, checkpoint)
        return self.remapper

    def _remap_parameter_name(self, checkpoint, state_dict, is_opt):
        remapper = self._get_remapper(checkpoint)
        if remapper is None:
            return state_dict
        if is_opt:
            return remapper.remap_optimizer_state(state_dict)
        else:
            return remapper.remap_model_state(state_dict)

    def set_optimizer(self, optimizer):
        self.optimizer = optimizer

    def load_state_dict_from_checkpoint_with_reshard(
        self, checkpoint, base_weight_name, model_wrapped, opt_state_dict=None
    ):
        """load state_dict from_checkpoint with reshard, Only load model state dict.
        Args:
            checkpoint (str): The directory of the checkpoint.
            base_weight_name (str): The name of the checkpoint file.
            model_wrapped (nn.Layer): The wrapped model.
        """
        group_getter = GroupGetter(self.model)
        gids = group_getter.get_group_ids()

        parallel_config = self._load_distributed_strategy(checkpoint)
        pp_degree = parallel_config["pp_degree"]
        mp_degree = parallel_config["mp_degree"]
        sharding_degree = parallel_config["sharding_degree"]
        assert (
            self.args.tensor_parallel_degree == mp_degree
        ), f"mp_degree of the script {self.args.tensor_parallel_degree} and mp of the model {mp_degree} are not matched"
        cur_sharding_degree = self.args.sharding_parallel_degree
        cur_pp_degree = self.args.pipeline_parallel_degree
        if pp_degree > 1:
            assert cur_pp_degree > 1, "can not reshard from pp to non pp"
        if pp_degree <= 1:
            assert cur_pp_degree <= 1, "can not reshard from non pp to pp"

        def load_model_slices():
            model_state = {gid: reshard_util.NodeModelState(group=group_getter.get_group_by_id(gid)) for gid in gids}
            for j in range(self.args.pipeline_parallel_rank, pp_degree, cur_pp_degree):
                cur_sharding_meta = self._load_sharding_meta(checkpoint, j)
                assert "structure_name_mapping" in cur_sharding_meta
                structure_name_map = cur_sharding_meta["structure_name_mapping"]
                structure_name_map = split_structure_name_mapping(structure_name_map, group_getter)
                for i in range(self.args.sharding_parallel_rank, sharding_degree, cur_sharding_degree):
                    tmp = self._load_one_state_dict_from_checkpoint(
                        checkpoint,
                        base_weight_name,
                        self.args.sharded_name_suffix(i, j, sharding_parallel_degree=sharding_degree),
                    )
                    tmp = split_model_state(tmp, group_getter)
                    for gid in gids:
                        sub_tmp = tmp.get(gid, {})
                        node_model_state_tmp = reshard_util.NodeModelState(group=group_getter.get_group_by_id(gid))
                        node_model_state_tmp.add_weights(sub_tmp)
                        node_model_state_tmp.pack_keys(structure_name_map.get(gid, {}))
                        model_state[gid].merge_from(node_model_state_tmp, i)
            return model_state

        node_model_state = load_model_slices()

        if self._need_reshard_pp(checkpoint):
            meta = self._load_model_meta(checkpoint)
            reshard_context = pp_reshard.build_pipeline_context(meta, model_wrapped)
            node_model_state = pp_reshard.reshard(node_model_state, reshard_context, self.hcg)

        if opt_state_dict is None:
            opt_state_dict = self.optimizer.state_dict()
        opt_state_dict = split_opt_state(opt_state_dict, group_getter)

        res_state_dict = OrderedDict()
        for gid, nms in node_model_state.items():
            nms.drop_rank()
            nms.unpack_keys()
            state_dict = nms.model_weights

            def filter_func(name):
                return True

            state_dict = reshard_util.all_gather_state_dict(state_dict, filter_func, nms.group)

            if self.args.bf16:
                state_dict = self._recover_params_from_master_weights(
                    state_dict,
                    opt_state_dict=opt_state_dict.get(gid, {}),
                    group=group_getter.get_group_by_id(gid),
                )

            res_state_dict.update(state_dict)

        return res_state_dict

    def _load_one_state_dict_from_checkpoint(self, resume_from_checkpoint, base_weight_name, weight_name_suffix):
        """
        load state_dict of one shard from_checkpoint, Only load model state dict.
        """
        if self.is_ema:
            base_weight_name = base_weight_name.replace("model_state", "ema").replace("pdparams", "pdopt")
        file_path = os.path.join(resume_from_checkpoint, _add_variant(base_weight_name, weight_name_suffix))
        if not os.path.isfile(file_path):
            raise ValueError(f"Can't find a valid checkpoint at {resume_from_checkpoint}, no {file_path}")

        logger.info(f"Loading model from {file_path}.")
        # We load the model state dict on the CPU to avoid an OOM error.
        state_dict = paddle.load(file_path, return_numpy=True)
        if self.is_ema:
            state_dict.pop("master_weights", None)
        state_dict = self._remap_parameter_name(resume_from_checkpoint, state_dict, is_opt=False)
        return state_dict

    def _load_optimizer_state_of_one_shard(self, checkpoint, base_opt_name, optimizer_name_suffix, group_getter=None):
        if self.is_ema:
            base_opt_name = base_opt_name.replace("optimizer", "ema")
        optimizer_name = _add_variant(base_opt_name, optimizer_name_suffix)
        path = os.path.join(checkpoint, optimizer_name)
        logger.info(f"load optimizer state from {path}")
        if os.path.isfile(path):
            opt_state = paddlenlp_load(path, map_location="cpu")
            if self.is_ema:
                opt_state = {"master_weights": opt_state.get("master_weights", {})}
            return self._remap_parameter_name(
                checkpoint,
                self._modify_ckpt_for_compatibility(opt_state),
                is_opt=True,
            )
        logger.info(f"{path} not exists")
        return None

    def _modify_ckpt_for_compatibility(self, ckpt):
        master_weights = ckpt.get("master_weights", None)
        if master_weights:
            for k, v in master_weights.items():
                assert isinstance(v, paddle.Tensor), v
                if not v.name.startswith(k):
                    new_name = k + "_fp32_master_0"
                    logger.info(f"Modify master weights {v.name} -> {new_name}")
                    v.name = new_name
        return ckpt

    def _need_reshard(self, checkpoint):
        if self._need_reshard_pp(checkpoint):
            return True
        parallel_config = self._load_distributed_strategy(checkpoint)
        sharding_meta = self._load_sharding_meta(checkpoint)
        sharding_degree = parallel_config["sharding_degree"]
        sharding_strategy = SHARDING_STRATEGY_V1
        if "sharding_strategy" in sharding_meta:
            sharding_strategy = sharding_meta["sharding_strategy"]
        cur_sharding_degree = self.args.sharding_parallel_degree
        cur_sharding_strategy = reshard_util.get_sharding_strategy(self.optimizer)
        if sharding_degree != cur_sharding_degree or sharding_strategy != cur_sharding_strategy:
            return True
        if sharding_strategy == SHARDING_STRATEGY_V1:
            param2rank = sharding_meta["param2rank"]
            optimizer = unwrap_optimizer(self.optimizer, DygraphShardingOptimizer)
            if self.args.sharding_parallel_degree > 1:
                assert optimizer is not None
            else:
                assert optimizer is None
            if len(param2rank) == 0 or optimizer is None:
                logger.warning("The param2rank is empty or sharding degree is 1. Force reshard would be performed.")
                return True
            assert len(param2rank) == len(optimizer._param2rank)
            for (k, v) in param2rank.items():
                assert k in optimizer._param2rank
                if optimizer._param2rank[k] != int(v):
                    return True
        else:
            pp_overlap = None
            # backward compatibility
            if "enable_overlap" in sharding_meta:
                pp_overlap = sharding_meta["enable_overlap"]

            cur_pp_overlap = unwrap_optimizer(self.optimizer, DygraphShardingOptimizerV2).pp_overlap
            return pp_overlap != cur_pp_overlap

        return False

    def _need_reshard_pp(self, checkpoint):
        parallel_config = self._load_distributed_strategy(checkpoint)
        pp_degree = parallel_config["pp_degree"]
        cur_pp_degree = self.args.pipeline_parallel_degree
        if pp_degree != cur_pp_degree:
            return True
        # vpp、segment method changes is not auto supported yet
        return self.args.force_reshard_pp

    def load_optimizer_state_with_reshard(self, checkpoint, base_opt_name, model_wrapped):
        """load state_dict of multiple shard from_checkpoint, Only load model state dict."""

        parallel_config = self._load_distributed_strategy(checkpoint)
        sharding_meta = self._load_sharding_meta(checkpoint)
        pp_degree = parallel_config["pp_degree"]
        mp_degree = parallel_config["mp_degree"]
        sharding_degree = parallel_config["sharding_degree"]
        assert sharding_degree > 1, "sharding degree of the checkpoint should be larger than 1"
        sharding_strategy = SHARDING_STRATEGY_V1
        if "sharding_strategy" in sharding_meta:
            sharding_strategy = sharding_meta["sharding_strategy"]
        assert self.args.tensor_parallel_degree == mp_degree
        cur_pp_degree = self.args.pipeline_parallel_degree

        if pp_degree > 1:
            assert cur_pp_degree > 1, "can not reshard from pp to non pp"
        if pp_degree <= 1:
            assert cur_pp_degree <= 1, "can not reshard from non pp to pp"

        cur_sharding_degree = self.args.sharding_parallel_degree
        cur_sharding_strategy = reshard_util.get_sharding_strategy(self.optimizer)

        group_getter = GroupGetter(self.model)

        if not self._need_reshard(checkpoint):
            one_shard_opt_state_dict = self._load_optimizer_state_of_one_shard(
                checkpoint,
                base_opt_name,
                self.args.sharded_name_suffix(sharding_parallel_degree=sharding_degree),
                group_getter=group_getter,
            )

            if sharding_strategy == SHARDING_STRATEGY_V2 and cur_sharding_strategy == SHARDING_STRATEGY_V2:
                is_matched = reshard_util.sharding_v2.is_matched_optimizer_state_dict(
                    one_shard_opt_state_dict, self.optimizer, model_wrapped
                )
                is_matched = paddle.to_tensor([is_matched], dtype=paddle.int32)
                dp_group = fleet.get_hybrid_communicate_group().get_data_parallel_group()
                dp_src_rank = fleet.get_hybrid_communicate_group().get_data_parallel_group_src_rank()
                dist.broadcast(is_matched, src=dp_src_rank, group=dp_group)
                is_matched = bool(is_matched[0])
            else:
                is_matched = True

            if is_matched:
                logger.info("do not need reshard")
                return one_shard_opt_state_dict
        else:
            one_shard_opt_state_dict = None

        logger.info("reshard optimizer state")
        gids = group_getter.get_group_ids()

        def load_model_slices():
            model_state = {gid: reshard_util.NodeModelState(group=group_getter.get_group_by_id(gid)) for gid in gids}
            for j in range(self.args.pipeline_parallel_rank, pp_degree, cur_pp_degree):
                cur_sharding_meta = self._load_sharding_meta(checkpoint, j)
                assert "structure_name_mapping" in cur_sharding_meta
                structure_name_map = cur_sharding_meta["structure_name_mapping"]
                structure_name_map = split_structure_name_mapping(structure_name_map, group_getter)
                for i in range(self.args.sharding_parallel_rank, sharding_degree, cur_sharding_degree):
                    sharded_name_suffix = self.args.sharded_name_suffix(i, j, sharding_parallel_degree=sharding_degree)
                    if one_shard_opt_state_dict is None:
                        tmp = self._load_optimizer_state_of_one_shard(checkpoint, base_opt_name, sharded_name_suffix)
                    else:
                        assert (
                            self.args.optimizer_name_suffix == sharded_name_suffix
                        ), f"{self.args.optimizer_name_suffix} vs {sharded_name_suffix}"
                        tmp = one_shard_opt_state_dict

                    tmp = split_opt_state(tmp, group_getter)
                    for gid in gids:
                        sub_tmp = tmp.get(gid, {})
                        node_model_state_tmp = reshard_util.NodeModelState(group=group_getter.get_group_by_id(gid))
                        node_model_state_tmp.add_opts(sub_tmp)
                        node_model_state_tmp.pack_keys(structure_name_map.get(gid, {}))
                        model_state[gid].merge_from(node_model_state_tmp, i)
            return model_state

        def reshard_pp(model_state):
            # pp reshard
            if self._need_reshard_pp(checkpoint):
                assert len(model_state) == 1, "only support one group reshard"
                key = list(model_state.keys())[0]
                tmp = model_state[key]
                meta = self._load_model_meta(checkpoint)
                reshard_context = pp_reshard.build_pipeline_context(meta, model_wrapped)
                model_state = {key: pp_reshard.reshard(tmp, reshard_context, self.hcg)}
            return model_state

        def reshard_sharding(node_model_state):
            # shard reshard
            restore_func = (
                reshard_util.sharding_v1.restore
                if sharding_strategy == SHARDING_STRATEGY_V1
                else reshard_util.sharding_v2.restore
            )

            for gid in gids:
                node_model_state[gid] = restore_func(node_model_state[gid], self.model, self.optimizer)

            shard_func = (
                reshard_util.sharding_v1.shard
                if cur_sharding_strategy == SHARDING_STRATEGY_V1
                else reshard_util.sharding_v2.shard
            )

            ret_opt_state_dict = OrderedDict()
            for gid in gids:
                node_model_state[gid] = shard_func(node_model_state[gid], model_wrapped, self.optimizer)
                # drop structural name in the key
                node_model_state[gid].unpack_keys()
                ret_opt_state_dict[gid] = node_model_state[gid].get_opt_state_dict()
            return merge_opt_state(ret_opt_state_dict)

        node_model_state = load_model_slices()
        node_model_state = reshard_pp(node_model_state)
        return reshard_sharding(node_model_state)

    def manipulate_state_dict_and_config(self, model_to_save, merge_tensor_parallel=False, state_dict=None):
        weight_name_suffix = self.args.sharded_name_suffix()
        group_getter = GroupGetter(model_to_save)
        gids = group_getter.get_group_ids()

        if state_dict is None:
            state_dict = model_to_save.state_dict()
            if self.args.should_save_sharding_stage1_model:
                state_dict = split_model_state(state_dict, group_getter)
                for gid in gids:
                    state_dict[gid] = filter_sharded_params(
                        state_dict.get(gid, {}),
                        self.optimizer,
                        self.sharding_group,
                        self.args.save_sharding_stage1_model_include_freeze_params,
                    )
                state_dict = merge_model_state(state_dict)

        config_to_save = None
        merge_tensor_parallel = merge_tensor_parallel and self.args.use_hybrid_parallel
        if merge_tensor_parallel:
            dtype = get_parameter_dtype(model_to_save)
            assert hasattr(model_to_save, "config")
            model_to_save.config.dtype = str(dtype).split(".")[1]
            config_to_save = copy.deepcopy(model_to_save.config)
            if config_to_save.tensor_parallel_degree > 1:
                state_dict = model_to_save.merge_tensor_parallel(state_dict, config_to_save)
                config_to_save.tensor_parallel_degree = 1
                if config_to_save.tensor_parallel_rank != 0:
                    logger.info("Saving with merge_tensor_parallel, tensor_parallel_rank > 0 don't need save")
                    return
                # if variant is not None and "tp" in variant:
                if "tp" in weight_name_suffix:
                    weight_name_suffix = "_".join([x for x in weight_name_suffix.split("_") if "tp" not in x])

        if self.args.bf16 and self.args.should_save_sharding_stage1_model:
            param_names_in_master_weights = []
            optimzier_state_dict = self.optimizer.state_dict()
            optimzier_state_dict = split_opt_state(optimzier_state_dict, group_getter)
            state_dict = split_model_state(state_dict, group_getter)
            for gid in gids:
                sub_opt_state = optimzier_state_dict.get(gid, {})
                param_names_in_master_weights = list(sub_opt_state.get("master_weights", {}).keys())
                state_dict[gid] = exclude_parameters_in_state_dict(
                    state_dict.get(gid, {}),
                    param_names_in_master_weights,
                    group_getter.get_group_by_id(gid),
                )
            state_dict = merge_model_state(state_dict)
            logger.info(
                "param_names_in_master_weights len:{}, bf16 state_dict len:{}, :{}".format(
                    len(param_names_in_master_weights), len(state_dict), state_dict.keys()
                )
            )
        return state_dict, config_to_save, weight_name_suffix

    def gather_distributed_model_meta(self):
        if not self.args.use_hybrid_parallel:
            return None

        if not self.args.should_save_sharding_stage1_model:
            return None

        nranks = dist.get_world_size()
        if nranks <= 1:
            return None

        model_meta = {}
        model_meta["parallel_config"] = self._get_distributed_strategy()
        model_meta["sharding_metas"] = self._gather_sharding_metas()

        return model_meta

    def _check_distributed_strategy(self, parallel_config):
        ep_degree = parallel_config.get("ep_degree", 1)
        if ep_degree > 1:
            tp_degree = parallel_config["mp_degree"]
            sharding_degree = parallel_config["sharding_degree"]
            moe_sharding_degree = parallel_config.get("moe_sharding_degree", 1)
            assert tp_degree * sharding_degree == ep_degree * moe_sharding_degree, "mismatch parallel degree settings"

    def check_same_strategy(self, resume_from_checkpoint=None):
        if resume_from_checkpoint:
            cur_config = self._get_distributed_strategy()
            old_config = self._load_model_meta_impl(resume_from_checkpoint)["parallel_config"]
            keys = list(old_config.keys())
            for key in keys:
                if key not in cur_config:
                    return False, f"missing {key}"
                else:
                    old_value = old_config[key]
                    cur_value = cur_config[key]
                    if old_value != cur_value:
                        return False, f"{key} not match: {old_value} vs {cur_value}"
        return True, None

    def _get_distributed_strategy(self):
        pp_degree = 1
        mp_degree = 1
        sharding_degree = 1
        ep_degree = 1
        moe_sharding_degree = 1
        nranks = dist.get_world_size()
        if self.args.use_hybrid_parallel and nranks > 1:
            hcg = fleet.get_hybrid_communicate_group()
            mp_degree = hcg.get_model_parallel_world_size()
            pp_degree = hcg.get_pipe_parallel_world_size()
            sharding_degree = hcg.get_sharding_parallel_world_size()
            if hasattr(hcg, "get_expert_parallel_world_size"):
                ep_degree = hcg.get_expert_parallel_world_size()
            if hasattr(hcg, "get_moe_sharding_parallel_world_size"):
                moe_sharding_degree = hcg.get_moe_sharding_parallel_world_size()

        parallel_config = {
            "pp_degree": pp_degree,
            "mp_degree": mp_degree,
            "sharding_degree": sharding_degree,
            "ep_degree": ep_degree,
            "moe_sharding_degree": moe_sharding_degree,
        }
        self._check_distributed_strategy(parallel_config)
        return parallel_config

    def _recover_params_from_master_weights(self, state_dict, opt_state_dict=None, group=None):
        if group is None:
            group = self.sharding_group
        if opt_state_dict is None:
            opt_state_dict = self.optimizer.state_dict()
        assert "master_weights" in opt_state_dict, opt_state_dict.keys()
        master_weights = opt_state_dict["master_weights"]
        tmp = OrderedDict()
        (master_weights, tmp) = (tmp, master_weights)
        # cast to before
        for (k, v) in tmp.items():
            name = v.name
            master_weights[k] = paddle.cast(to_device(v), paddle.bfloat16).cpu()
            master_weights[k].name = name

        structure_name_map = {k: v.name for (k, v) in self.model.state_dict().items()}
        node_model_state = reshard_util.NodeModelState(group=group)
        node_model_state_tmp = reshard_util.NodeModelState(group=group)
        node_model_state_tmp.add_master_weights(master_weights)
        node_model_state_tmp.pack_keys(structure_name_map)
        node_model_state.merge_from(node_model_state_tmp, max(group.rank, 0))
        del node_model_state_tmp
        sharding_strategy = reshard_util.get_sharding_strategy(self.optimizer)
        restore_func = (
            reshard_util.sharding_v1.restore
            if sharding_strategy == SHARDING_STRATEGY_V1
            else reshard_util.sharding_v2.restore
        )
        node_model_state = restore_func(node_model_state, self.model, self.optimizer)
        node_model_state.unpack_keys()
        master_weights = node_model_state.master_weights

        def filter_func(name):
            return True

        master_weights = reshard_util.all_gather_state_dict(master_weights, filter_func, group)
        model_state_dict = self.model.state_dict()
        logger.info(f"state-dict-keys: {state_dict.keys()}, nums: {len(state_dict.keys())}")
        logger.info("before recover, model_state_dict number: {}".format(len(model_state_dict)))
        for key, param in model_state_dict.items():
            if param.name in master_weights:
                assert param.shape == master_weights[param.name].shape
                paddle.assign(
                    paddle.cast(to_device(master_weights[param.name]), paddle.bfloat16), model_state_dict[key]
                )
            elif key in state_dict:
                logger.info(f"key: {key} is in state_dict, but not in master_weights")
                paddle.assign(state_dict[key], model_state_dict[key])
            else:
                logger.info(f"key: {key} is not in state_dict and master_weights")
        logger.info("after recover, casted model_state_dict number: {}".format(len(model_state_dict)))
        state_dict.update(model_state_dict)
        return state_dict

    def _all_gather_simple_object(self, obj, group=None):
        if group is None:
            group = self.hcg.get_sharding_parallel_group()
        res = []
        if group.nranks < 2:
            return [obj]
        paddle.distributed.all_gather_object(res, obj, group)
        return res

    def _load_model_meta_impl(self, dir):
        meta_path = os.path.join(dir, MODEL_META_NAME)
        assert os.path.exists(meta_path), f"{meta_path} not exist"
        with open(meta_path, "r") as handle:
            model_dist_meta = json.load(handle)

        assert "parallel_config" in model_dist_meta
        self._check_distributed_strategy(model_dist_meta["parallel_config"])
        return model_dist_meta

    def _load_model_meta(self, dir):
        model_meta = self._load_model_meta_impl(dir)
        remapper = self._get_remapper(dir)
        if remapper is not None:
            suffix = self._sharding_meta_suffix()
            sharding_metas = model_meta["sharding_metas"]
            cur_sharding_metas = sharding_metas.pop(suffix)
            sharding_metas.clear()
            sharding_metas[suffix] = cur_sharding_metas
            cur_sharding_metas["structure_name_mapping"] = remapper.new_mapping
            if "param2rank" in cur_sharding_metas:
                new_param2rank = {}
                for k, rank in cur_sharding_metas["param2rank"].items():
                    new_k = remapper.p_name_map[k]
                    new_param2rank[new_k] = rank
                cur_sharding_metas["param2rank"] = new_param2rank
        return model_meta

    def _sharding_meta_suffix(self, tp_rank=None, pp_rank=None):
        if tp_rank is None:
            tp_rank = self.args.tensor_parallel_rank
        if pp_rank is None:
            pp_rank = self.args.pipeline_parallel_rank
        suffix = f"tp{tp_rank:0>2d}_pp{pp_rank:0>2d}"
        if self.args.expert_parallel_degree > 1:
            ep_rank = self.args.expert_parallel_rank
            return f"{suffix}_ep{ep_rank:0>2d}"
        else:
            return suffix

    def _load_distributed_strategy(self, dir):
        model_dist_meta = self._load_model_meta(dir)
        parallel_config = model_dist_meta["parallel_config"]
        assert "pp_degree" in parallel_config
        assert "mp_degree" in parallel_config
        assert "sharding_degree" in parallel_config
        return parallel_config

    def _load_sharding_meta(self, dir, pp_rank=None):
        suffix = self._sharding_meta_suffix(pp_rank=pp_rank)
        distributed_model_meta = self._load_model_meta(dir)
        if "sharding_metas" in distributed_model_meta:
            sharding_metas = distributed_model_meta["sharding_metas"]
            assert suffix in sharding_metas
            sharding_meta = sharding_metas[suffix]
            assert "param2rank" in sharding_meta
            return sharding_meta

        # for backward compatibility
        meta_path = os.path.join(dir, _add_variant(SHARDING_META_NAME, suffix))
        assert os.path.exists(meta_path), f"{meta_path} not exist"
        with open(meta_path, "r") as f:
            sharding_meta = json.load(f)
        assert "param2rank" in sharding_meta
        return sharding_meta

    def _map_optimizer_state_to_param(self, optimizer_state_names):
        optimizer = unwrap_optimizer(self.optimizer, DygraphShardingOptimizer)
        all_names = list(optimizer._param2rank.keys())
        all_names.extend(list(optimizer_state_names))
        all_names.sort()
        pre_p_name = ""
        opt_to_p = {}
        for n in all_names:
            if n in optimizer._param2rank:
                # we get a param
                pre_p_name = n
            else:
                assert pre_p_name, n
                opt_to_p[n] = pre_p_name
        return opt_to_p

    def _gather_sharding_metas(self):
        nranks = dist.get_world_size()
        if not self.args.use_hybrid_parallel or nranks <= 1:
            return None
        if not reshard_util.is_sharding_opt(self.optimizer):
            return None

        sharding_strategy = reshard_util.get_sharding_strategy(self.optimizer)
        param2rank = {}
        pp_overlap = False
        if sharding_strategy == SHARDING_STRATEGY_V1:
            optimizer = unwrap_optimizer(self.optimizer, DygraphShardingOptimizer)
            param2rank = {k: v for (k, v) in optimizer._param2rank.items()}
        else:
            pp_overlap = unwrap_optimizer(self.optimizer, DygraphShardingOptimizerV2).pp_overlap

        model = self.model
        structure_name_mapping = {}
        param_meta = {}
        for k, v in model.state_dict().items():
            structure_name_mapping[k] = v.name
            is_distributed = getattr(v, "is_distributed", False)
            no_sync = getattr(v, "no_sync", False)
            param_meta[k] = (v.shape, int(v.dtype), is_distributed, no_sync)

        sharding_metas = {}
        sharding_meta = {}

        sharding_meta["param2rank"] = param2rank
        sharding_meta["structure_name_mapping"] = structure_name_mapping
        sharding_meta["param_meta"] = param_meta
        sharding_meta["param_meta_keys"] = ["shape", "dtype", "is_distributed", "no_sync"]
        sharding_meta["sharding_strategy"] = sharding_strategy
        sharding_meta["enable_overlap"] = pp_overlap
        dp_metas_list = self._all_gather_simple_object(sharding_meta, self.hcg.get_data_parallel_group())
        for e in dp_metas_list:
            for key in ["structure_name_mapping", "param_meta"]:
                sharding_meta[key].update(e[key])
        suffix = self._sharding_meta_suffix()
        sharding_metas[suffix] = sharding_meta
        sharding_metas_list = self._all_gather_simple_object(sharding_metas, self.hcg.get_model_parallel_group())
        sharding_metas = {k: v for e in sharding_metas_list for (k, v) in e.items()}
        sharding_metas_list = self._all_gather_simple_object(sharding_metas, self.hcg.get_pipe_parallel_group())
        sharding_metas = {k: v for e in sharding_metas_list for (k, v) in e.items()}
        if self.args.expert_parallel_degree > 1:
            sharding_metas_list = self._all_gather_simple_object(sharding_metas, self.hcg.get_expert_parallel_group())
            sharding_metas = {k: v for e in sharding_metas_list for (k, v) in e.items()}
        return sharding_metas
