Refactor

Signed-off-by: Austin Liu <[email protected]>
linkedin · Nov 14, 2024 · b88708d · b88708d
1 parent 0e71106
commit b88708d
Show file tree

Hide file tree

Showing 5 changed files with 343 additions and 376 deletions.
diff --git a/benchmark/scripts/benchmark_dpo_loss.py b/benchmark/scripts/benchmark_dpo_loss.py
@@ -1,5 +1,6 @@
+from test.chunked_loss.test_dpo_loss import HF_DPO_Loss
+
 import torch
-import triton
 from utils import (
     QUANTILES,
     SingleBenchmarkRunInput,
@@ -9,7 +10,8 @@
     run_benchmarks,
 )
 
-from liger_kernel.alignment.dpo_loss import HF_DPO_Loss, LigerFusedLinearDPOFunction
+import triton
+from liger_kernel.chunked_loss.dpo_loss import LigerFusedLinearDPOFunction
 
 
 class TorchDPOLoss(torch.nn.Module):

diff --git a/src/liger_kernel/alignment/dpo_loss.py b/src/liger_kernel/alignment/dpo_loss.py