fix #3625

2024-05-08 17:12:56 +08:00 · 2024-05-08 17:12:56 +08:00 · d9cdddd19c
parent 48ee46dac1
commit d9cdddd19c
1 changed files with 4 additions and 0 deletions
--- a/src/llmtuner/model/utils/valuehead.py
+++ b/src/llmtuner/model/utils/valuehead.py
@ -57,3 +57,7 @@ def prepare_valuehead_model(model: "PreTrainedModel") -> None:
    if getattr(model.config, "model_type", None) == "chatglm":
        setattr(model, "lm_head", model.transformer.output_layer)
        setattr(model, "_keys_to_ignore_on_save", ["lm_head.weight"])
+
+    if getattr(model.config, "model_type", None) == "internlm2":
+        setattr(model, "lm_head", model.output)
+        setattr(model, "_keys_to_ignore_on_save", ["lm_head.weight"])