[fix] resolve conflict from main

model/reward/instructor/configs/deberta-v2-xlarge.yml

@@ -0,0 +1,15 @@
+model_name: microsoft/deberta-v2-xlarge
+learning_rate: 1e-5
+freeze_layer: 15
+scheduler: cosine
+gradient_checkpointing: false
+gradient_accumulation_steps: 16
+per_device_train_batch_size: 1
+warmup_steps: 600
+eval_steps: 200
+save_steps: 500
+max_length: 512
+num_train_epochs: 2
+datasets:
+  - webgpt
+  - hfsummary

model/reward/instructor/configs/deberta-v3-base.yml

@@ -0,0 +1,14 @@
+model_name: microsoft/deberta-v3-base
+learning_rate: 1e-5
+scheduler: cosine
+gradient_checkpointing: false
+gradient_accumulation_steps: 32
+per_device_train_batch_size: 2
+warmup_steps: 600
+eval_steps: 200
+save_steps: 500
+max_length: 512
+num_train_epochs: 2
+datasets:
+  - webgpt
+  - hfsummary

model/reward/instructor/configs/deberta-v3-large-squad2.yml

@@ -0,0 +1,13 @@
+model_name: deepset/deberta-v3-large-squad2
+learning_rate: 1e-5
+gradient_checkpointing: false
+gradient_accumulation_steps: 32
+per_device_train_batch_size: 1
+warmup_steps: 600
+eval_steps: 200
+save_steps: 500
+max_length: 512
+num_train_epochs: 2
+datasets:
+  - webgpt
+  - hfsummary

model/reward/instructor/configs/deberta-v3-large.yml

@@ -0,0 +1,14 @@
+model_name: microsoft/deberta-v3-large
+learning_rate: 1e-5
+scheduler: cosine
+gradient_checkpointing: false
+gradient_accumulation_steps: 32
+per_device_train_batch_size: 1
+warmup_steps: 600
+eval_steps: 200
+save_steps: 500
+max_length: 512
+num_train_epochs: 2
+datasets:
+  - webgpt
+  - hfsummary

model/reward/instructor/trainer.py

@@ -8,7 +8,15 @@ import torch
 from models import RankGenModel
 from rank_datasets import DataCollatorForPairRank, RankGenCollator
 from torch import nn
-from transformers import AutoModelForSequenceClassification, PreTrainedModel, Trainer, TrainingArguments
+from transformers import (
+    AdamW,
+    AutoModelForSequenceClassification,
+    PreTrainedModel,
+    Trainer,
+    TrainingArguments,
+    get_cosine_schedule_with_warmup,
+    get_linear_schedule_with_warmup,
+)
 from utils import argument_parsing, freeze_top_n_layers, get_datasets, get_tokenizer
 
 os.environ["WANDB_PROJECT"] = "reward-model"
@@ -144,7 +152,7 @@ if __name__ == "__main__":
         evaluation_strategy="steps",
         eval_steps=training_conf["eval_steps"],
         save_steps=1000,
-        report_to="local",
+        report_to="wandb",
     )
 
     tokenizer = get_tokenizer(training_conf["tokenizer_name"])
@@ -154,6 +162,21 @@ if __name__ == "__main__":
     else:
         collate_fn = DataCollatorForPairRank(tokenizer, max_length=training_conf["max_length"])
     assert len(evals) > 0
+
+    optimizer = AdamW(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay)
+    scheduler = None
+    if "scheduler" in training_conf:
+        if training_conf["scheduler"] == "linear":
+            scheduler = get_linear_schedule_with_warmup()
+        elif training_conf["scheduler"] == "cosine":
+            scheduler = get_cosine_schedule_with_warmup(
+                optimizer,
+                num_warmup_steps=args.warmup_steps,
+                num_training_steps=len(train)
+                * args.num_train_epochs
+                / (args.per_device_train_batch_size * args.gradient_accumulation_steps),
+            )
+
     trainer = RankTrainer(
         model=model,
         model_name=model_name,
@@ -164,6 +187,7 @@ if __name__ == "__main__":
         data_collator=collate_fn,
         tokenizer=tokenizer,
         compute_metrics=compute_metrics,
+        optimizers=(optimizer, scheduler),
     )
     # trainer.evaluate()
     trainer.train()