liger-kernel-nightly 0.3.1.dev20241025165359__tar.gz → 0.3.1.dev20241030045958__tar.gz

{liger_kernel_nightly-0.3.1.dev20241025165359 → liger_kernel_nightly-0.3.1.dev20241030045958}/PKG-INFO

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

{liger_kernel_nightly-0.3.1.dev20241025165359 → liger_kernel_nightly-0.3.1.dev20241030045958}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.3.1.dev20241025165359"
+version = "0.3.1.dev20241030045958"
 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.3.1.dev20241025165359 → liger_kernel_nightly-0.3.1.dev20241030045958}/src/liger_kernel/ops/fused_linear_cross_entropy.py

@@ -2,7 +2,7 @@ import torch
 import triton
 
 from liger_kernel.ops.cross_entropy import liger_cross_entropy_kernel
-from liger_kernel.ops.utils import element_mul_kernel
+from liger_kernel.ops.utils import amp_custom_bwd, amp_custom_fwd, element_mul_kernel
 
 # The hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576 https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19
 # However, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling
@@ -19,9 +19,7 @@ def fused_linear_cross_entropy_forward(
     label_smoothing=0.0,
     reduction="mean",
 ):
-    dtype = (
-        torch.get_autocast_gpu_dtype() if torch.is_autocast_enabled() else _input.dtype
-    )
+    dtype = _input.dtype
     device = _input.device
 
     # inputs have shape: BT x H
@@ -189,6 +187,7 @@ def fused_linear_cross_entropy_backward(
 
 class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
     @staticmethod
+    @amp_custom_fwd
     def forward(
         ctx,
         _input,
@@ -228,6 +227,7 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
         return loss
 
     @staticmethod
+    @amp_custom_bwd
     def backward(ctx, grad_output):
         (grad_input, grad_weight, grad_bias) = ctx.saved_tensors
         grad_input, grad_weight, grad_bias = fused_linear_cross_entropy_backward(

{liger_kernel_nightly-0.3.1.dev20241025165359 → liger_kernel_nightly-0.3.1.dev20241030045958}/src/liger_kernel/ops/fused_linear_jsd.py

@@ -92,7 +92,9 @@ def fused_linear_jsd_forward(
             dX_ptr=student_prob_chunk,
             dX_stride=student_prob_chunk.stride(-2),
             label_ptr=(
-                shift_labels if has_label else torch.empty(1, device=device)
+                shift_labels[start_idx:end_idx]
+                if has_label
+                else torch.empty(1, device=device)
             ),  # dummy ptr if no label
             beta=jsd_beta,
             n_non_ignore=n_non_ignore,

{liger_kernel_nightly-0.3.1.dev20241025165359 → liger_kernel_nightly-0.3.1.dev20241030045958}/src/liger_kernel/ops/jsd.py

@@ -19,7 +19,7 @@ def _jsd_kernel(
     dX_stride,
     label_ptr,
     beta,
-    n_non_ignore,
+    n_non_ignore: int,
     ignore_index: tl.constexpr,
     n_cols,
     BLOCK_SIZE: tl.constexpr,

{liger_kernel_nightly-0.3.1.dev20241025165359 → liger_kernel_nightly-0.3.1.dev20241030045958}/src/liger_kernel/ops/utils.py

@@ -12,6 +12,7 @@ Modifications made by Yanning Chen, 2024.
 
 import functools
 import importlib
+import operator
 from typing import Callable
 
 import torch
@@ -63,6 +64,18 @@ def compare_version(package: str, operator: Callable, target: str):
     return operator(pkg_version, Version(target))
 
 
+def get_amp_custom_fwd_bwd() -> Callable:
+    if compare_version("torch", operator.ge, "2.4.0"):
+        return (
+            functools.partial(torch.amp.custom_fwd, device_type="cuda"),
+            functools.partial(torch.amp.custom_bwd, device_type="cuda"),
+        )
+    return torch.cuda.amp.custom_fwd, torch.cuda.amp.custom_bwd
+
+
+amp_custom_fwd, amp_custom_bwd = get_amp_custom_fwd_bwd()
+
+
 torch_to_triton_dtype = {
     torch.float32: tl.float32,
     torch.float16: tl.float16,

{liger_kernel_nightly-0.3.1.dev20241025165359 → liger_kernel_nightly-0.3.1.dev20241030045958}/src/liger_kernel_nightly.egg-info/PKG-INFO

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