Former-commit-id: e131bc03e05ccae3c6ad8bb42ccf2cdcc2cf3cea
This commit is contained in:
hiyouga
2023-09-14 17:56:58 +08:00
parent 3e5555502a
commit 0ca36a0f8d
4 changed files with 4 additions and 3 deletions

View File

@@ -10,8 +10,8 @@ if TYPE_CHECKING:
def replace_model(model: "AutoModelForCausalLMWithValueHead", target: Literal["default", "reward"]) -> None:
if target == "reward": # save default head temporarily
valuehead_state_dict = model.v_head.state_dict()
setattr(model, "default_head_weight", valuehead_state_dict["summary.weight"].clone())
setattr(model, "default_head_bias", valuehead_state_dict["summary.bias"].clone())
setattr(model, "default_head_weight", valuehead_state_dict["summary.weight"].detach().clone())
setattr(model, "default_head_bias", valuehead_state_dict["summary.bias"].detach().clone())
model.pretrained_model.set_adapter(target) # set the LoRA adapter to be active
model.v_head.load_state_dict({