423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-08-22 13:42:51 +08:00
Code Issues Packages Projects Releases Wiki Activity

upcast logits

Browse Source 
Former-commit-id: c13ae2df19ed4cdc849bef55d04225e1a98c19b5
This commit is contained in:
hiyouga 2024-07-02 22:32:05 +08:00
parent cc31014002
commit 4828bed837
2 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/llamafactory/train/ppo/trainer.py
View File

@ -407,7 +407,7 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
values = torch.transpose(values, 0, 1)
rewards = values.gather(dim=-1, index=(batch["attention_mask"].sum(dim=-1, keepdim=True) - 1))
return rewards.to(torch.float32).detach().cpu() # use fp32 type
return rewards.float().detach() # use fp32 type
@PPODecorators.empty_device_cache()
def batched_forward_pass(

2
src/llamafactory/train/rm/trainer.py
View File

@ -99,7 +99,7 @@ class PairwiseTrainer(Trainer):
chosen_scores = chosen_rewards.gather(dim=-1, index=(chosen_masks.sum(dim=-1, keepdim=True) - 1))
rejected_scores = rejected_rewards.gather(dim=-1, index=(rejected_masks.sum(dim=-1, keepdim=True) - 1))
chosen_scores, rejected_scores = chosen_scores.squeeze(), rejected_scores.squeeze()
loss = -torch.nn.functional.logsigmoid(chosen_scores - rejected_scores).mean()
loss = -torch.nn.functional.logsigmoid(chosen_scores.float() - rejected_scores.float()).mean()
if return_outputs:
return loss, (loss, chosen_scores, rejected_scores)
else: