From 568a42066a80198f197fb0ac42c24af3cb334795 Mon Sep 17 00:00:00 2001 From: Bobak Hashemi Date: Tue, 3 Jan 2023 00:53:07 -0500 Subject: [PATCH] FP32 Training Works --- model/reward/instructor/configs/rankgen-t5-base.yml | 3 ++- model/reward/instructor/models.py | 6 ++++-- model/reward/instructor/rank_datasets.py | 1 + model/reward/instructor/trainer.py | 4 +--- 4 files changed, 8 insertions(+), 6 deletions(-) diff --git a/model/reward/instructor/configs/rankgen-t5-base.yml b/model/reward/instructor/configs/rankgen-t5-base.yml index 7dd39777..6776ad47 100644 --- a/model/reward/instructor/configs/rankgen-t5-base.yml +++ b/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 diff --git a/model/reward/instructor/models.py b/model/reward/instructor/models.py index 699f3566..dc7692bf 100644 --- a/model/reward/instructor/models.py +++ b/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 \ No newline at end of file diff --git a/model/reward/instructor/rank_datasets.py b/model/reward/instructor/rank_datasets.py index 3b995a7d..965893ce 100644 --- a/model/reward/instructor/rank_datasets.py +++ b/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 = [] diff --git a/model/reward/instructor/trainer.py b/model/reward/instructor/trainer.py index 5bb1017a..c6f58f66 100644 --- a/model/reward/instructor/trainer.py +++ b/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"],