Skip to content

Commit

Permalink
return chosen_rewards, rejected_rewards
Browse files Browse the repository at this point in the history
  • Loading branch information
kashif committed Jan 15, 2025
1 parent 5a26e54 commit 142fa87
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 2 deletions.
5 changes: 4 additions & 1 deletion src/liger_kernel/chunked_loss/dpo_loss.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,9 +45,12 @@ def preference_loss_fn(
chosen_logratios = chosen_logps - ref_chosen_logps
rejected_logratios = rejected_logps - ref_rejected_logps

chosen_rewards = beta * (chosen_logps - ref_chosen_logps)
rejected_rewards = beta * (rejected_logps - ref_rejected_logps)

logits_diff = beta * (chosen_logratios - rejected_logratios)
loss = -F.logsigmoid(logits_diff).sum() / (full_target.shape[0] // 2)
return loss
return loss, chosen_rewards, rejected_rewards

@staticmethod
def forward(
Expand Down
5 changes: 4 additions & 1 deletion test/chunked_loss/test_dpo_loss.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,9 +56,12 @@ def alignment_loss(
chosen_logratios = policy_chosen_logps - ref_chosen_logps
rejected_logratios = policy_rejected_logps - ref_rejected_logps

chosen_rewards = self.beta * (policy_chosen_logps - ref_chosen_logps)
rejected_rewards = self.beta * (policy_rejected_logps - ref_rejected_logps)

logits_diff = self.beta * (chosen_logratios - rejected_logratios)
losses = -F.logsigmoid(logits_diff)
return losses
return losses, chosen_rewards, rejected_rewards


class TorchLMHeadDPO(torch.nn.Module):
Expand Down

0 comments on commit 142fa87

Please sign in to comment.