Upload model.py

model.py

@@ -0,0 +1,210 @@
+# 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
+