From ba76dedf05643644770e403494f949ceba382849 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=9C=D0=B0=D0=BB=D0=B0=D1=85=D0=BE=D0=B2=20=D0=90=D0=BB?=
 =?UTF-8?q?=D0=B5=D0=BA=D1=81=D0=B5=D0=B9=20=D0=9F=D0=B0=D0=B2=D0=BB=D0=BE?=
 =?UTF-8?q?=D0=B2=D0=B8=D1=87?= <a.p.malakhov@tinkoff.ru>
Date: Sun, 29 Sep 2024 15:16:49 +0300
Subject: [PATCH] move liger kernels to modeling

---
 turbo_alignment/common/tf/loaders/model/model.py            | 4 +---
 turbo_alignment/modeling/liger_kernels/__init__.py          | 1 +
 .../{common/tf => modeling}/liger_kernels/cross_entropy.py  | 0
 .../{common/tf => modeling}/liger_kernels/geglu.py          | 2 +-
 .../tf => modeling}/liger_kernels/monkey_patch_liger.py     | 6 +++---
 .../{common/tf => modeling}/liger_kernels/rope.py           | 0
 .../{common/tf => modeling}/liger_kernels/utils.py          | 0
 7 files changed, 6 insertions(+), 7 deletions(-)
 create mode 100644 turbo_alignment/modeling/liger_kernels/__init__.py
 rename turbo_alignment/{common/tf => modeling}/liger_kernels/cross_entropy.py (100%)
 rename turbo_alignment/{common/tf => modeling}/liger_kernels/geglu.py (98%)
 rename turbo_alignment/{common/tf => modeling}/liger_kernels/monkey_patch_liger.py (82%)
 rename turbo_alignment/{common/tf => modeling}/liger_kernels/rope.py (100%)
 rename turbo_alignment/{common/tf => modeling}/liger_kernels/utils.py (100%)

diff --git a/turbo_alignment/common/tf/loaders/model/model.py b/turbo_alignment/common/tf/loaders/model/model.py
index 797d9ce..747799a 100755
--- a/turbo_alignment/common/tf/loaders/model/model.py
+++ b/turbo_alignment/common/tf/loaders/model/model.py
@@ -2,13 +2,11 @@
 from peft import PeftModel, get_peft_model, prepare_model_for_int8_training
 from transformers import PreTrainedModel, PreTrainedTokenizerBase
 
-from turbo_alignment.common.tf.liger_kernels.monkey_patch_liger import (
-    apply_liger_kernel_to_gemma2,
-)
 from turbo_alignment.common.tf.loaders.model.registry import (
     PeftConfigRegistry,
     TransformersAutoModelRegistry,
 )
+from turbo_alignment.modeling.liger_kernels import apply_liger_kernel_to_gemma2
 from turbo_alignment.settings.model import (
     ModelForPeftSettings,
     PreTrainedAdaptersModelSettings,
diff --git a/turbo_alignment/modeling/liger_kernels/__init__.py b/turbo_alignment/modeling/liger_kernels/__init__.py
new file mode 100644
index 0000000..baf9401
--- /dev/null
+++ b/turbo_alignment/modeling/liger_kernels/__init__.py
@@ -0,0 +1 @@
+from monkey_patch_liger import *
\ No newline at end of file
diff --git a/turbo_alignment/common/tf/liger_kernels/cross_entropy.py b/turbo_alignment/modeling/liger_kernels/cross_entropy.py
similarity index 100%
rename from turbo_alignment/common/tf/liger_kernels/cross_entropy.py
rename to turbo_alignment/modeling/liger_kernels/cross_entropy.py
diff --git a/turbo_alignment/common/tf/liger_kernels/geglu.py b/turbo_alignment/modeling/liger_kernels/geglu.py
similarity index 98%
rename from turbo_alignment/common/tf/liger_kernels/geglu.py
rename to turbo_alignment/modeling/liger_kernels/geglu.py
index 3039223..3bb85c8 100644
--- a/turbo_alignment/common/tf/liger_kernels/geglu.py
+++ b/turbo_alignment/modeling/liger_kernels/geglu.py
@@ -5,7 +5,7 @@
 import triton
 import triton.language as tl
 
-from turbo_alignment.common.tf.liger_kernels.utils import (
+from turbo_alignment.modeling.liger_kernels.utils import (
     calculate_settings,
     compare_version,
     ensure_contiguous,
diff --git a/turbo_alignment/common/tf/liger_kernels/monkey_patch_liger.py b/turbo_alignment/modeling/liger_kernels/monkey_patch_liger.py
similarity index 82%
rename from turbo_alignment/common/tf/liger_kernels/monkey_patch_liger.py
rename to turbo_alignment/modeling/liger_kernels/monkey_patch_liger.py
index ce915bb..bfce35d 100644
--- a/turbo_alignment/common/tf/liger_kernels/monkey_patch_liger.py
+++ b/turbo_alignment/modeling/liger_kernels/monkey_patch_liger.py
@@ -1,9 +1,9 @@
 from transformers import PretrainedConfig, PreTrainedModel
 
 from turbo_alignment.common.logging import get_project_logger
-from turbo_alignment.common.tf.liger_kernels.cross_entropy import LigerCrossEntropyLoss
-from turbo_alignment.common.tf.liger_kernels.geglu import LigerGEGLUMLP
-from turbo_alignment.common.tf.liger_kernels.rope import liger_rotary_pos_emb
+from turbo_alignment.modeling.liger_kernels.cross_entropy import LigerCrossEntropyLoss
+from turbo_alignment.modeling.liger_kernels.geglu import LigerGEGLUMLP
+from turbo_alignment.modeling.liger_kernels.rope import liger_rotary_pos_emb
 
 logger = get_project_logger()
 
diff --git a/turbo_alignment/common/tf/liger_kernels/rope.py b/turbo_alignment/modeling/liger_kernels/rope.py
similarity index 100%
rename from turbo_alignment/common/tf/liger_kernels/rope.py
rename to turbo_alignment/modeling/liger_kernels/rope.py
diff --git a/turbo_alignment/common/tf/liger_kernels/utils.py b/turbo_alignment/modeling/liger_kernels/utils.py
similarity index 100%
rename from turbo_alignment/common/tf/liger_kernels/utils.py
rename to turbo_alignment/modeling/liger_kernels/utils.py