Add files via upload

Author: maxmatical

ML tips/NLP/reward_model/__init__.py

@@ -0,0 +1,18 @@
+import torch
+from torch import nn
+from transformers import GPT2Model
+
+class RewardModel(nn.Module):
+    """
+    Using GPT for reward model like in instructGPT
+    """
+    def __init__(self, base_model: GPT2Model):
+        super().__init__()
+        self.model = base_model
+        self.reward_head = nn.Linear(base_model.config.hidden_size, 1, bias=False)
+    
+    def forward(self, input_ids, attention_mask):
+        out = self.model(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state  # bs, seq_len, hid_dim 
+        logits = self.reward_head(out[:, -1, :]).squeeze(-1) # bs
+        return logits
+        

ML tips/NLP/reward_model/reward_model.py

@@ -0,0 +1,100 @@
+from transformers import Trainer, PreTrainedModel
+from rlhf.data.data import RewardDataCollatorWithPadding
+from torch import nn
+from rlhf.optimizer.lion import DecoupledLionW
+from transformers.trainer_pt_utils import get_parameter_names 
+import torch
+from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
+
+
+class RewardTrainer(Trainer):
+    def __init__(
+        self, 
+        model: PreTrainedModel,
+        args,
+        data_collator = None,
+        train_dataset = None,
+        eval_dataset = None,
+        tokenizer = None,
+        model_init = None,
+        compute_metrics = None,
+        callbacks = None,
+        optimizers = (None, None),
+        preprocess_logits_for_metrics = None,
+        max_length: int = 512,
+        use_lion: bool = False
+    ):
+        # data_collator = RewardDataCollatorWithPadding(tokenizer, max_length=max_length)
+        super().__init__(
+            model=model,
+            args=args,
+            data_collator=data_collator,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+            tokenizer=tokenizer,
+            model_init=model_init,
+            compute_metrics=compute_metrics,
+            callbacks=callbacks,
+            optimizers=optimizers,
+            preprocess_logits_for_metrics=preprocess_logits_for_metrics,
+        )
+        self.use_lion = use_lion
+
+    def compute_loss(self, model, inputs):
+        rewards_chosen = model(input_ids=inputs["input_ids_chosen"], attention_mask=inputs["attention_mask_chosen"])
+        rewards_rejected = model(
+            input_ids=inputs["input_ids_rejected"], attention_mask=inputs["attention_mask_rejected"]
+        )
+        loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean()
+        return loss
+
+
+    def create_optimizer(self):
+        """
+        Setup the optimizer.
+
+        We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
+        Trainer's init through `optimizers`, or subclass and override this method in a subclass.
+        """
+        opt_model = self.model
+
+        decay_parameters = get_parameter_names(opt_model, ALL_LAYERNORM_LAYERS)
+        decay_parameters = [name for name in decay_parameters if "bias" not in name]
+        optimizer_grouped_parameters = [
+            {
+                "params": [
+                    p for n, p in opt_model.named_parameters() if (n in decay_parameters and p.requires_grad)
+                ],
+                "weight_decay": self.args.weight_decay,
+            },
+            {
+                "params": [
+                    p for n, p in opt_model.named_parameters() if (n not in decay_parameters and p.requires_grad)
+                ],
+                "weight_decay": 0.0,
+            },
+        ]
+
+        # lion or adam
+        if not self.use_lion:
+            optimizer_cls = torch.optim.AdamW
+            optimizer_kwargs = {
+                "lr": self.args.learning_rate,
+                "betas": (self.args.adam_beta1, self.args.adam_beta2),
+                "eps": self.args.adam_epsilon,
+            }
+        else:
+            optimizer_cls = DecoupledLionW
+            optimizer_kwargs = {
+                "lr": self.args.learning_rate,
+                "betas": (self.args.adam_beta1, self.args.adam_beta2),
+            }
+        
+
+        self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
+
+
+        # print(f"Using optimizer {self.optimizer}")
+
+        return self.optimizer
+    

ML tips/NLP/reward_model/reward_trainer.py

@@ -0,0 +1,61 @@
+from transformers import HfArgumentParser, TrainingArguments
+from transformers import AutoTokenizer, GPT2Model
+from datasets import load_dataset 
+
+from dataclasses import dataclass, field
+from rlhf.reward_model.reward_trainer import RewardTrainer
+from rlhf.reward_model.reward_model import RewardModel
+from rlhf.data.data import RewardDataCollatorWithPadding, preprocess_function
+
+
+@dataclass
+class RMTrainerArguments(TrainingArguments):
+    """
+    Adapt transformers.TrainingArguments
+    """
+    use_lion: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "Will use the LION optimizer instead of AdamW "
+                "uses same hyperparam args (betas, wd, lr) as AdamW, except eps)."
+            )
+        },
+    )
+def main():
+    parser = HfArgumentParser(RMTrainerArguments)
+
+    trainer_args = parser.parse_args_into_dataclasses()[0]
+
+    tokenizer = AutoTokenizer.from_pretrained("gpt2")
+
+    tokenizer.pad_token = tokenizer.eos_token
+
+    collator = RewardDataCollatorWithPadding(tokenizer=tokenizer)
+
+    lm_model = GPT2Model.from_pretrained("gpt2")
+    lm_model.config.pad_token_id = tokenizer.eos_token_id
+
+    model = RewardModel(lm_model)
+
+    dataset = load_dataset("Anthropic/hh-rlhf")
+    train_ds = dataset["train"]
+
+    train_dataset = train_ds.map(
+        preprocess_function, 
+        batched=True, num_proc=1, 
+        fn_kwargs={"tokenizer": tokenizer}
+    )
+
+
+    trainer = RewardTrainer(
+        model=model,
+        args = trainer_args,
+        train_dataset = train_dataset,
+        data_collator = collator
+    )
+
+    trainer.train()
+
+if __name__ == '__main__':
+    main()