fix #2161

Former-commit-id: 9acd5a2b678cd07f8e3b48eca76c4cbacb559e37

src/llmtuner/extras/callbacks.py

@@ -1,78 +1,23 @@
 import os
 import json
 import time
-import torch
-from typing import TYPE_CHECKING, Dict
+from typing import TYPE_CHECKING
 from datetime import timedelta
-
-from transformers import PreTrainedModel, TrainerCallback
-from transformers.utils import WEIGHTS_NAME, SAFE_WEIGHTS_NAME
+from transformers import TrainerCallback
 from transformers.trainer_utils import has_length, PREFIX_CHECKPOINT_DIR
-from peft import PeftModel
 
-from llmtuner.extras.constants import LOG_FILE_NAME, V_HEAD_WEIGHTS_NAME, V_HEAD_SAFE_WEIGHTS_NAME
+from llmtuner.extras.constants import LOG_FILE_NAME
 from llmtuner.extras.logging import get_logger
+from llmtuner.extras.misc import fix_valuehead_checkpoint
 
 
 if TYPE_CHECKING:
     from transformers import TrainingArguments, TrainerState, TrainerControl
-    from trl import AutoModelForCausalLMWithValueHead
 
 
 logger = get_logger(__name__)
 
 
-def _fix_valuehead_checkpoint(
-    model: "AutoModelForCausalLMWithValueHead",
-    output_dir: str,
-    safe_serialization: bool
-) -> None:
-    r"""
-    The model is already unwrapped.
-
-    There are three cases:
-    1. full tuning without ds_zero3: state_dict = {"model.layers.*": ..., "v_head.summary.*": ...}
-    2. lora tuning without ds_zero3: state_dict = {"v_head.summary.*": ...}
-    3. under deepspeed zero3: state_dict = {"pretrained_model.model.layers.*": ..., "v_head.summary.*": ...}
-
-    We assume `stage3_gather_16bit_weights_on_model_save=true`.
-    """
-    if not isinstance(model.pretrained_model, (PreTrainedModel, PeftModel)):
-        return
-
-    if safe_serialization:
-        from safetensors import safe_open
-        from safetensors.torch import save_file
-        path_to_checkpoint = os.path.join(output_dir, SAFE_WEIGHTS_NAME)
-        with safe_open(path_to_checkpoint, framework="pt", device="cpu") as f:
-            state_dict: Dict[str, torch.Tensor] = {key: f.get_tensor(key) for key in f.keys()}
-    else:
-        path_to_checkpoint = os.path.join(output_dir, WEIGHTS_NAME)
-        state_dict: Dict[str, torch.Tensor] = torch.load(path_to_checkpoint, map_location="cpu")
-
-    decoder_state_dict = {}
-    v_head_state_dict = {}
-    for name, param in state_dict.items():
-        if name.startswith("v_head."):
-            v_head_state_dict[name] = param
-        else:
-            decoder_state_dict[name.replace("pretrained_model.", "")] = param
-
-    os.remove(path_to_checkpoint)
-    model.pretrained_model.save_pretrained(
-        output_dir,
-        state_dict=decoder_state_dict or None,
-        safe_serialization=safe_serialization
-    )
-
-    if safe_serialization:
-        save_file(v_head_state_dict, os.path.join(output_dir, V_HEAD_SAFE_WEIGHTS_NAME), metadata={"format": "pt"})
-    else:
-        torch.save(v_head_state_dict, os.path.join(output_dir, V_HEAD_WEIGHTS_NAME))
-
-    logger.info("Value head model saved at: {}".format(output_dir))
-
-
 class FixValueHeadModelCallback(TrainerCallback):
 
     def on_save(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
@@ -80,21 +25,12 @@ class FixValueHeadModelCallback(TrainerCallback):
         Event called after a checkpoint save.
         """
         if args.should_save:
-            _fix_valuehead_checkpoint(
+            fix_valuehead_checkpoint(
                 model=kwargs.pop("model"),
                 output_dir=os.path.join(args.output_dir, "{}-{}".format(PREFIX_CHECKPOINT_DIR, state.global_step)),
                 safe_serialization=args.save_safetensors
             )
 
-    def on_train_end(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
-        r"""
-        Event called at the end of training.
-        """
-        if args.should_save:
-            _fix_valuehead_checkpoint(
-                model=kwargs.pop("model"), output_dir=args.output_dir, safe_serialization=args.save_safetensors
-            )
-
 
 class LogCallback(TrainerCallback):