FP32 Training Works

model/reward/instructor/configs/rankgen-t5-base.yml

@@ -2,8 +2,9 @@ model_name: kalpeshk2011/rankgen-t5-base-all
 tokenizer_name: google/t5-v1_1-base
 learning_rate: 6e-6
 gradient_checkpointing: false
+fp16: false
 gradient_accumulation_steps: 16
-per_device_train_batch_size: 3
+per_device_train_batch_size: 2
 warmup_steps: 600
 freeze_layer: 20
 eval_steps: 200

model/reward/instructor/models.py

@@ -12,11 +12,13 @@ class RankGenModel(torch.nn.Module):
     self.model = AutoModel.from_pretrained(self.rankgen_hf_hub, trust_remote_code=True)
 
   def forward(self, prefixes, suffixes):
+    # print(list(self.model.parameters()))
+    # raise Exception("stop")
     embedded_prefixes = self.model(**prefixes)
     embedded_suffixes = self.model(**suffixes)
     # take dot product of each row independently
     dot_products = torch.sum(embedded_prefixes * embedded_suffixes, dim=1)
     
-    print(f"{prefixes=}, {suffixes=}, {embedded_prefixes=}, {embedded_suffixes=}, {dot_products=}")
-
+    # print(f"{embedded_prefixes.shape=}, {embedded_suffixes.shape=}, {prefixes['input_ids'].shape=}, {suffixes['input_ids'].shape=}, {embedded_prefixes=}, {embedded_suffixes=}, {dot_products=}")
+    # raise Exception("stop")
     return dot_products

model/reward/instructor/rank_datasets.py

@@ -33,6 +33,7 @@ class RankGenCollator():
     tokenizer: PreTrainedTokenizerBase
     padding: Union[bool, str, PaddingStrategy] = True
     max_length: Optional[int] = None
+    max_examples: Optional[int] = None
 
     def __call__(self, batch : list[dict[str, str]]) -> dict[str, torch.Tensor]:
         prefixes = []

model/reward/instructor/trainer.py

@@ -50,7 +50,6 @@ class RankLoss(nn.Module):
 
     def forward(self, pos, neg):
         loss = -self.log_sigmoid(pos - neg + self.eps).mean()
-        print(f"in loss {pos=}, {neg=}, {loss=}")
         return loss
 
 
@@ -90,7 +89,6 @@ class RankTrainer(Trainer):
     def compute_loss(self, model, inputs, return_outputs=False):
         # forward pass
         if "rankgen" in self.model_name:
-            print(f"{inputs=}")
             positive_outputs = model(inputs["prefix"], inputs["positive"])
             negative_outputs = model(inputs["prefix"], inputs["negative"])
             if self.loss_function == "rank":
@@ -171,7 +169,7 @@ if __name__ == "__main__":
         loss_function=training_conf["loss"],
         learning_rate=training_conf["learning_rate"],
         # half_precision_backend="apex",
-        fp16=True,
+        fp16=training_conf["fp16"] if "fp16" in training_conf else True,
         gradient_checkpointing=training_conf["gradient_checkpointing"],
         gradient_accumulation_steps=training_conf["gradient_accumulation_steps"],
         per_device_train_batch_size=training_conf["per_device_train_batch_size"],