fix rm server

Former-commit-id: 81bc1638682a9fd01518f9f25250a6b584d2a9e6
2024-01-03 15:30:46 +08:00
parent e85c205a81
commit ccc5b324fe
3 changed files with 4 additions and 2 deletions
--- a/src/llmtuner/model/patcher.py
+++ b/src/llmtuner/model/patcher.py
@@ -276,5 +276,6 @@ def patch_valuehead_model(model: "AutoModelForCausalLMWithValueHead") -> None:

    ignore_modules = [name for name, _ in model.named_parameters() if "pretrained_model" in name]
    setattr(model, "_keys_to_ignore_on_save", ignore_modules)
+    setattr(model, "_no_split_modules", getattr(model.pretrained_model, "_no_split_modules", None))
    setattr(model, "tie_weights", MethodType(tie_weights, model))
    setattr(model, "get_input_embeddings", MethodType(get_input_embeddings, model))