liger-kernel-nightly 0.5.3.dev20250221162633__py3-none-any.whl → 0.5.3.dev20250221233257__py3-none-any.whl

liger_kernel/ops/tvd.py

@@ -1,4 +1,5 @@
-from typing import Literal, Optional
+from typing import Literal
+from typing import Optional
 
 import torch
 import triton
@@ -178,15 +179,13 @@ class LigerTVDLossFunction(torch.autograd.Function):
         """
         has_label = False
         if shift_labels is not None:
-            assert shift_labels.shape == (
-                p.shape[0],
-            ), f"the shape of shift_labels must be (BT,). Got: {shift_labels.shape}"
+            assert shift_labels.shape == (p.shape[0],), (
+                f"the shape of shift_labels must be (BT,). Got: {shift_labels.shape}"
+            )
             shift_labels = shift_labels.contiguous()
             has_label = True
 
-        loss, grads = tv_distance_forward_triton(
-            p, q, shift_labels, reduction, ignore_index, has_label
-        )
+        loss, grads = tv_distance_forward_triton(p, q, shift_labels, reduction, ignore_index, has_label)
         ctx.save_for_backward(grads)
         return loss
 

liger_kernel/transformers/__init__.py

@@ -19,7 +19,7 @@ from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen2
 from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen2_vl  # 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.tvd import LigerTVDLoss  # noqa: F401
 from liger_kernel.transformers.swiglu import LigerBlockSparseTop2MLP  # noqa: F401
 from liger_kernel.transformers.swiglu import LigerPhi3SwiGLUMLP  # noqa: F401
 from liger_kernel.transformers.swiglu import LigerSwiGLUMLP  # noqa: F401
+from liger_kernel.transformers.tvd import LigerTVDLoss  # noqa: F401

liger_kernel/transformers/functional.py

@@ -14,6 +14,7 @@ from liger_kernel.ops.rope import LigerRopeFunction
 from liger_kernel.ops.swiglu import LigerSiLUMulFunction
 from liger_kernel.ops.tvd import LigerTVDLossFunction
 
+
 # conform to the function signature in https://pytorch.org/docs/stable/generated/torch.nn.functional.cross_entropy.html
 # `weight` and `size_average` are placeholders and not implemented yet
 def liger_cross_entropy(
@@ -156,6 +157,7 @@ def liger_kl_div(
         eps,
     )
 
+
 def liger_tvd(
     input,
     target,
@@ -169,7 +171,8 @@ def liger_tvd(
         shift_labels,
         reduction,
         ignore_index,
-    ) 
+    )
+
 
 def liger_layer_norm(X, W, B, eps):
     return LigerLayerNormFunction.apply(X, W, B, eps)

liger_kernel/transformers/tvd.py

@@ -10,6 +10,4 @@ class LigerTVDLoss(nn.Module):
         self.ignore_index = ignore_index
 
     def forward(self, p, q, shift_labels=None):
-        return LigerTVDLossFunction.apply(
-            p, q, shift_labels, self.reduction, self.ignore_index
-        )
+        return LigerTVDLossFunction.apply(p, q, shift_labels, self.reduction, self.ignore_index)

liger_kernel/utils.py

@@ -9,5 +9,54 @@ def infer_device():
         return "cuda"
     elif torch.xpu.is_available():
         return "xpu"
+    elif torch.hip.is_available():
+        return "hip"
     else:
         return "cpu"
+
+
+def transformers_version_dispatch(
+    required_version: str,
+    before_fn,
+    after_fn,
+    before_args: tuple = (),
+    after_args: tuple = (),
+    before_kwargs: dict = None,
+    after_kwargs: dict = None,
+):
+    """
+    Dispatches to different functions based on package version comparison.
+
+    Args:
+        required_version: Version to compare against (e.g. "4.48.0")
+        before_fn: Function to call if package_version < required_version
+        after_fn: Function to call if package_version >= required_version
+        before_args: Positional arguments for before_fn
+        after_args: Positional arguments for after_fn
+        before_kwargs: Keyword arguments for before_fn
+        after_kwargs: Keyword arguments for after_fn
+
+    Returns:
+        Result from either before_fn or after_fn
+
+    Example:
+        >>> rotary_emb = transformers_version_dispatch(
+        ...     "4.48.0",
+        ...     LlamaRotaryEmbedding,
+        ...     LlamaRotaryEmbedding,
+        ...     before_args=(head_dim,),
+        ...     after_args=(LlamaConfig(head_dim=head_dim),),
+        ...     before_kwargs={'device': device},
+        ...     after_kwargs={'device': device}
+        ... )
+    """
+    from packaging import version
+    from transformers import __version__ as transformers_version
+
+    before_kwargs = before_kwargs or {}
+    after_kwargs = after_kwargs or {}
+
+    if version.parse(transformers_version) < version.parse(required_version):
+        return before_fn(*before_args, **before_kwargs)
+    else:
+        return after_fn(*after_args, **after_kwargs)

{liger_kernel_nightly-0.5.3.dev20250221162633.dist-info → liger_kernel_nightly-0.5.3.dev20250221233257.dist-info}/METADATA

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

{liger_kernel_nightly-0.5.3.dev20250221162633.dist-info → liger_kernel_nightly-0.5.3.dev20250221233257.dist-info}/RECORD

@@ -1,6 +1,6 @@
 liger_kernel/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 liger_kernel/env_report.py,sha256=uhdEC8OydxoZlb7B6YYcAaBF3crGFdIck-4cxaW4NJY,1728
-liger_kernel/utils.py,sha256=HJa-xVKOohDn6pLVIx-Fv0V9h0QAL3qZGQNRICI-OpI,249
+liger_kernel/utils.py,sha256=FtVUkCGBT1UNasTl6HMNycWwiwHayK6tx-ZDdA-sNX4,1884
 liger_kernel/chunked_loss/README.md,sha256=0FmkFC3hKBqyoDT5uTlIYmrvRkF-EOCR1y-EBU1LpWU,2248
 liger_kernel/chunked_loss/__init__.py,sha256=ATu-xX5Fc49Cr6yBOGBRNTo593ZrU5ZCsIuvoIbJWw4,603
 liger_kernel/chunked_loss/cpo_loss.py,sha256=OdBR8WYdHTKpLI_c9DcuwqKSWPeAAeTyREz46Vu_cAY,3682
@@ -28,14 +28,14 @@ liger_kernel/ops/qwen2vl_mrope.py,sha256=3GExhYpLgB4VUtyZyjRk8XjEur3W4EWF6HQ67ML
 liger_kernel/ops/rms_norm.py,sha256=PWLJcdIKU5e-8BuYFHd9Cqlq6wmr6fUXKi9zQD4LetU,11727
 liger_kernel/ops/rope.py,sha256=ofmBOkUpZZO-Q8Z5B_LOFYYLD-YT-8WnJ4vGOrDYouI,8943
 liger_kernel/ops/swiglu.py,sha256=KmgMjaJQnbLLgZn2nEpbwHU_xpnYRweCyrLQSVvM1vA,3015
-liger_kernel/ops/tvd.py,sha256=9wVCijj2vBtgiLeUHhl7hy_LAiJ3liPIYOGMSU3P1ro,6407
+liger_kernel/ops/tvd.py,sha256=FHJtLQI95ijqgg9UtaHpMAjSCiPxB6CduPwPMcGxelc,6405
 liger_kernel/ops/utils.py,sha256=uoFKQqo-34N2TWQNvXMFywqGiOMMXNEVBxVojzlUAa0,3836
 liger_kernel/ops/experimental/embedding.py,sha256=tolj3tItkzpSb30zWqDN2_yX4ectflaQ8HMyKyFIQc8,4172
 liger_kernel/ops/experimental/mm_int8int2.py,sha256=TrS9lpwekrik_w5qE7AhMJD1bcq-OidjtbsW80oZ6IM,13314
-liger_kernel/transformers/__init__.py,sha256=i6GPkP5-esFBh205nF4MluNrL7KNugseGiUKdSHGW70,2172
+liger_kernel/transformers/__init__.py,sha256=MGgdJkohu0tQS6owEBHuRVYhRUPXRFP9OiVc1fcjkjc,2172
 liger_kernel/transformers/auto_model.py,sha256=0qCTRZt280Bj_LcFdzo9hlaR-BWNazawXOGgoCZjgEg,1545
 liger_kernel/transformers/cross_entropy.py,sha256=z3KTWQnFxr_IZaVjtYt0ZNEWQdDdYThN35xWkHlDGH0,1683
-liger_kernel/transformers/functional.py,sha256=zahXVCjA2NxcVFpAgajILIRN0GO6mrbfLPgONUkTrY8,4940
+liger_kernel/transformers/functional.py,sha256=ShLD3eb--XKNtllznCrOYTbo4f-1KVwzi0KLMICdrn4,4942
 liger_kernel/transformers/fused_linear_cross_entropy.py,sha256=09Rt7FZzLH42VOcIbQ4dlQd0o3Rlb4vk6fqiOQ7WTD8,1778
 liger_kernel/transformers/fused_linear_jsd.py,sha256=bZ4otCvWBuOnA5XdQL-FzZVItJlDt-ht9e_pG7PG93E,3999
 liger_kernel/transformers/geglu.py,sha256=mrgqzIUVd6lN7fkDKLkw5YaESDxDtFgbot430WwPVOQ,1107
@@ -49,7 +49,7 @@ liger_kernel/transformers/rms_norm.py,sha256=GqCEJuGt0YdqqlMcToE0Wp4A8YFquDa4UUS
 liger_kernel/transformers/rope.py,sha256=ZTrTORSAyfcFIKjk6XEeYmk4ROH7xXED9L4g2NFntlE,999
 liger_kernel/transformers/swiglu.py,sha256=i9WTqcNRqReU4XJs391IPbl-I5X0wG4T72D4pqGFfJg,2422
 liger_kernel/transformers/trainer_integration.py,sha256=W3ON51O5GkyzNJsItz0y5rKx-uy2f2cFfveZpqbUdhw,123
-liger_kernel/transformers/tvd.py,sha256=lt730XDR4IYEkn-HeWS7WU6AGssg90ubg8pqPAX2lbE,472
+liger_kernel/transformers/tvd.py,sha256=XrRfyJIqN6HFxXk8MYyFVZM1OLz3mtSbRZvWfZ_JerQ,450
 liger_kernel/transformers/experimental/embedding.py,sha256=2P0QYdlFyFrG5OqTzTa1wcRgDSyjBMv5i1a7BrDPDQw,881
 liger_kernel/transformers/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 liger_kernel/transformers/model/gemma.py,sha256=ky89b3aWPaeTGRMC-745KgixtQIRXzNAiCORAMLn9yo,9654
@@ -65,9 +65,9 @@ liger_kernel/transformers/trainer/__init__.py,sha256=p7yQfklV8-467qSz_ZMimkbDF7H
 liger_kernel/transformers/trainer/orpo_trainer.py,sha256=pdekW7l6Qg_aqa5SYKYlSWUF8m3lkOFvFLcIMEHrz9s,8338
 liger_kernel/triton/__init__.py,sha256=qCiCamzCRv6lpV8IqpAc9YMdNKC7GKurClWceQPnlis,92
 liger_kernel/triton/monkey_patch.py,sha256=Rd0hUHAzDkFfHvnX7-PBaNK5EKnZhtfM_h-fgQH9HPY,1568
-liger_kernel_nightly-0.5.3.dev20250221162633.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
-liger_kernel_nightly-0.5.3.dev20250221162633.dist-info/METADATA,sha256=9kT0AmFeMP_D-n-dGhoOSBX8ThMPM8ypnFT4f3O1FGc,22093
-liger_kernel_nightly-0.5.3.dev20250221162633.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
-liger_kernel_nightly-0.5.3.dev20250221162633.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-liger_kernel_nightly-0.5.3.dev20250221162633.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
-liger_kernel_nightly-0.5.3.dev20250221162633.dist-info/RECORD,,
+liger_kernel_nightly-0.5.3.dev20250221233257.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
+liger_kernel_nightly-0.5.3.dev20250221233257.dist-info/METADATA,sha256=ZsBEdMtozVk3GNw6IRdb-wu7XjzLVefruhKVD7JJjdE,22093
+liger_kernel_nightly-0.5.3.dev20250221233257.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
+liger_kernel_nightly-0.5.3.dev20250221233257.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+liger_kernel_nightly-0.5.3.dev20250221233257.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
+liger_kernel_nightly-0.5.3.dev20250221233257.dist-info/RECORD,,