fix

Author: DesmonDay

paddlenlp/trainer/unified_checkpoint/sharding_split_param_utils.py

@@ -20,13 +20,20 @@
 import paddle
 import paddle.distributed as dist
 from paddle.distributed import fleet
+from safetensors import safe_open
 from tqdm.auto import tqdm
 
 from paddlenlp.peft import LoRAModel, PrefixModelForCausalLM
-from paddlenlp.transformers.model_utils import load_state_dict, unwrap_model
+from paddlenlp.transformers.model_utils import (
+ _add_variant,
+ load_state_dict,
+ unwrap_model,
+)
 from paddlenlp.utils.env import (
  SAFE_MASTER_WEIGHTS_INDEX_NAME,
+ SAFE_MASTER_WEIGHTS_NAME,
  SAFE_OPTIMIZER_INDEX_NAME,
+ SAFE_OPTIMIZER_NAME,
 )
 from paddlenlp.utils.nested import nested_copy
 
@@ -175,6 +182,26 @@ def gather_splited_param_for_optimizer(optimizer, ckpt_quant_stage="O0"):
  return optim_state_dict, master_weights
 
 
+def get_params_info(comm_buffer_list):
+ expected_keys = []
+ param_slice_info = {}
+ param_shape_info = {}
+
+ for buffer in comm_buffer_list:
+ for key in buffer._sharding_param_grad_view.keys():
+ begin = buffer._sharding_param_grad_view[key]._param_begin
+ end = buffer._sharding_param_grad_view[key]._param_end
+ if end > begin:
+ expected_keys.append(key)
+ shape = buffer._sharding_param_grad_view[key]._param.shape
+ numel = buffer._sharding_param_grad_view[key]._param.numel().item()
+ index = buffer._sharding_param_grad_view[key]._index
+ padded_size = buffer._sharding_param_grad_view[key]._padded_size
+ param_slice_info[key] = (begin, end)
+ param_shape_info[key] = (shape, numel, index, padded_size)
+ return expected_keys, param_slice_info, param_shape_info
+
+
 def load_unified_optimizer_split_param(args, model, optimizer, resume_from_checkpoint, ckpt_quant_stage="O0"):
  returned_optim_state_dict = nested_copy(optimizer.state_dict())
 
@@ -196,28 +223,12 @@ def load_unified_optimizer_split_param(args, model, optimizer, resume_from_check
  static2struct_name_mappings = {v.name: k for k, v in model_state_dict.items()} # get optimizer param mappings
  struct2static_name_mappings = {k: v.name for k, v in model_state_dict.items()}
 
- expected_keys = []
- param_slice_info = {}
- param_shape_info = {}
-
  comm_buffer_list = optimizer._inner_opt._comm_buffer_list
  if hasattr(args, "enable_sharding_comm_overlap") and args.enable_sharding_comm_overlap:
  comm_buffer_list = list(chain(*model._chunk_2_comm_buffers.values()))
  model = unwrap_model(model)
 
- for buffer in comm_buffer_list:
- for key in buffer._sharding_param_grad_view.keys():
- begin = buffer._sharding_param_grad_view[key]._param_begin
- end = buffer._sharding_param_grad_view[key]._param_end
- if end > begin:
- expected_keys.append(key)
- shape = buffer._sharding_param_grad_view[key]._param.shape
- numel = buffer._sharding_param_grad_view[key]._param.numel().item()
- index = buffer._sharding_param_grad_view[key]._index
- padded_size = buffer._sharding_param_grad_view[key]._padded_size
- param_slice_info[key] = (begin, end)
- param_shape_info[key] = (shape, numel, index, padded_size)
-
+ expected_keys, param_slice_info, param_shape_info = get_params_info(comm_buffer_list)
  expected_keys = set([static2struct_name_mappings.get(name, None) for name in expected_keys])
  expected_keys_optim = []
  for key in expected_keys:
@@ -291,7 +302,7 @@ def load_resolved_archive_file(
 
  if int(state_dict_optim[key].numel()) > 1:
  begin, end = param_slice_info[static_name]
- shape, numel, index, padded_size = param_shape_info[static_name]
+ _, numel, index, padded_size = param_shape_info[static_name]
  state_dict_optim[key] = state_dict_optim[key].reshape([-1])
  state_dict_optim[key] = state_dict_optim[key][begin - index : end - index]
 
@@ -330,7 +341,7 @@ def load_resolved_archive_file(
  static_name = struct2static_name_mappings.get(key, None)
  if int(state_dict_master_weight[key].numel()) > 1:
  begin, end = param_slice_info[static_name]
- shape, numel, index, padded_size = param_shape_info[static_name]
+ _, numel, index, padded_size = param_shape_info[static_name]
  state_dict_master_weight[key] = state_dict_master_weight[key].reshape([-1])
  state_dict_master_weight[key] = state_dict_master_weight[key][begin - index : end - index]
 
@@ -357,3 +368,122 @@ def load_resolved_archive_file(
  returned_optim_state_dict["master_weights"][static_name].name = "_".join([static_name, FP32_MASTER])
 
  return returned_optim_state_dict
+
+
+def load_non_merge_optimizer_with_split_param(args, model, optimizer, resume_from_checkpoint, ckpt_quant_stage="O0"):
+ returned_optim_state_dict = nested_copy(optimizer.state_dict())
+
+ optimizer_name = _add_variant(SAFE_OPTIMIZER_NAME, args.optimizer_name_suffix)
+ master_weights_name = _add_variant(SAFE_MASTER_WEIGHTS_NAME, args.optimizer_name_suffix)
+ optimizer_path = os.path.join(resume_from_checkpoint, optimizer_name)
+ master_weights_path = os.path.join(resume_from_checkpoint, master_weights_name)
+
+ # no quantization & no master weight represent O1 AMP strategy.
+ is_amp_o1 = args.fp16_opt_level == "O1"
+
+ model_state_dict = get_expected_state_dict(model)
+ static2struct_name_mappings = {v.name: k for k, v in model_state_dict.items()} # get optimizer param mappings
+ struct2static_name_mappings = {k: v.name for k, v in model_state_dict.items()}
+
+ comm_buffer_list = optimizer._inner_opt._comm_buffer_list
+ if hasattr(args, "enable_sharding_comm_overlap") and args.enable_sharding_comm_overlap:
+ comm_buffer_list = list(chain(*model._chunk_2_comm_buffers.values()))
+
+ expected_keys, param_slice_info, param_shape_info = get_params_info(comm_buffer_list)
+ expected_keys = set([static2struct_name_mappings.get(name, None) for name in expected_keys])
+ expected_keys_optim = []
+ typename_set = set()
+ with safe_open(optimizer_path, framework="numpy") as f:
+ optim_keys = f.keys()
+ for key in optim_keys:
+ _, typename = key.split("/")
+ typename_set.add(typename)
+ for key in expected_keys:
+ for typename in typename_set:
+ expected_keys_optim.append(f"{key}/{typename}")
+ expected_keys_optim = set(expected_keys_optim)
+
+ optimizer_state_dict = load_state_dict(
+ optimizer_path, None, None, device="expected", ckpt_quant_stage=ckpt_quant_stage
+ )
+ master_weights = {}
+ # normal AMP O2
+ if not is_amp_o1 and os.path.isfile(master_weights_path):
+ master_weights = load_state_dict(master_weights_path, None, None, device="expected")
+
+ # Get other param slice which maybe in other shard files.
+ unfound_keys = expected_keys_optim - optimizer_state_dict.keys()
+ if len(unfound_keys) > 0:
+ backup_files = []
+ files = os.listdir(resume_from_checkpoint)
+ for f in files:
+ if f.startswith("optimizer") and f.endswith("safetensors"):
+ backup_files.append(f)
+ print(backup_files)
+ raise ValueError
+
+ for key in list(optimizer_state_dict.keys()):
+ key_name = key.split("/")
+ static_name = struct2static_name_mappings.get(key_name[0], None)
+
+ if int(optimizer_state_dict[key].numel()) > 1:
+ begin, end = param_slice_info[static_name]
+ _, numel, index, padded_size = param_shape_info[static_name]
+ optimizer_state_dict[key] = optimizer_state_dict[key].reshape([-1])
+ optimizer_state_dict[key] = optimizer_state_dict[key][begin - index : end - index]
+
+ padding_start = max(begin, index + numel)
+ padding_end = min(end, index + padded_size)
+ if padding_start < padding_end:
+ optimizer_state_dict[key] = paddle.concat(
+ (
+ optimizer_state_dict[key],
+ paddle.zeros([padding_end - padding_start], dtype=optimizer_state_dict[key].dtype),
+ )
+ )
+
+ # rename and move to paddle.Tensor
+ for key in list(optimizer_state_dict.keys()):
+ key_name = key.split("/")
+ model_weight_key = key_name[0]
+ static_name = struct2static_name_mappings[key_name[0]]
+ if not is_amp_o1:
+ if model_state_dict[key_name[0]].dtype != paddle.float32:
+ key_name = "_".join([static_name, FP32_MASTER, key_name[1]])
+ else:
+ key_name = "_".join([static_name, key_name[1]])
+ else:
+ key_name = "_".join([static_name, key_name[1]])
+ returned_optim_state_dict[key_name] = optimizer_state_dict.pop(key)
+ returned_optim_state_dict[key_name].name = key_name
+
+ # master weight cast (only in AMP O2 + remove_master_weight)
+ if not is_amp_o1 and not os.path.isfile(master_weights_path):
+ master_weights[model_weight_key] = paddle.cast(model_state_dict[model_weight_key], dtype=paddle.float32)
+
+ if not is_amp_o1:
+ for key in list(master_weights.keys()):
+ static_name = struct2static_name_mappings.get(key, None)
+ if int(master_weights[key].numel()) > 1:
+ begin, end = param_slice_info[static_name]
+ _, numel, index, padded_size = param_shape_info[static_name]
+ master_weights[key] = master_weights[key].reshape([-1])
+ master_weights[key] = master_weights[key][begin - index : end - index]
+
+ padding_start = max(begin, index + numel)
+ padding_end = min(end, index + padded_size)
+ if padding_start < padding_end:
+ master_weights[key] = paddle.concat(
+ (
+ master_weights[key],
+ paddle.zeros([padding_end - padding_start], dtype=master_weights[key].dtype),
+ )
+ )
+
+ returned_optim_state_dict["master_weights"] = {}
+ for key in list(master_weights.keys()):
+ static_name = struct2static_name_mappings[key]
+ returned_optim_state_dict["master_weights"][static_name] = master_weights.pop(key)
+ returned_optim_state_dict["master_weights"][static_name].name = "_".join([static_name, FP32_MASTER])
+
+ return returned_optim_state_dict

paddlenlp/trainer/unified_checkpoint/unified_checkpoint.py

@@ -62,7 +62,10 @@
  save_single_card_checkpoint,
  save_single_card_optimizer,
 )
-from .sharding_split_param_utils import gather_splited_param_for_optimizer
+from .sharding_split_param_utils import (
+ gather_splited_param_for_optimizer,
+ load_non_merge_optimizer_with_split_param,
+)
 from .utils import (
  FP32_MASTER,
  UnifiedCheckpointOption,
@@ -263,6 +266,23 @@ def save_non_merge_optimizer(self, model, optim_state_dict, master_weights, outp
  )
 
  def load_non_merge_optimizer(self, model, optimizer, resume_from_checkpoint, ckpt_quant_stage="O0"):
+ """load non merge optimizer
+
+ Args:
+ model (PretrainedModel): model used to get key mapping.
+ optimizer (Optimizer): optimizer to load
+ resume_from_checkpoint (str): path of the checkpoint to load
+ ckpt_quant_stage (str): ckpt quant stage
+
+ Returns:
+ dict: optimizer state dict
+ """
+
+ if is_sharding_split_param_mode(self.args):
+ return load_non_merge_optimizer_with_split_param(
+ self.args, model, optimizer, resume_from_checkpoint, ckpt_quant_stage
+ )
+
  # init and get optimizer LR_Scheduler
  returned_optim_state_dict = nested_copy(optimizer.state_dict())