[Unified Checkpoint] Fix split param loading directly when using ignore_merge_optimizer (#9935)

* fix * refine code * update typename

Author: DesmonDay

paddlenlp/trainer/unified_checkpoint/sharding_split_param_utils.py

@@ -15,18 +15,27 @@
 
 import gc
 import os
+import re
 from itertools import chain
 
 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.transformers.utils import device_guard
 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 +184,49 @@ 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 reshape_params(state_dict, struct2static_name_mappings, param_shape_info, param_slice_info):
+ """Reshape params to 1-D tensors"""
+ for key in list(state_dict.keys()):
+ key_name = key.split("/")[0]
+ static_name = struct2static_name_mappings.get(key_name, None)
+ if int(state_dict[key].numel()) > 1:
+ begin, end = param_slice_info[static_name]
+ _, numel, index, padded_size = param_shape_info[static_name]
+ state_dict[key] = state_dict[key].reshape([-1])
+ state_dict[key] = 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:
+ state_dict[key] = paddle.concat(
+ (
+ state_dict[key],
+ paddle.zeros([padding_end - padding_start], dtype=state_dict[key].dtype),
+ )
+ )
+ return state_dict
+
+
 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 +248,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:
@@ -285,25 +321,10 @@ def load_resolved_archive_file(
  )
 
  # need to split param for different sharding rank, maybe need to deal with oom issue.
+ reshape_params(state_dict_optim, struct2static_name_mappings, param_shape_info, param_slice_info)
  for key in list(state_dict_optim.keys()):
  key_name = key.split("/")
  static_name = struct2static_name_mappings.get(key_name[0], None)
-
- 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]
- state_dict_optim[key] = state_dict_optim[key].reshape([-1])
- state_dict_optim[key] = state_dict_optim[key][begin - index : end - index]
-
- padding_start = max(begin, index + numel)
- padding_end = min(end, index + padded_size)
- if padding_start < padding_end:
- state_dict_optim[key] = paddle.concat(
- (
- state_dict_optim[key],
- paddle.zeros([padding_end - padding_start], dtype=state_dict_optim[key].dtype),
- )
- )
  if has_master_weights:
  if model_state_dict[key_name[0]].dtype != paddle.float32:
  key_name = "_".join([static_name, FP32_MASTER, key_name[1]])
@@ -325,24 +346,10 @@ def load_resolved_archive_file(
  expected_keys,
  is_master_weights=True,
  )
+ reshape_params(state_dict_master_weight, struct2static_name_mappings, param_shape_info, param_slice_info)
 
  for key in list(state_dict_master_weight.keys()):
  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]
- 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]
-
- padding_start = max(begin, index + numel)
- padding_end = min(end, index + padded_size)
- if padding_start < padding_end:
- state_dict_master_weight[key] = paddle.concat(
- (
- state_dict_master_weight[key],
- paddle.zeros([padding_end - padding_start], dtype=state_dict_master_weight[key].dtype),
- )
- )
  state_dict_master_weight[key] = state_dict_master_weight[key]._copy_to(
  paddle.framework._current_expected_place(), False
  )
@@ -357,3 +364,113 @@ 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 = []
+ sharding_typename_set, typename_set = [], []
+ with safe_open(optimizer_path, framework="numpy") as f:
+ optim_keys = f.keys()
+ for key in optim_keys:
+ _, typename = key.split("/")
+ typename_set.append(typename)
+
+ # To avoid incomplete typename in some shard files, communication is performed.
+ hcg = fleet.get_hybrid_communicate_group()
+ sharding_group = hcg.get_sharding_parallel_group()
+ dist.all_gather_object(sharding_typename_set, typename_set, sharding_group)
+ typename_set = set(chain(*sharding_typename_set))
+ 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")
+
+ def get_unfound_params(unfound_keys, state_dict, is_optimizer=True):
+ if len(unfound_keys) > 0:
+ backup_files = []
+ files = os.listdir(resume_from_checkpoint)
+ name = optimizer_name if is_optimizer else master_weights_name
+ name_without_shard = re.sub(r"_?shard\d+_?", "", name)
+ name_ = "optimizer" if is_optimizer else "master_weights"
+ for f in files:
+ if f.startswith(name_) and f.endswith("safetensors") and f != name:
+ if re.sub(r"_?shard\d+_?", "", f) == name_without_shard:
+ backup_files.append(f)
+ for f in backup_files:
+ new_path = os.path.join(resume_from_checkpoint, f)
+ with safe_open(new_path, framework="numpy") as fin:
+ keys = fin.keys()
+ for key in unfound_keys:
+ if key in keys:
+ tensor = fin.get_tensor(key)
+ with device_guard():
+ tensor = paddle.Tensor(tensor, zero_copy=True)
+ state_dict[key] = tensor._copy_to(paddle.framework._current_expected_place(), False)
+
+ # Get other optimizer paramsters which maybe in other shard files.
+ unfound_keys = expected_keys_optim - optimizer_state_dict.keys()
+ get_unfound_params(unfound_keys, optimizer_state_dict, True)
+
+ # Get other master weight parameters which maybe in other shard files.
+ if master_weights != {}:
+ unfound_keys = expected_keys - master_weights.keys()
+ get_unfound_params(unfound_keys, master_weights, False)
+ reshape_params(optimizer_state_dict, struct2static_name_mappings, param_shape_info, param_slice_info)
+
+ # 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:
+ reshape_params(master_weights, struct2static_name_mappings, param_shape_info, param_slice_info)
+
+ 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())