# 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