Update checkpoint_converter.py

PaddlePaddle · Jul 31, 2024 · 5fc66c2 · 5fc66c2
1 parent 3bdedab
commit 5fc66c2
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/paddlenlp/trainer/checkpoint_converter.py b/paddlenlp/trainer/checkpoint_converter.py
@@ -206,9 +206,10 @@ def gen_matadata_for_optimizer(self):
                 for k, v in state_dict.items():
                     for prame_name, opt_name in structure_name_mapping.items():
                         if opt_name in k:
-                            new_key = k.replace(opt_name, prame_name) + "_tp" + "{:02d}".format(distributed_rank[0])
-                        else:
-                            new_key = k.replace(opt_name, prame_name)
+                            if v.shape[0] > 1:
+                                new_key = k.replace(opt_name, prame_name) + "_tp" + "{:02d}".format(distributed_rank[0])
+                            else:
+                                new_key = k.replace(opt_name, prame_name)
                             renamed_state_dict[new_key] = v
                             # Calculate the local_shape.
                             cur_rank_sharded_tensor_infos[(new_key, file)] = [v.shape, str(v.dtype)]