[Unified Checkpoint] Checkpoint compression

paddlenlp/peft/lora/lora_model.py

@@ -262,7 +262,7 @@ def from_pretrained(cls, model, lora_path, **kwargs):
                     pre_tensor_parallel_split = True
                     tp_actions = lora_model._get_tensor_parallel_convert_actions(loaded_keys, is_split=True)
                 state_dict = load_state_dict(
-                    shard_file, tp_actions if pre_tensor_parallel_split else None, expected_keys
+                    shard_file, tp_actions if pre_tensor_parallel_split else None, expected_keys, ckpt_quant_stage=model.config.ckpt_quant_stage
                 )
                 error_msgs += _load_state_dict_into_model(lora_model.model, state_dict, "")
                 del state_dict

paddlenlp/peft/prefix/prefix_model.py

@@ -333,7 +333,7 @@ def from_pretrained(
                     pre_tensor_parallel_split = True
                     tp_actions = prefix_model._get_tensor_parallel_convert_actions(is_split=True)
                 state_dict = load_state_dict(
-                    shard_file, tp_actions if pre_tensor_parallel_split else None, expected_keys
+                    shard_file, tp_actions if pre_tensor_parallel_split else None, expected_keys, ckpt_quant_stage=model.config.ckpt_quant_stage
                 )
                 error_msgs += _load_state_dict_into_model(prefix_model.prefix_encoder, state_dict, "")
                 del state_dict