File size: 8,284 Bytes
b65f8a1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 |
# Modified from Huggingface trl package AutoModelForCausalLMWithValueHead class
# Enabling better customization for generalizable reward modeling
import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM
from trl import PreTrainedModelWrapper
class ValueHead(nn.Module):
def __init__(self, config, **kwargs):
super().__init__()
if not hasattr(config, "summary_dropout_prob"):
summary_dropout_prob = kwargs.pop("summary_dropout_prob", 0.1)
else:
summary_dropout_prob = config.summary_dropout_prob
self.dropout = nn.Dropout(summary_dropout_prob) if summary_dropout_prob else nn.Identity()
# some models such as OPT have a projection layer before the word embeddings - e.g. OPT-350m
if hasattr(config, "hidden_size"):
hidden_size = config.hidden_size
if hasattr(config, "word_embed_proj_dim"):
hidden_size = config.word_embed_proj_dim
elif hasattr(config, "is_encoder_decoder"):
if config.is_encoder_decoder and hasattr(config, "decoder"):
if hasattr(config.decoder, "hidden_size"):
hidden_size = config.decoder.hidden_size
# get vhead config
if hasattr(config, "vhead_layer_type"): # config from json first
self.layer_type = config.vhead_layer_type
else:
self.layer_type = kwargs.pop("vhead_layer_type", 'mlp')
if hasattr(config, 'vhead_num_neurons'):
num_neurons = config.vhead_num_neurons
else:
num_neurons = kwargs.pop("vhead_num_neurons", 1024)
if hasattr(config, 'vhead_num_layers'):
num_layers = config.vhead_num_layers
else:
num_layers = kwargs.pop("vhead_num_layers", 1)
if self.layer_type == 'linear':
self.summary = nn.Linear(hidden_size, 1)
else:
module_lis = []
input_neurons = hidden_size
for i in range(num_layers):
module_lis.extend([nn.Linear(input_neurons, num_neurons), nn.ReLU()])
input_neurons = num_neurons
module_lis.append(nn.Linear(num_neurons, 1))
self.summary = nn.Sequential(*module_lis)
self.flatten = nn.Flatten()
def forward(self, hidden_states):
output = self.dropout(hidden_states)
if (self.layer_type == 'linear' and output.dtype != self.summary.weight.dtype):
output = output.to(self.summary.weight.dtype)
elif (self.layer_type != 'linear' and output.dtype != self.summary[0].weight.dtype):
output = output.to(self.summary[0].weight.dtype)
output = self.summary(output)
return output
class AutoModelForCausalLMWithValueHead(PreTrainedModelWrapper):
transformers_parent_class = AutoModelForCausalLM
lm_head_namings = ["lm_head", "embed_out"]
supported_args = (
"summary_dropout_prob",
"v_head_initializer_range",
"v_head_init_strategy",
"layer_type",
'num_neurons',
'num_layers',
)
def __init__(self, pretrained_model, **kwargs):
r"""
Initializes the model.
"""
super().__init__(pretrained_model, **kwargs)
v_head_kwargs, _, _ = self._split_kwargs(kwargs)
if not any(hasattr(self.pretrained_model, attribute) for attribute in self.lm_head_namings):
raise ValueError("The model does not have a language model head, please use a model that has one.")
self.v_head = ValueHead(self.pretrained_model.config, **v_head_kwargs)
self._init_weights(**v_head_kwargs)
def _init_weights(self, **kwargs):
r"""
Initializes the weights of the value head.
"""
initializer_range = kwargs.pop("v_head_initializer_range", 0.2)
# random init by default
init_strategy = kwargs.pop("v_head_init_strategy", None)
if init_strategy is None:
# do nothing
pass
elif init_strategy == "normal":
self.v_head.summary.weight.data.normal_(mean=0.0, std=initializer_range)
self.v_head.summary.bias.data.zero_()
def forward(
self,
input_ids=None,
past_key_values=None,
attention_mask=None,
**kwargs,
):
kwargs["output_hidden_states"] = True # this had already been set in the LORA / PEFT examples
kwargs["past_key_values"] = past_key_values
if self.is_peft_model and self.pretrained_model.active_peft_config.peft_type == "PREFIX_TUNING":
kwargs.pop("past_key_values")
base_model_output = self.pretrained_model(
input_ids=input_ids,
attention_mask=attention_mask,
**kwargs,
)
last_hidden_state = base_model_output.hidden_states[-1]
lm_logits = base_model_output.logits
loss = base_model_output.loss
if (hasattr(self.v_head.summary, 'weight') and last_hidden_state.device != self.v_head.summary.weight.device):
last_hidden_state = last_hidden_state.to(self.v_head.summary.weight.device)
elif not hasattr(self.v_head.summary, 'weight') and (last_hidden_state.device != self.v_head.summary[0].weight.device):
last_hidden_state = last_hidden_state.to(self.v_head.summary[0].weight.device)
# use the last token value as reward
if torch.any(attention_mask[:, 0] == 0):
# left padding
last_index = attention_mask.shape[-1] - 1
else:
# right padding
last_index = attention_mask.sum(dim=-1) - 1
value = self.v_head(last_hidden_state).squeeze(-1)[torch.arange(len(last_hidden_state)), last_index]
# force upcast in fp32 if logits are in half-precision
if lm_logits.dtype != torch.float32:
lm_logits = lm_logits.float()
return (lm_logits, loss, value)
def generate(self, *args, **kwargs):
return self.pretrained_model.generate(*args, **kwargs)
def state_dict(self, *args, **kwargs):
pretrained_model_state_dict = self.pretrained_model.state_dict(*args, **kwargs)
v_head_state_dict = self.v_head.state_dict(*args, **kwargs)
for k, v in v_head_state_dict.items():
pretrained_model_state_dict[f"v_head.{k}"] = v
return pretrained_model_state_dict
def push_to_hub(self, *args, **kwargs):
setattr(self.pretrained_model, "v_head", self.v_head)
return self.pretrained_model.push_to_hub(*args, **kwargs)
def post_init(self, state_dict):
for k in list(state_dict.keys()):
if "v_head." in k:
state_dict[k.replace("v_head.", "")] = state_dict.pop(k)
self.v_head.load_state_dict(state_dict, strict=False)
del state_dict
if hasattr(self.pretrained_model, "hf_device_map"):
if (
"cpu" in self.pretrained_model.hf_device_map.values()
or "disk" in self.pretrained_model.hf_device_map.values()
):
raise ValueError(
"The model is offloaded on CPU or disk - CPU & disk offloading is not supported for ValueHead models."
)
first_device = list(set(self.pretrained_model.hf_device_map.values()))[0]
self.v_head = self.v_head.to(first_device)
def set_device_hook(module, input, outputs):
new_output = ()
for output in outputs:
if isinstance(output, torch.Tensor):
new_output += (output.to(first_device),)
else:
new_output += (output,)
return new_output
self.register_forward_hook(set_device_hook)
self.is_sequential_parallel = True
@classmethod
def register_for_auto_class(cls, auto_class="AutoModel"):
if not isinstance(auto_class, str):
auto_class = auto_class.__name__
import transformers.models.auto as auto_module
if not hasattr(auto_module, auto_class):
raise ValueError(f"{auto_class} is not a valid auto class.")
cls._auto_class = auto_class
|