num items in batch

Author: shivam15s

src/liger_kernel/chunked_loss/cpo_loss.py

@@ -9,7 +9,7 @@
 class LigerFusedLinearCPOFunction(LigerFusedLinearPreferenceBase):
 
     @staticmethod
-    def preference_loss_fn(chosen_logps, rejected_logps, full_target, beta=0.1):
+    def preference_loss_fn(chosen_logps, rejected_logps, num_items_in_batch, beta=0.1):
         """
         Paper: https://arxiv.org/pdf/2401.08417
 
@@ -28,11 +28,11 @@ def preference_loss_fn(chosen_logps, rejected_logps, full_target, beta=0.1):
         Args:
             chosen_logps (torch.Tensor): Avg log probabilities of chosen tokens. Shape: (batch_size,).
             rejected_logps (torch.Tensor): Avg log probabilities of rejected tokens. Shape: (batch_size,).
-            full_target (torch.Tensor): Non chunked full target tensor
+            num_items_in_batch (int): Number of items in the batch.
             beta (float): Weight for the CPO loss
         """
         logits = beta * (chosen_logps - rejected_logps)
-        loss = F.logsigmoid(logits).sum() / (full_target.shape[0] // 2)
+        loss = F.logsigmoid(logits).sum() / (num_items_in_batch // 2)
         return loss
 
     @staticmethod

src/liger_kernel/chunked_loss/dpo_loss.py

@@ -12,7 +12,7 @@ class LigerFusedLinearDPOFunction(LigerFusedLinearPreferenceBase):
     def preference_loss_fn(
         chosen_logps,
         rejected_logps,
-        full_target,
+        num_items_in_batch,
         ref_chosen_logps=None,
         ref_rejected_logps=None,
         beta=0.1,
@@ -34,7 +34,7 @@ def preference_loss_fn(
         Args:
             chosen_logps: Log probabilities of chosen tokens (batch_size,)
             rejected_logps: Log probabilities of rejected tokens (batch_size,)
-            full_target: Non chunked full target tensor
+            num_items_in_batch (int): Number of items in the batch.
             ref_chosen_logps: Reference log probs of chosen tokens (batch_size,)
             ref_rejected_logps: Reference log probs of rejected tokens (batch_size,)
             beta: Weight for the direct preference loss
@@ -49,7 +49,7 @@ def preference_loss_fn(
         rejected_logratios = rejected_logps - ref_rejected_logps
 
         logits_diff = beta * (chosen_logratios - rejected_logratios)
-        loss = -F.logsigmoid(logits_diff).sum() / (full_target.shape[0] // 2)
+        loss = -F.logsigmoid(logits_diff).sum() / (num_items_in_batch // 2)
         return loss
 
     @staticmethod

src/liger_kernel/chunked_loss/fused_linear_preference.py

@@ -387,7 +387,7 @@ def _compute_loss(
             loss_kwargs["ref_rejected_logps"] = ref_rejected_logps
 
         preference_loss_outputs = preference_loss_fn(
-            chosen_logps, rejected_logps, full_target, beta=beta, **loss_kwargs
+            chosen_logps, rejected_logps, full_target.shape[0], beta=beta, **loss_kwargs
         )
         if isinstance(preference_loss_outputs, tuple):
             preference_loss, *aux_outputs = preference_loss_outputs

src/liger_kernel/chunked_loss/orpo_loss.py

@@ -9,7 +9,7 @@
 class LigerFusedLinearORPOFunction(LigerFusedLinearPreferenceBase):
 
     @staticmethod
-    def preference_loss_fn(chosen_logps, rejected_logps, full_target, beta=0.1):
+    def preference_loss_fn(chosen_logps, rejected_logps, num_items_in_batch, beta=0.1):
         """
         Paper: https://arxiv.org/pdf/2403.07691
 
@@ -28,21 +28,21 @@ def preference_loss_fn(chosen_logps, rejected_logps, full_target, beta=0.1):
         Args:
             chosen_logps (torch.Tensor): Avg log probabilities of chosen tokens. Shape: (batch_size,).
             rejected_logps (torch.Tensor): Avg log probabilities of rejected tokens. Shape: (batch_size,).
-            full_target (torch.Tensor): Non chunked full target tensor
+            num_items_in_batch (int): Number of items in the batch.
             beta (float): Weight for the odds ratio loss.
         """
         log_odds = (chosen_logps - rejected_logps) - (
             torch.log1p(-torch.exp(chosen_logps))
             - torch.log1p(-torch.exp(rejected_logps))
         )
         ratio = F.logsigmoid(log_odds)
-        loss = beta * ratio.sum() / (full_target.shape[0] // 2)
+        loss = beta * ratio.sum() / (num_items_in_batch // 2)
 
         chosen_rewards = beta * chosen_logps
         rejected_rewards = beta * rejected_logps
 
-        log_odds_ratio = torch.sum(ratio) / (full_target.shape[0] // 2)
-        log_odds_chosen = torch.sum(log_odds) / (full_target.shape[0] // 2)
+        log_odds_ratio = torch.sum(ratio) / (num_items_in_batch // 2)
+        log_odds_chosen = torch.sum(log_odds) / (num_items_in_batch // 2)
 
         return loss, chosen_rewards, rejected_rewards, log_odds_ratio, log_odds_chosen
 

src/liger_kernel/chunked_loss/simpo_loss.py

@@ -10,7 +10,7 @@ class LigerFusedLinearSimPOFunction(LigerFusedLinearPreferenceBase):
 
     @staticmethod
     def preference_loss_fn(
-        chosen_logps, rejected_logps, full_target, beta=0.1, gamma=0.5
+        chosen_logps, rejected_logps, num_items_in_batch, beta=0.1, gamma=0.5
     ):
         """
         Paper: https://arxiv.org/pdf/2405.14734
@@ -30,12 +30,12 @@ def preference_loss_fn(
         Args:
             chosen_logps (torch.Tensor): Avg log probabilities of chosen tokens. Shape: (batch_size,).
             rejected_logps (torch.Tensor): Avg log probabilities of rejected tokens. Shape: (batch_size,).
-            full_target: Non chunked full target tensor
+            num_items_in_batch (int): Number of items in the batch.
             beta (float): beta weight
             gamma (float): gemma margin term
         """
         logits = beta * (chosen_logps - rejected_logps) - gamma
-        loss = F.logsigmoid(logits).sum() / (full_target.shape[0] // 2)
+        loss = F.logsigmoid(logits).sum() / (num_items_in_batch // 2)
         return loss
 
     @staticmethod