liger-kernel 0.6.2__py3-none-any.whl → 0.6.3__py3-none-any.whl

liger_kernel/ops/poly_norm.py

@@ -0,0 +1,386 @@
+import operator
+
+import torch
+import triton
+import triton.language as tl
+
+from liger_kernel.ops.utils import calculate_settings
+from liger_kernel.ops.utils import compare_version
+from liger_kernel.ops.utils import ensure_contiguous
+
+if compare_version("triton", operator.ge, "3.0.0"):
+    try:
+        from triton.language.extra.libdevice import rsqrt
+    except ModuleNotFoundError:
+        from triton.language.extra.cuda.libdevice import rsqrt
+else:
+    from triton.language.math import rsqrt
+
+
+@triton.jit
+def _poly_norm_forward_kernel(
+    Y_ptr,
+    Y_row_stride,
+    X_ptr,
+    X_row_stride,
+    W_ptr,  # weight: [3] for [w0, w1, w2]
+    B_ptr,  # bias: scalar
+    RSTD_ptr,  # cache rstd for backward: shape (n_rows, 3)
+    RSTD_row_stride,
+    n_cols,
+    eps,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    PolyNorm formula:
+        y = w₀·norm(x³) + w₁·norm(x²) + w₂·norm(x) + b
+        where norm(u) = u / sqrt(mean(u²) + ε)
+
+    Reference:
+    1. https://github.com/BryceZhuo/PolyCom/
+    2. https://arxiv.org/pdf/2411.03884
+
+    Cache rstd values for backward pass
+    """
+    row_idx = tl.program_id(0).to(tl.int64)
+    col_offsets = tl.arange(0, BLOCK_SIZE)
+    mask = col_offsets < n_cols
+
+    # Load pointers
+    Y_ptr += row_idx * Y_row_stride
+    X_ptr += row_idx * X_row_stride
+    RSTD_ptr += row_idx * RSTD_row_stride
+
+    # Load input row
+    X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0.0)
+
+    # Load weights and bias
+    w0 = tl.load(W_ptr + 0)
+    w1 = tl.load(W_ptr + 1)
+    w2 = tl.load(W_ptr + 2)
+    b = tl.load(B_ptr)
+
+    # Compute x³, x², x
+    X_pow3 = X_row * X_row * X_row
+    X_pow2 = X_row * X_row
+    X_pow1 = X_row
+
+    # Compute norm(x³): norm(u) = u * rsqrt(mean(u²) + eps)
+    mean_square_3 = tl.sum(X_pow3 * X_pow3, axis=0) / n_cols
+    rstd_3 = rsqrt(mean_square_3 + eps)
+    norm_x3 = X_pow3 * rstd_3
+
+    # Compute norm(x²)
+    mean_square_2 = tl.sum(X_pow2 * X_pow2, axis=0) / n_cols
+    rstd_2 = rsqrt(mean_square_2 + eps)
+    norm_x2 = X_pow2 * rstd_2
+
+    # Compute norm(x)
+    mean_square_1 = tl.sum(X_pow1 * X_pow1, axis=0) / n_cols
+    rstd_1 = rsqrt(mean_square_1 + eps)
+    norm_x1 = X_pow1 * rstd_1
+
+    # Cache rstd values for backward
+    tl.store(RSTD_ptr + 0, rstd_3)
+    tl.store(RSTD_ptr + 1, rstd_2)
+    tl.store(RSTD_ptr + 2, rstd_1)
+
+    # Compute output: y = w₀·norm(x³) + w₁·norm(x²) + w₂·norm(x) + b
+    Y_row = w0 * norm_x3 + w1 * norm_x2 + w2 * norm_x1 + b
+
+    # Store output
+    tl.store(Y_ptr + col_offsets, Y_row, mask=mask)
+
+
+@triton.jit
+def _poly_norm_backward_kernel(
+    dY_ptr,
+    dY_row_stride,
+    dX_ptr,
+    dX_row_stride,
+    X_ptr,
+    X_row_stride,
+    W_ptr,
+    RSTD_ptr,
+    RSTD_row_stride,
+    dW_ptr,  # shape: (n_programs, 3)
+    dW_row_stride,
+    dB_ptr,  # shape: (n_programs,)
+    n_rows,
+    n_cols,
+    rows_per_program: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    PolyNorm Backward Kernel Gradient:
+        ∂L/∂x_i = Σ_p w_p * [p*x_i^(p-1) * grad_i/D_p - (p/d)*x_i^(2p-1) * S_p/(D_p³)]
+
+    where:
+        - D_p = RMS(x^p) = 1/rstd_p
+        - S_p = sum(grad * x^p) over the row
+        - d = n_cols
+        - p ∈ {3, 2, 1}
+    """
+    row_block_id = tl.program_id(0).to(tl.int64)
+    row_start = row_block_id * rows_per_program
+    row_end = min((row_block_id + 1) * rows_per_program, n_rows)
+    col_offsets = tl.arange(0, BLOCK_SIZE)
+    mask = col_offsets < n_cols
+
+    # Initialize accumulators for weight and bias gradients (scalars)
+    dW0_acc = 0.0
+    dW1_acc = 0.0
+    dW2_acc = 0.0
+    dB_acc = 0.0
+
+    # Load weights
+    w0 = tl.load(W_ptr + 0).to(tl.float32)
+    w1 = tl.load(W_ptr + 1).to(tl.float32)
+    w2 = tl.load(W_ptr + 2).to(tl.float32)
+
+    dY_ptr += row_start * dY_row_stride
+    dX_ptr += row_start * dX_row_stride
+    X_ptr += row_start * X_row_stride
+    RSTD_ptr += row_start * RSTD_row_stride
+
+    for _ in range(row_start, row_end):
+        # Load input and gradient
+        dY_row = tl.load(dY_ptr + col_offsets, mask=mask, other=0.0).to(tl.float32)
+        X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0.0).to(tl.float32)
+
+        # Load cached rstd values
+        rstd_3 = tl.load(RSTD_ptr + 0).to(tl.float32)
+        rstd_2 = tl.load(RSTD_ptr + 1).to(tl.float32)
+        rstd_1 = tl.load(RSTD_ptr + 2).to(tl.float32)
+
+        # Compute powers
+        X_pow3 = X_row * X_row * X_row
+        X_pow2 = X_row * X_row
+        X_pow1 = X_row
+
+        # Accumulate bias gradient: dB = sum(dY)
+        dB_acc += tl.sum(dY_row, axis=0)
+
+        # Compute gradient w.r.t. input using closed-form formula
+        # For p=3: ∂L/∂x from w0 * norm(x³)
+        S_3 = tl.sum(dY_row * X_pow3, axis=0)  # scalar
+        grad_x_3 = w0 * (
+            3.0 * X_pow2 * rstd_3 * dY_row
+            - (3.0 / n_cols) * X_row * X_row * X_row * X_row * X_row * (rstd_3 * rstd_3 * rstd_3) * S_3
+        )
+
+        # For p=2: ∂L/∂x from w1 * norm(x²)
+        S_2 = tl.sum(dY_row * X_pow2, axis=0)  # scalar
+        grad_x_2 = w1 * (
+            2.0 * X_row * rstd_2 * dY_row - (2.0 / n_cols) * X_row * X_row * X_row * (rstd_2 * rstd_2 * rstd_2) * S_2
+        )
+
+        # For p=1: ∂L/∂x from w2 * norm(x)
+        S_1 = tl.sum(dY_row * X_pow1, axis=0)  # scalar
+        grad_x_1 = w2 * (1.0 * rstd_1 * dY_row - (1.0 / n_cols) * X_row * (rstd_1 * rstd_1 * rstd_1) * S_1)
+
+        # Accumulate weight gradients using closed-form: dW_p = rstd_p * S_p
+        dW0_acc += rstd_3 * S_3
+        dW1_acc += rstd_2 * S_2
+        dW2_acc += rstd_1 * S_1
+
+        # Total gradient
+        dX_row = grad_x_3 + grad_x_2 + grad_x_1
+
+        # Store gradient
+        tl.store(dX_ptr + col_offsets, dX_row, mask=mask)
+
+        # Update pointers
+        dY_ptr += dY_row_stride
+        dX_ptr += dX_row_stride
+        X_ptr += X_row_stride
+        RSTD_ptr += RSTD_row_stride
+
+    # Store accumulated gradients (scalars)
+    tl.store(dW_ptr + row_block_id * dW_row_stride + 0, dW0_acc)
+    tl.store(dW_ptr + row_block_id * dW_row_stride + 1, dW1_acc)
+    tl.store(dW_ptr + row_block_id * dW_row_stride + 2, dW2_acc)
+    tl.store(dB_ptr + row_block_id, dB_acc)
+
+
+def poly_norm_forward(X, W, B, eps=1e-6):
+    """
+    PolyNorm Forward Pass
+
+    Args:
+        X: input tensor of shape (*, H) where H is hidden dimension
+        W: weight tensor of shape (3,) for [w0, w1, w2]
+        B: bias scalar tensor
+        eps: epsilon for numerical stability
+
+    Returns:
+        Y: output tensor of same shape as X
+        X: reshaped input (for backward)
+        RSTD: cached rstd values (for backward)
+        BLOCK_SIZE: block size used
+        num_warps: number of warps used
+    """
+    shape = X.shape
+    dim = shape[-1]
+    X = X.view(-1, dim)
+    n_rows, n_cols = X.shape
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+
+    # RSTD is to cache rstd for each row
+    Y = torch.empty((n_rows, n_cols), dtype=X.dtype, device=X.device)
+    RSTD = torch.empty((n_rows, 3), dtype=torch.float32, device=X.device)
+
+    # Check constraints
+    assert W.shape[0] == 3, "Weight tensor must have shape (3,)"
+    assert B.numel() == 1, "Bias must be a scalar"
+
+    # XPU-specific optimization
+    kernel_args = {}
+    if X.device.type == "xpu":
+        kernel_args["grf_mode"] = "large"
+
+    # Launch kernel
+    _poly_norm_forward_kernel[(n_rows,)](
+        Y,
+        Y.stride(0),
+        X,
+        X.stride(0),
+        W,
+        B,
+        RSTD,
+        RSTD.stride(0),
+        n_cols,
+        eps,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+        **kernel_args,
+    )
+
+    return Y.view(*shape), X, RSTD, BLOCK_SIZE, num_warps
+
+
+def poly_norm_backward(dY, X, W, RSTD, BLOCK_SIZE, num_warps, in_place):
+    """
+    PolyNorm Backward Pass
+
+    Args:
+        dY: gradient of output
+        X: input tensor (already reshaped to 2D)
+        W: weight tensor
+        RSTD: cached rstd values from forward
+        BLOCK_SIZE: block size from forward
+        num_warps: number of warps from forward
+        in_place: whether to in-place modify dY to store dX (saves memory)
+
+    Returns:
+        dX: gradient w.r.t. input
+        dW: gradient w.r.t. weight
+        dB: gradient w.r.t. bias
+    """
+    shape = dY.shape
+    dim = shape[-1]
+    dY = dY.view(-1, dim)
+    n_rows, n_cols = dY.shape
+
+    # Get number of SMs for parallelization
+    import math
+
+    sm_count = 1
+    if X.device.type == "cuda":
+        sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count
+    elif X.device.type == "xpu":
+        sm_count = torch.xpu.get_device_properties(X.device).gpu_eu_count
+
+    # Allocate or reuse gradients
+    if in_place is True:
+        dX = dY
+    else:
+        dX = torch.zeros_like(dY)
+
+    _dW = torch.empty((sm_count, 3), dtype=torch.float32, device=W.device)
+    _dB = torch.empty((sm_count,), dtype=torch.float32, device=W.device)
+
+    rows_per_program = math.ceil(n_rows / sm_count)
+    grid = (sm_count,)
+
+    # XPU-specific optimization
+    kernel_args = {}
+    if X.device.type == "xpu":
+        kernel_args["grf_mode"] = "large"
+
+    # Launch backward kernel
+    _poly_norm_backward_kernel[grid](
+        dY,
+        dY.stride(0),
+        dX,
+        dX.stride(0),
+        X,
+        X.stride(0),
+        W,
+        RSTD,
+        RSTD.stride(0),
+        _dW,
+        _dW.stride(0),
+        _dB,
+        n_rows,
+        n_cols,
+        rows_per_program,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+        **kernel_args,
+    )
+
+    # Reduce gradients across SMs
+    dX = dX.view(*shape)
+    dW = _dW.sum(dim=0).to(W.dtype)
+    dB = _dB.sum().to(W.dtype)
+
+    return dX, dW, dB
+
+
+class LigerPolyNormFunction(torch.autograd.Function):
+    """
+    PolyNorm Function with forward and backward pass
+
+    PolyNorm formula:
+        y = w₀·norm(x³) + w₁·norm(x²) + w₂·norm(x) + b
+        where norm(u) = u / sqrt(mean(u²) + ε)
+
+    Backward uses closed-form gradient:
+        ∂L/∂x_i = Σ_p w_p * [p*x_i^(p-1) * grad_i/D_p - (p/d)*x_i^(2p-1) * S_p/(D_p³)]
+    """
+
+    @staticmethod
+    @ensure_contiguous
+    def forward(ctx, X, W, B, eps=1e-6, in_place=True):
+        """
+        Args:
+            X: input tensor of shape (B, T, H) or (BxT, H)
+            W: weight tensor of shape (3,) for [w0, w1, w2]
+            B: bias scalar
+            eps: epsilon for numerical stability
+            in_place: whether to in-place modify grad_output in backward (saves memory)
+
+        Returns:
+            Y: output tensor of same shape as X
+        """
+        Y, X, RSTD, BLOCK_SIZE, num_warps = poly_norm_forward(X, W, B, eps)
+        ctx.BLOCK_SIZE = BLOCK_SIZE
+        ctx.num_warps = num_warps
+        ctx.in_place = in_place
+        ctx.save_for_backward(X, W, RSTD)
+        return Y
+
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, grad_output):
+        """
+        Args:
+            grad_output: gradient of output
+
+        Returns:
+            dX, dW, dB: gradients w.r.t. X, W, B
+        """
+        X, W, RSTD = ctx.saved_tensors
+        dX, dW, dB = poly_norm_backward(grad_output, X, W, RSTD, ctx.BLOCK_SIZE, ctx.num_warps, ctx.in_place)
+        return dX, dW, dB, None, None

liger_kernel/transformers/__init__.py

@@ -15,6 +15,7 @@ from liger_kernel.transformers.layer_norm import LigerLayerNorm  # noqa: F401
 from liger_kernel.transformers.llama4_rope import liger_llama4_text_rotary_pos_emb  # noqa: F401
 from liger_kernel.transformers.llama4_rope import liger_llama4_vision_rotary_pos_emb  # noqa: F401
 from liger_kernel.transformers.multi_token_attention import LigerMultiTokenAttention  # noqa: F401
+from liger_kernel.transformers.poly_norm import LigerPolyNorm  # 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.softmax import LigerSoftmax  # noqa: F401
@@ -30,13 +31,16 @@ if TYPE_CHECKING:
     from liger_kernel.transformers.auto_model import AutoLigerKernelForCausalLM  # noqa: F401
     from liger_kernel.transformers.monkey_patch import _apply_liger_kernel  # noqa: F401
     from liger_kernel.transformers.monkey_patch import _apply_liger_kernel_to_instance  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_falcon_h1  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma2  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma3  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma3_text  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_glm4  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_glm4v  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_glm4v_moe  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_granite  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_internvl  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_llama  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_llama4  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_llava  # noqa: F401
@@ -51,7 +55,9 @@ if TYPE_CHECKING:
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen2_vl  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen3  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen3_moe  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen3_next  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_smollm3  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_smolvlm  # noqa: F401
 
 
 # Check if 'transformers' is installed
@@ -89,13 +95,16 @@ def __getattr__(name: str):
     monkey_patch_symbols = {
         "_apply_liger_kernel",
         "_apply_liger_kernel_to_instance",
+        "apply_liger_kernel_to_falcon_h1",
         "apply_liger_kernel_to_gemma",
         "apply_liger_kernel_to_gemma2",
         "apply_liger_kernel_to_gemma3",
         "apply_liger_kernel_to_gemma3_text",
         "apply_liger_kernel_to_glm4",
         "apply_liger_kernel_to_glm4v",
+        "apply_liger_kernel_to_glm4v_moe",
         "apply_liger_kernel_to_granite",
+        "apply_liger_kernel_to_internvl",
         "apply_liger_kernel_to_llama",
         "apply_liger_kernel_to_llava",
         "apply_liger_kernel_to_llama4",
@@ -110,7 +119,9 @@ def __getattr__(name: str):
         "apply_liger_kernel_to_qwen2_vl",
         "apply_liger_kernel_to_qwen3",
         "apply_liger_kernel_to_qwen3_moe",
+        "apply_liger_kernel_to_qwen3_next",
         "apply_liger_kernel_to_smollm3",
+        "apply_liger_kernel_to_smolvlm",
     }
 
     if name in monkey_patch_symbols:
@@ -131,6 +142,7 @@ __all__ = [
     "LigerJSD",
     "LigerLayerNorm",
     "LigerFusedAddRMSNorm",
+    "LigerPolyNorm",
     "LigerRMSNorm",
     "liger_rotary_pos_emb",
     "liger_llama4_text_rotary_pos_emb",
@@ -153,13 +165,16 @@ if _TRANSFORMERS_AVAILABLE:
             "AutoLigerKernelForCausalLM",
             "_apply_liger_kernel",
             "_apply_liger_kernel_to_instance",
+            "apply_liger_kernel_to_falcon_h1",
             "apply_liger_kernel_to_gemma",
             "apply_liger_kernel_to_gemma2",
             "apply_liger_kernel_to_gemma3",
             "apply_liger_kernel_to_gemma3_text",
             "apply_liger_kernel_to_glm4",
             "apply_liger_kernel_to_glm4v",
+            "apply_liger_kernel_to_glm4v_moe",
             "apply_liger_kernel_to_granite",
+            "apply_liger_kernel_to_internvl",
             "apply_liger_kernel_to_llama",
             "apply_liger_kernel_to_llava",
             "apply_liger_kernel_to_llama4",
@@ -174,6 +189,8 @@ if _TRANSFORMERS_AVAILABLE:
             "apply_liger_kernel_to_qwen2_vl",
             "apply_liger_kernel_to_qwen3",
             "apply_liger_kernel_to_qwen3_moe",
+            "apply_liger_kernel_to_qwen3_next",
             "apply_liger_kernel_to_smollm3",
+            "apply_liger_kernel_to_smolvlm",
         ]
     )

liger_kernel/transformers/functional.py

@@ -12,6 +12,7 @@ from liger_kernel.ops.jsd import LigerJSDFunction
 from liger_kernel.ops.kl_div import LigerKLDivLossFunction
 from liger_kernel.ops.layer_norm import LigerLayerNormFunction
 from liger_kernel.ops.multi_token_attention import LigerMultiTokenAttentionFunction
+from liger_kernel.ops.poly_norm import LigerPolyNormFunction
 from liger_kernel.ops.qwen2vl_mrope import LigerQwen2VLMRopeFunction
 from liger_kernel.ops.rms_norm import LigerRMSNormFunction
 from liger_kernel.ops.rope import LigerRopeFunction
@@ -65,6 +66,7 @@ def liger_fused_linear_cross_entropy(
     softcap: Optional[float] = None,
     return_z_loss: bool = False,
     accum_dtype=None,
+    use_token_scaling: bool = False,
 ):
     loss, z_loss = LigerFusedLinearCrossEntropyFunction.apply(
         input,
@@ -79,6 +81,7 @@ def liger_fused_linear_cross_entropy(
         softcap,
         return_z_loss,
         accum_dtype,
+        use_token_scaling,
     )
     if not return_z_loss:
         return loss
@@ -256,6 +259,10 @@ def liger_rms_norm(X, W, eps, offset: float = 0.0, casting_mode: str = "llama",
     return LigerRMSNormFunction.apply(X, W, eps, offset, casting_mode, in_place)
 
 
+def liger_poly_norm(X, W, B, eps=1e-6, in_place=True):
+    return LigerPolyNormFunction.apply(X, W, B, eps, in_place)
+
+
 def liger_fused_add_rms_norm(X, R, W, eps, offset: float = 0.0, casting_mode: str = "llama", in_place: bool = True):
     return LigerFusedAddRMSNormFunction.apply(X, R, W, eps, offset, casting_mode, in_place)
 

liger_kernel/transformers/fused_linear_cross_entropy.py

@@ -16,6 +16,7 @@ class LigerFusedLinearCrossEntropyLoss(torch.nn.Module):
         softcap: Optional[float] = None,
         return_z_loss: bool = False,
         accum_dtype: Optional[torch.dtype] = None,
+        use_token_scaling: bool = False,
     ):
         super().__init__()
         assert (label_smoothing >= 0) and (label_smoothing <= 1), (
@@ -24,7 +25,8 @@ class LigerFusedLinearCrossEntropyLoss(torch.nn.Module):
         assert reduction in {
             "mean",
             "sum",
-        }, f"reduction must be 'mean' or 'sum'. Got: {reduction}"
+            "none",
+        }, f"reduction must be 'mean' or 'sum' or 'none'. Got: {reduction}"
         assert softcap is None or softcap > 0, f"softcap must greater than 0.0 or None. Got: {softcap}"
         self.ce_weight = ce_weight
         self.ignore_index = ignore_index
@@ -34,6 +36,7 @@ class LigerFusedLinearCrossEntropyLoss(torch.nn.Module):
         self.softcap = softcap
         self.return_z_loss = return_z_loss
         self.accum_dtype = accum_dtype
+        self.use_token_scaling = use_token_scaling
 
     def forward(self, lin_weight, _input, target, bias=None):
         loss, z_loss = LigerFusedLinearCrossEntropyFunction.apply(
@@ -49,6 +52,7 @@ class LigerFusedLinearCrossEntropyLoss(torch.nn.Module):
             self.softcap,
             self.return_z_loss,
             self.accum_dtype,
+            self.use_token_scaling,
         )
         if not self.return_z_loss:
             return loss

liger_kernel/transformers/model/falcon_h1.py

@@ -0,0 +1,108 @@
+from typing import TYPE_CHECKING
+from typing import Optional
+from typing import Union
+
+import torch
+
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+if TYPE_CHECKING:
+    from transformers.models.falcon_h1.modeling_falcon_h1 import FalconHybridMambaAttentionDynamicCache
+
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
+
+
+def lce_forward(
+    self,
+    input_ids: torch.LongTensor = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional["FalconHybridMambaAttentionDynamicCache"] = None,
+    inputs_embeds: Optional[torch.FloatTensor] = None,
+    labels: Optional[torch.LongTensor] = None,
+    use_cache: Optional[bool] = None,
+    output_attentions: Optional[bool] = None,
+    output_hidden_states: Optional[bool] = None,
+    cache_position: Optional[torch.LongTensor] = None,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
+    skip_logits: Optional[bool] = None,
+    **kwargs,
+) -> Union[tuple, CausalLMOutputWithPast]:
+    r"""
+    labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+        Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+        config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+        (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+    Example:
+
+    ```python
+    >>> from transformers import AutoTokenizer, FalconH1ForCausalLM
+
+    >>> model = FalconH1ForCausalLM.from_pretrained("...")
+    >>> tokenizer = AutoTokenizer.from_pretrained("...")
+
+    >>> prompt = "Hey, are you conscious? Can you talk to me?"
+    >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+    >>> # Generate
+    >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+    >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+    "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+    ```"""
+    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+    output_hidden_states = (
+        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+    )
+
+    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+    outputs = self.model(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        position_ids=position_ids,
+        past_key_values=past_key_values,
+        inputs_embeds=inputs_embeds,
+        use_cache=use_cache,
+        output_attentions=output_attentions,
+        output_hidden_states=output_hidden_states,
+        cache_position=cache_position,
+        **kwargs,
+    )
+
+    hidden_states = outputs[0]
+    # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+    kept_hidden_states = hidden_states[:, slice_indices, :]
+
+    shift_labels = kwargs.pop("shift_labels", None)
+    logits = None
+    loss = None
+    # if in training mode, don't materialize logits
+    if skip_logits and labels is None:
+        raise ValueError("skip_logits is True, but labels and shift_labels are None")
+
+    if skip_logits is None:
+        # By default, if in training mode, don't materialize logits
+        skip_logits = self.training and labels is not None
+
+    if skip_logits:
+        loss = LigerForCausalLMLoss(
+            hidden_states=kept_hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            shift_labels=shift_labels,
+            hidden_size=self.config.hidden_size,
+            **kwargs,
+        )
+    else:
+        logits = self.lm_head(kept_hidden_states)
+        if labels is not None or shift_labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+    return CausalLMOutputWithPast(
+        loss=loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+    )

liger_kernel/transformers/model/gemma.py

@@ -228,10 +228,11 @@ def lce_forward(
         )
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

liger_kernel/transformers/model/gemma2.py

@@ -252,8 +252,14 @@ def lce_forward(
             logits = logits * self.config.final_logit_softcapping
 
         loss = None
-        if labels is not None:
-            loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+        if labels is not None or shift_labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                shift_labels=shift_labels,
+                vocab_size=self.vocab_size,
+                **kwargs,
+            )
 
     if not return_dict:
         output = (logits,) + outputs[1:]