Former-commit-id: 8c0f5d1db29862277d84aa128b424b7d0f2b187f
This commit is contained in:
hiyouga 2024-05-08 17:12:56 +08:00
parent 400ae144a4
commit 4ce4172c87

View File

@ -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"])