[Cherry-pick] fix multi-threading load, fix single card load. (#9560)

* Fix multi-threading load_state_dict (#9464)

* Update model_utils.py

* Update model_utils.py

* [Unified Checkpoint] fix single card loading without master weights (#9540)

Author: DesmonDay

paddlenlp/trainer/unified_checkpoint/load_save_single_card.py

@@ -225,6 +225,8 @@ def load_single_card_optimizer(model, optimizer, resume_from_checkpoint: str):
                 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] = state_dict_optim.pop(key)
         returned_optim_state_dict[key_name].name = key_name
     if has_master_weights:

paddlenlp/transformers/model_utils.py

@@ -437,7 +437,7 @@ def load_state_dict(
             thread_num = int(os.environ.get("LOAD_STATE_DICT_THREAD_NUM", "1"))
             if thread_num > 1:
                 logger.info(f"Set loading state_dict thread num to {thread_num}")
-            state_dict = {}
+            state_dict, scale_dict = {}, {}
             if thread_num <= 1:
                 with safe_open(checkpoint_file, framework="np") as f:
                     state_dict, scale_dict = _load_part_state_dict(
@@ -464,9 +464,9 @@ def load_state_dict(
                         for keys in keys_groups
                     }
                     for future in concurrent.futures.as_completed(future_to_key):
-                        state_dict, scale_dict = future.result()
-                        state_dict.update(state_dict)
-                        scale_dict.update(scale_dict)
+                        res_state_dict, res_scale_dict = future.result()
+                        state_dict.update(res_state_dict)
+                        scale_dict.update(res_scale_dict)
 
             if device == "cpu":
                 for k in list(state_dict.keys()):