Revert "Add ref_input parameter to support separate inputs for refere…

…nce model" (#469) Reverts #467 until the test is fixed cc @shivam15s
linkedin · Dec 11, 2024 · 969ce3a · 969ce3a
1 parent eee40c5
commit 969ce3a
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/src/liger_kernel/chunked_loss/fused_linear_preference.py b/src/liger_kernel/chunked_loss/fused_linear_preference.py
@@ -29,7 +29,7 @@ def forward(
         compute_nll_loss=True,
         compiled=True,
         use_ref_model=False,
-        ref_input=None,
+        # TODO: ref input
         ref_weight=None,
         ref_bias=None,
         **loss_kwargs,
@@ -59,7 +59,6 @@ def forward(
             compute_nll_loss (bool): Whether to compute NLL loss.
             compiled (bool): Whether to use torch compile for chunk accumulation.
             use_ref_model (bool): Whether to use a reference model for the alignment loss.
-            ref_input (torch.Tensor): Reference input tensor. Shape: (batch_size, seq_len, hidden_size).
             ref_weight (torch.Tensor): Reference weight tensor. Shape: (vocab_size, hidden_size).
             ref_bias (torch.Tensor, optional): Reference bias tensor. Shape: (vocab_size,).
             loss_kwargs (dict): Other possible arguments that a loss function might need
@@ -93,7 +92,6 @@ def forward(
             compute_nll_loss=compute_nll_loss,
             full_target=target,
             use_ref_model=use_ref_model,
-            ref_input=ref_input,
             ref_weight=ref_weight,
             ref_bias=ref_bias,
             **loss_kwargs,
@@ -303,7 +301,6 @@ def _compute_loss(
         beta=0.1,
         compute_nll_loss=True,
         use_ref_model=False,
-        ref_input=None,
         ref_weight=None,
         ref_bias=None,
         **loss_kwargs,
@@ -322,7 +319,6 @@ def _compute_loss(
             beta (float): Weight for the preference loss.
             compute_nll_loss (bool): Whether to compute NLL loss.
             use_ref_model (bool): Whether to use a reference model for the alignment loss.
-            ref_input (torch.Tensor): Reference input tensor. Shape: (2 * chunk_size, sequence_length, hidden_size).
             ref_weight (torch.Tensor): Reference weight tensor. Shape: (vocab_size, hidden_size).
             ref_bias (torch.Tensor, optional): Reference bias tensor. Shape: (vocab_size,).
             loss_kwargs (dict): Additional arguments for the loss function.
@@ -361,7 +357,7 @@ def _compute_loss(
                     ref_rejected_logits,
                     ref_chosen_nll_loss,
                 ) = LigerFusedLinearPreferenceBase.chunk_forward(
-                    ref_input,
+                    input_chunk,
                     ref_weight,
                     target_chunk,
                     ref_bias,