[Unified Checkpoint] Fix tie_weights save and load

paddlenlp/trainer/plugins/unified_checkpoint.py

@@ -226,6 +226,10 @@
 
     model_state_dict = get_expected_state_dict(model)
     expected_keys = set(list(model_state_dict.keys()))
+    if hasattr(model, "_tied_weights_keys") and model._tied_weights_keys is not None:
+        for key in model._tied_weights_keys:
+            expected_keys.remove(key)
+
     missing_keys = expected_keys - set(loaded_keys)
 
     if len(missing_keys) > 0:
@@ -607,6 +611,12 @@
     static2struct_name_mappings = {}
     state_dict = get_expected_state_dict(model)
     for k, v in state_dict.items():
+        if (
+            hasattr(model, "_tied_weights_keys")
+            and model._tied_weights_keys is not None
+            and k in model._tied_weights_keys
+        ):
+            continue
         static2struct_name_mappings[v.name] = k
 
     # rename optimizer param
@@ -739,6 +749,12 @@
         need_files = set()
         state_dict = get_expected_state_dict(model)
         for key in state_dict.keys():
+            if (
+                hasattr(model, "_tied_weights_keys")
+                and model._tied_weights_keys is not None
+                and key in model._tied_weights_keys
+            ):
+                continue
             filename = index["weight_map"][key]
             need_files.add(filename)
         diff_filelist = list(need_files.difference(set(existed_files)))
@@ -829,6 +845,12 @@
             need_files = set()
             state_dict = get_expected_state_dict(model)
             for key in state_dict.keys():
+                if (
+                    hasattr(model, "_tied_weights_keys")
+                    and model._tied_weights_keys is not None
+                    and key in model._tied_weights_keys
+                ):
+                    continue
                 if sharding_group.nranks > 1:
                     static_name = struct2static_name_mappings.get(key, None)
                     param_rank = param2rank.get(static_name, None)
@@ -893,6 +915,12 @@
     index_weight_file = {}
     total_size = 0
     for key, weight in state_dict.items():
+        if (
+            hasattr(model_to_save, "_tied_weights_keys")
+            and model_to_save._tied_weights_keys is not None
+            and key in model_to_save._tied_weights_keys
+        ):
+            continue
         index_weight_file[key] = weight_filename
         total_size += weight.numel().item() * dtype_byte_size(weight.dtype)
     sharded_index_json = {}
@@ -926,6 +954,12 @@
     static2struct_name_mappings = {}
     state_dict = get_expected_state_dict(model)
     for k, v in state_dict.items():
+        if (
+            hasattr(model, "_tied_weights_keys")
+            and model._tied_weights_keys is not None
+            and k in model._tied_weights_keys
+        ):
+            continue
         static2struct_name_mappings[v.name] = k
 
     # rename optimizer param
@@ -1023,6 +1057,12 @@
     if args.dataset_rank == 0:
         state_dict = get_expected_state_dict(model)
         for (k, v) in state_dict.items():
+            if (
+                hasattr(model, "_tied_weights_keys")
+                and model._tied_weights_keys is not None
+                and k in model._tied_weights_keys
+            ):
+                continue
             if hasattr(v, "is_distributed") and v.is_distributed:
                 recv_table[k] = [(dist.get_rank(), tp_rank)]
             else:
@@ -1069,6 +1109,12 @@
     if args.data_parallel_rank == 0:
         state_dict = get_expected_state_dict(model)
         for (k, v) in state_dict.items():
+            if (
+                hasattr(model, "_tied_weights_keys")
+                and model._tied_weights_keys is not None
+                and k in model._tied_weights_keys
+            ):
+                continue
             if sharding_group.nranks > 1:
                 static_name = struct2static_name_mappings[k]
                 param_rank = param2rank.get(static_name, None)
@@ -1196,7 +1242,15 @@
         _, typename = key.split("/")
         typename_set.add(typename)
     struct2static_name_mappings = {k: v.name for k, v in get_expected_state_dict(model).items()}
-    static2struct_name_mappings = {v.name: k for k, v in get_expected_state_dict(model).items()}
+    static2struct_name_mappings = {}
+    for k, v in get_expected_state_dict(model).items():
+        if (
+            hasattr(model, "_tied_weights_keys")
+            and model._tied_weights_keys is not None
+            and k in model._tied_weights_keys
+        ):
+            continue
+        static2struct_name_mappings[v.name] = k
     # Get send_table and recv_table. The send table indicates which workers are responsible for sending tensors, and the recv table indicates which workers should receive the tensors.
     send_table, recv_table = create_optimizer_dispatch_table(
         args,
@@ -1349,6 +1403,9 @@
     loaded_keys = sharded_metadata["all_checkpoint_keys"]
     model_state_dict = get_expected_state_dict(model)
     expected_keys = set(list(model_state_dict.keys()))
+    if hasattr(model, "_tied_weights_keys") and model._tied_weights_keys is not None:
+        for key in model._tied_weights_keys:
+            expected_keys.remove(key)
     missing_keys = expected_keys - set(loaded_keys)
 
     if len(missing_keys) > 0:
@@ -1656,6 +1713,12 @@
         tensor_bytes_dict = {}
         model_state_dict = get_expected_state_dict(model_to_save)
         for (k, v) in state_dict.items():
+            if (
+                hasattr(model_to_save, "_tied_weights_keys")
+                and model_to_save._tied_weights_keys is not None
+                and k in model_to_save._tied_weights_keys
+            ):
+                continue
             model_v = model_state_dict[k.split("/")[0]] if is_optimizer else v
             if hasattr(model_v, "is_distributed") and model_v.is_distributed:
                 tensor_bytes_dict[k] = v.numel().item() * tp_size * dtype_byte_size(v.dtype)

paddlenlp/trainer/trainer.py

@@ -2457,6 +2457,7 @@
             # Load in optimizer and scheduler states
             self.optimizer.set_state_dict(opt_state_dict)
         else:
+            optimizer_name = _add_variant(OPTIMIZER_NAME, self.args.optimizer_name_suffix)
             raise ValueError(f"optimizer-state-dict not found, opt: {os.path.join(checkpoint, optimizer_name)}.")
 
         if not self.args.ignore_load_lr_and_optim:

paddlenlp/transformers/chatglm/modeling.py

@@ -789,7 +789,7 @@ def forward(self, hidden_states):
 
 class ChatGLMForCausalLM(ChatGLMPretrainedModel):
     _keys_to_ignore_on_save = [r"lm_head.decoder_weight"]
-    _tied_weights_keys = ["lm_head.weight"]
+    _tied_weights_keys = ["lm_head.decoder_weight"]
 
     def __init__(self, config: ChatGLMConfig):
         super(ChatGLMForCausalLM, self).__init__(config)