liger-kernel 0.5.0__tar.gz → 0.5.2__tar.gz

{liger_kernel-0.5.0/src/liger_kernel.egg-info → liger_kernel-0.5.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel
-Version: 0.5.0
+Version: 0.5.2
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 
@@ -34,9 +34,10 @@ Requires-Dist: torch>=2.1.2
 Requires-Dist: triton>=2.3.1
 Provides-Extra: transformers
 Requires-Dist: transformers~=4.0; extra == "transformers"
+Provides-Extra: trl
+Requires-Dist: trl>=0.11.0; extra == "trl"
 Provides-Extra: dev
 Requires-Dist: transformers>=4.44.2; extra == "dev"
-Requires-Dist: trl>=0.11.0; extra == "dev"
 Requires-Dist: matplotlib>=3.7.2; extra == "dev"
 Requires-Dist: flake8>=4.0.1.1; extra == "dev"
 Requires-Dist: black>=24.4.2; extra == "dev"

{liger_kernel-0.5.0 → liger_kernel-0.5.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel"
-version = "0.5.0"
+version = "0.5.2"
 description = "Efficient Triton kernels for LLM Training"
 urls = { "Homepage" = "https://github.com/linkedin/Liger-Kernel" }
 readme = { file = "README.md", content-type = "text/markdown" }
@@ -20,9 +20,12 @@ transformers = [
     "transformers~=4.0"
 ]
 
+trl = [
+    "trl>=0.11.0",
+]
+
 dev = [
     "transformers>=4.44.2",
-    "trl>=0.11.0",
     "matplotlib>=3.7.2",
     "flake8>=4.0.1.1",
     "black>=24.4.2",

{liger_kernel-0.5.0 → liger_kernel-0.5.2}/src/liger_kernel/ops/qwen2vl_mrope.py

@@ -10,6 +10,7 @@ def _triton_qwen2vl_mrope(
     cos,
     sin,
     sl,
+    bs: tl.constexpr,
     n_qh: tl.constexpr,
     n_kh: tl.constexpr,
     hd: tl.constexpr,
@@ -41,13 +42,12 @@ def _triton_qwen2vl_mrope(
     t_end = mrope_section_t
     h_end = t_end + mrope_section_h
 
-    cos_row_idx = pid % sl
-    t_cos = cos + cos_row_idx * hd
-    h_cos = t_cos + sl * hd
-    w_cos = h_cos + sl * hd
-    t_sin = sin + cos_row_idx * hd
-    h_sin = t_sin + sl * hd
-    w_sin = h_sin + sl * hd
+    t_cos = cos + pid * hd
+    h_cos = t_cos + bs * sl * hd
+    w_cos = h_cos + bs * sl * hd
+    t_sin = sin + pid * hd
+    h_sin = t_sin + bs * sl * hd
+    w_sin = h_sin + bs * sl * hd
 
     cos_offsets = tl.arange(0, pad_hd // 2)
     t_mask = cos_offsets < t_end
@@ -151,6 +151,7 @@ def qwen2vl_mrope_forward(q, k, cos, sin, mrope_section):
         cos,
         sin,
         seq_len,
+        batch_size,
         n_q_head,
         n_kv_head,
         head_dim,
@@ -189,6 +190,7 @@ def qwen2vl_mrope_backward(dq, dk, cos, sin, mrope_section):
         cos,
         sin,
         seq_len,
+        batch_size,
         n_q_head,
         n_kv_head,
         head_dim,
@@ -216,8 +218,8 @@ class LigerQwen2VLMRopeFunction(torch.autograd.Function):
         """
         q size: (bsz, n_q_head, seq_len, head_dim)
         k size: (bsz, n_kv_head, seq_len, head_dim)
-        cos size: (3, 1, seq_len, head_dim)
-        sin size: (3, 1, seq_len, head_dim)
+        cos size: (3, bsz, seq_len, head_dim)
+        sin size: (3, bsz, seq_len, head_dim)
         """
         q, k, cos, sin = qwen2vl_mrope_forward(q, k, cos, sin, mrope_section)
         ctx.save_for_backward(cos, sin)
@@ -228,10 +230,9 @@ class LigerQwen2VLMRopeFunction(torch.autograd.Function):
         """
         dq size: (bsz, n_q_head, seq_len, head_dim)
         dk size: (bsz, n_kv_head, seq_len, head_dim)
-        cos size: (3, 1, seq_len, head_dim)
-        sin size: (3, 1, seq_len, head_dim)
+        cos size: (3, bsz, seq_len, head_dim)
+        sin size: (3, bsz, seq_len, head_dim)
         """
-
         cos, sin = ctx.saved_tensors
         mrope_section = ctx.mrope_section
         dq, dk = qwen2vl_mrope_backward(dq, dk, cos, sin, mrope_section)

{liger_kernel-0.5.0 → liger_kernel-0.5.2}/src/liger_kernel/transformers/__init__.py

@@ -22,7 +22,6 @@ from liger_kernel.transformers.monkey_patch import (  # noqa: F401
     apply_liger_kernel_to_qwen2,
     apply_liger_kernel_to_qwen2_vl,
 )
-from liger_kernel.transformers.orpo_trainer import LigerORPOTrainer  # noqa: F401
 from liger_kernel.transformers.rms_norm import LigerRMSNorm  # noqa: F401
 from liger_kernel.transformers.rope import liger_rotary_pos_emb  # noqa: F401
 from liger_kernel.transformers.swiglu import (  # noqa: F401

{liger_kernel-0.5.0 → liger_kernel-0.5.2}/src/liger_kernel/transformers/qwen2vl_mrope.py

@@ -8,8 +8,8 @@ def liger_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim
     Args:
         q (torch.Tensor): The query tensor of shape (bsz, n_q_head, seq_len, head_dim).
         k (torch.Tensor): The key tensor of shape (bsz, n_kv_head, seq_len, head_dim).
-        cos (torch.Tensor): The cosine tensor of shape (3, 1, seq_len, head_dim).
-        sin (torch.Tensor): The sine tensor of shape (3, 1, seq_len, head_dim).
+        cos (torch.Tensor): The cosine tensor of shape (3, bsz, seq_len, head_dim).
+        sin (torch.Tensor): The sine tensor of shape (3, bsz, seq_len, head_dim).
         mrope_section (List[int]): The multimodal rope section for channel dimension of temporal, height and width in rope calculation.
         unsqueeze_dim (int, optional): The dimension to unsqueeze. Defaults to 1.
 

liger_kernel-0.5.2/src/liger_kernel/transformers/trainer/__init__.py

@@ -0,0 +1,6 @@
+try:
+    from liger_kernel.transformers.trainer.orpo_trainer import (  # noqa: F401
+        LigerORPOTrainer,
+    )
+except ImportError:
+    raise ImportError("Please `pip install trl` to use LigerORPOTrainer")

{liger_kernel-0.5.0/src/liger_kernel/transformers → liger_kernel-0.5.2/src/liger_kernel/transformers/trainer}/orpo_trainer.py

@@ -76,9 +76,7 @@ class LigerORPOTrainer(ORPOTrainer):
             padding_value=self.padding_value,
             device=self.accelerator.device,
         )
-        # if self.accelerator.is_main_process:
-        #     import pdb; pdb.set_trace()
-        # torch.distributed.barrier()
+
         model_kwargs = (
             {
                 "decoder_input_ids": self._shift_right(

{liger_kernel-0.5.0 → liger_kernel-0.5.2/src/liger_kernel.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel
-Version: 0.5.0
+Version: 0.5.2
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 
@@ -34,9 +34,10 @@ Requires-Dist: torch>=2.1.2
 Requires-Dist: triton>=2.3.1
 Provides-Extra: transformers
 Requires-Dist: transformers~=4.0; extra == "transformers"
+Provides-Extra: trl
+Requires-Dist: trl>=0.11.0; extra == "trl"
 Provides-Extra: dev
 Requires-Dist: transformers>=4.44.2; extra == "dev"
-Requires-Dist: trl>=0.11.0; extra == "dev"
 Requires-Dist: matplotlib>=3.7.2; extra == "dev"
 Requires-Dist: flake8>=4.0.1.1; extra == "dev"
 Requires-Dist: black>=24.4.2; extra == "dev"

{liger_kernel-0.5.0 → liger_kernel-0.5.2}/src/liger_kernel.egg-info/SOURCES.txt

@@ -46,7 +46,6 @@ src/liger_kernel/transformers/jsd.py
 src/liger_kernel/transformers/kl_div.py
 src/liger_kernel/transformers/layer_norm.py
 src/liger_kernel/transformers/monkey_patch.py
-src/liger_kernel/transformers/orpo_trainer.py
 src/liger_kernel/transformers/qwen2vl_mrope.py
 src/liger_kernel/transformers/rms_norm.py
 src/liger_kernel/transformers/rope.py
@@ -63,5 +62,7 @@ src/liger_kernel/transformers/model/mllama.py
 src/liger_kernel/transformers/model/phi3.py
 src/liger_kernel/transformers/model/qwen2.py
 src/liger_kernel/transformers/model/qwen2_vl.py
+src/liger_kernel/transformers/trainer/__init__.py
+src/liger_kernel/transformers/trainer/orpo_trainer.py
 src/liger_kernel/triton/__init__.py
 src/liger_kernel/triton/monkey_patch.py

{liger_kernel-0.5.0 → liger_kernel-0.5.2}/src/liger_kernel.egg-info/requires.txt

@@ -9,7 +9,6 @@ triton>=3.0.0
 
 [dev]
 transformers>=4.44.2
-trl>=0.11.0
 matplotlib>=3.7.2
 flake8>=4.0.1.1
 black>=24.4.2
@@ -23,3 +22,6 @@ seaborn
 
 [transformers]
 transformers~=4.0
+
+[trl]
+trl>=0.11.0