Add hook for lm_head tuning.

app.py

@@ -3,6 +3,7 @@ import torch
 from transformers import AutoTokenizer, GenerationConfig, AutoModel
 
 model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True, revision="658202d").float()
+setattr(model, "lm_head_raw", model.lm_head)
 tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True, revision="658202d")
 
 from peft import PeftModel