liger-kernel-nightly 0.5.2.dev20241212030605__tar.gz → 0.5.2.dev20241212033924__tar.gz

{liger_kernel_nightly-0.5.2.dev20241212030605 → liger_kernel_nightly-0.5.2.dev20241212033924}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.5.2.dev20241212030605
+Version: 0.5.2.dev20241212033924
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 

{liger_kernel_nightly-0.5.2.dev20241212030605 → liger_kernel_nightly-0.5.2.dev20241212033924}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.5.2.dev20241212030605"
+version = "0.5.2.dev20241212033924"
 description = "Efficient Triton kernels for LLM Training"
 urls = { "Homepage" = "https://github.com/linkedin/Liger-Kernel" }
 readme = { file = "README.md", content-type = "text/markdown" }

{liger_kernel_nightly-0.5.2.dev20241212030605 → liger_kernel_nightly-0.5.2.dev20241212033924}/src/liger_kernel/chunked_loss/dpo_loss.py

@@ -59,6 +59,7 @@ class LigerFusedLinearDPOFunction(LigerFusedLinearPreferenceBase):
         weight,
         target,
         bias=None,
+        ref_input=None,
         ref_weight=None,
         ref_bias=None,
         ignore_index=-100,
@@ -79,6 +80,7 @@ class LigerFusedLinearDPOFunction(LigerFusedLinearPreferenceBase):
             compute_nll_loss=compute_nll_loss,
             compiled=compiled,
             use_ref_model=use_ref_model,
+            ref_input=ref_input,
             ref_weight=ref_weight,
             ref_bias=ref_bias,
         )
@@ -86,7 +88,7 @@ class LigerFusedLinearDPOFunction(LigerFusedLinearPreferenceBase):
     @staticmethod
     def backward(ctx, *grad_output):
         grads = LigerFusedLinearPreferenceBase.backward(ctx, grad_output)[:4]
-        return *grads, None, None, None, None, None, None, None
+        return *grads, None, None, None, None, None, None, None, None
 
 
 class LigerFusedLinearDPOLoss(torch.nn.Module):
@@ -118,13 +120,21 @@ class LigerFusedLinearDPOLoss(torch.nn.Module):
         self.use_ref_model = use_ref_model
 
     def forward(
-        self, lin_weight, _input, target, bias=None, ref_weight=None, ref_bias=None
+        self,
+        lin_weight,
+        _input,
+        target,
+        bias=None,
+        ref_input=None,
+        ref_weight=None,
+        ref_bias=None,
     ):
         return LigerFusedLinearDPOFunction.apply(
             _input,
             lin_weight,
             target,
             bias,
+            ref_input,
             ref_weight,
             ref_bias,
             self.ignore_index,

{liger_kernel_nightly-0.5.2.dev20241212030605 → liger_kernel_nightly-0.5.2.dev20241212033924}/src/liger_kernel/chunked_loss/fused_linear_preference.py

@@ -29,7 +29,7 @@ class LigerFusedLinearPreferenceBase(torch.autograd.Function):
         compute_nll_loss=True,
         compiled=True,
         use_ref_model=False,
-        # TODO: ref input
+        ref_input=None,
         ref_weight=None,
         ref_bias=None,
         **loss_kwargs,
@@ -97,20 +97,26 @@ class LigerFusedLinearPreferenceBase(torch.autograd.Function):
             **loss_kwargs,
         )
 
-        def fused_fwd_bwd(input_chunk, target_chunk):
+        def fused_fwd_bwd(input_chunk, target_chunk, ref_input_chunk):
             """
             Fused forward and backward pass for a chunk of input and target.
             """
             if bias is not None:
                 return torch.func.grad_and_value(
                     compute_loss, argnums=(0, 1, 3), has_aux=True
-                )(input_chunk, weight, target_chunk, bias)
+                )(
+                    input_chunk,
+                    weight,
+                    target_chunk,
+                    bias,
+                    ref_input_chunk=ref_input_chunk,
+                )
             else:
                 return torch.func.grad_and_value(
                     compute_loss, argnums=(0, 1), has_aux=True
-                )(input_chunk, weight, target_chunk)
+                )(input_chunk, weight, target_chunk, ref_input_chunk=ref_input_chunk)
 
-        def accumulate_chunk(input_chunk, target_chunk):
+        def accumulate_chunk(input_chunk, target_chunk, ref_input_chunk=None):
             if bias is not None:
                 (chunk_grad_input, chunk_grad_weight, chunk_grad_bias), (
                     chunk_loss,
@@ -122,7 +128,7 @@ class LigerFusedLinearPreferenceBase(torch.autograd.Function):
                         chunk_nll_loss,
                         *aux_outputs,
                     ),
-                ) = fused_fwd_bwd(input_chunk, target_chunk)
+                ) = fused_fwd_bwd(input_chunk, target_chunk, ref_input_chunk)
                 grad_bias.add_(chunk_grad_bias)  # accumulate bias gradient
             else:
                 (chunk_grad_input, chunk_grad_weight), (
@@ -135,7 +141,7 @@ class LigerFusedLinearPreferenceBase(torch.autograd.Function):
                         chunk_nll_loss,
                         *aux_outputs,
                     ),
-                ) = fused_fwd_bwd(input_chunk, target_chunk)
+                ) = fused_fwd_bwd(input_chunk, target_chunk, ref_input_chunk)
 
             # Accumulate gradients
             grad_weight.add_(chunk_grad_weight)
@@ -182,18 +188,43 @@ class LigerFusedLinearPreferenceBase(torch.autograd.Function):
         _rejected_input_chunks = torch.chunk(_input[len_chosen:], chunks=chunks, dim=0)
         _rejected_target_chunks = torch.chunk(target[len_chosen:], chunks=chunks, dim=0)
 
+        if use_ref_model:
+            _ref_chosen_input_chunks = torch.chunk(
+                ref_input[:len_chosen], chunks=chunks, dim=0
+            )
+            _ref_rejected_input_chunks = torch.chunk(
+                ref_input[len_chosen:], chunks=chunks, dim=0
+            )
+
         for (
             chosen_input_chunk,
             rejected_input_chunk,
             chosen_target_chunk,
             rejected_target_chunk,
+            ref_chosen_input_chunk,
+            ref_rejected_input_chunk,
         ) in zip(
             _chosen_input_chunks,
             _rejected_input_chunks,
             _chosen_target_chunks,
             _rejected_target_chunks,
+            (
+                _ref_chosen_input_chunks
+                if use_ref_model
+                else [None] * len(_chosen_input_chunks)
+            ),
+            (
+                _ref_rejected_input_chunks
+                if use_ref_model
+                else [None] * len(_rejected_input_chunks)
+            ),
         ):
             input_chunk = torch.cat([chosen_input_chunk, rejected_input_chunk], dim=0)
+            ref_input_chunk = (
+                torch.cat([ref_chosen_input_chunk, ref_rejected_input_chunk], dim=0)
+                if use_ref_model
+                else None
+            )
             target_chunk = torch.cat(
                 [chosen_target_chunk, rejected_target_chunk], dim=0
             )
@@ -202,9 +233,10 @@ class LigerFusedLinearPreferenceBase(torch.autograd.Function):
             torch._dynamo.mark_dynamic(input_chunk, 1)
             torch._dynamo.mark_dynamic(target_chunk, 1)
             torch._dynamo.mark_dynamic(target, 1)
+            torch._dynamo.mark_dynamic(ref_input_chunk, 1) if use_ref_model else None
 
             # accumulate loss, gradients, and metrics
-            accumulate_chunk(input_chunk, target_chunk)
+            accumulate_chunk(input_chunk, target_chunk, ref_input_chunk)
 
         # combine grad_chosen_inputs and grad_rejected_inputs
         grad_inputs = grad_chosen_inputs + grad_rejected_inputs
@@ -301,6 +333,7 @@ class LigerFusedLinearPreferenceBase(torch.autograd.Function):
         beta=0.1,
         compute_nll_loss=True,
         use_ref_model=False,
+        ref_input_chunk=None,
         ref_weight=None,
         ref_bias=None,
         **loss_kwargs,
@@ -357,7 +390,7 @@ class LigerFusedLinearPreferenceBase(torch.autograd.Function):
                     ref_rejected_logits,
                     ref_chosen_nll_loss,
                 ) = LigerFusedLinearPreferenceBase.chunk_forward(
-                    input_chunk,
+                    ref_input_chunk,
                     ref_weight,
                     target_chunk,
                     ref_bias,

{liger_kernel_nightly-0.5.2.dev20241212030605 → liger_kernel_nightly-0.5.2.dev20241212033924}/src/liger_kernel_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.5.2.dev20241212030605
+Version: 0.5.2.dev20241212033924
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 

{liger_kernel_nightly-0.5.2.dev20241212030605 → liger_kernel_nightly-0.5.2.dev20241212033924}/test/chunked_loss/test_dpo_loss.py

@@ -75,9 +75,15 @@ class TorchLMHeadDPO(torch.nn.Module):
             ignore_index=ignore_index, beta=beta, use_ref_model=True
         ).get_batch_loss_metrics
 
-    def forward(self, x, y):
+    def forward(self, x, ref_x, y):
         return self.dpo_loss(
-            self.lin.weight, x, y, self.lin.bias, self.ref_lin.weight, self.ref_lin.bias
+            self.lin.weight,
+            x,
+            y,
+            self.lin.bias,
+            ref_x,
+            self.ref_lin.weight,
+            self.ref_lin.bias,
         )
 
 
@@ -103,9 +109,15 @@ class LigerLMHeadDPO(torch.nn.Module):
             ignore_index=ignore_index, beta=beta, use_ref_model=True
         )
 
-    def forward(self, x, y):
+    def forward(self, x, ref_x, y):
         return self.dpo_loss(
-            self.lin.weight, x, y, self.lin.bias, self.ref_lin.weight, self.ref_lin.bias
+            self.lin.weight,
+            x,
+            y,
+            self.lin.bias,
+            ref_x,
+            self.ref_lin.weight,
+            self.ref_lin.bias,
         )
 
 
@@ -170,6 +182,10 @@ def test_correctness(
     input1 = _input.detach().clone().requires_grad_(True)
     input2 = _input.detach().clone().requires_grad_(True)
 
+    ref_input = (
+        torch.randn(B, T, H, device=device, dtype=dtype, requires_grad=False) * scalar
+    )
+
     target = torch.randint(
         0,
         V,
@@ -185,8 +201,8 @@ def test_correctness(
     indices_to_assign = torch.randperm(B * T)[:num_elements_to_assign]
     target.view(-1)[indices_to_assign] = ignore_index
 
-    loss1, aggregated_aux_outputs1 = torch_lm_head_dpo(input1, target)
-    loss2, aggregated_aux_outputs2 = liger_lm_head_dpo(input2, target)
+    loss1, aggregated_aux_outputs1 = torch_lm_head_dpo(input1, ref_input, target)
+    loss2, aggregated_aux_outputs2 = liger_lm_head_dpo(input2, ref_input, target)
 
     assert_verbose_allclose(loss1, loss2, atol=atol, rtol=rtol)
 
@@ -242,6 +258,10 @@ def test_correctness_functional(B, T, H, V, scalar, dtype, atol, rtol, bias, ref
     input1 = _input.detach().clone().requires_grad_(True)
     input2 = _input.detach().clone().requires_grad_(True)
 
+    ref_input = (
+        torch.randn(B, T, H, device=device, dtype=dtype, requires_grad=False) * scalar
+    )
+
     target = torch.randint(
         0,
         V,
@@ -270,10 +290,10 @@ def test_correctness_functional(B, T, H, V, scalar, dtype, atol, rtol, bias, ref
     ref_bias2 = _ref_bias.detach().clone().requires_grad_(True) if ref_bias else None
 
     loss1, aggregated_aux_outputs1 = LigerFusedLinearDPOFunction.apply(
-        input1, weight1, target, bias1, ref_weight1, ref_bias1
+        input1, weight1, target, bias1, ref_input, ref_weight1, ref_bias1
     )
     loss2, aggregated_aux_outputs2 = liger_fused_linear_dpo(
-        input2, weight2, target, bias2, ref_weight2, ref_bias2
+        input2, weight2, target, bias2, ref_input, ref_weight2, ref_bias2
     )
 
     assert_verbose_allclose(loss1, loss2, atol=atol, rtol=rtol)

{liger_kernel_nightly-0.5.2.dev20241212030605 → liger_kernel_nightly-0.5.2.dev20241212033924}/test/utils.py

@@ -478,6 +478,7 @@ class HFAlignmentLoss:
         _input: torch.FloatTensor,
         target: torch.LongTensor,
         bias: torch.FloatTensor = None,
+        ref_input: torch.FloatTensor = None,
         ref_weight: torch.FloatTensor = None,
         ref_bias: torch.FloatTensor = None,
         average_log_prob: bool = True,
@@ -498,7 +499,7 @@ class HFAlignmentLoss:
         loss_kwargs = {}
         if self.use_ref_model:
             ref_chosen_logps, ref_rejected_logps = self.get_ref_logps(
-                _input, ref_weight, target, ref_bias, average_log_prob
+                ref_input, ref_weight, target, ref_bias, average_log_prob
             )
             loss_kwargs["ref_chosen_logps"] = ref_chosen_logps
             loss_kwargs["ref_rejected_logps"] = ref_rejected_logps