liger-kernel 0.1.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

liger_kernel/ops/fused_linear_cross_entropy.py

@@ -1,7 +1,10 @@
 import torch
 import triton
 
-from liger_kernel.ops.cross_entropy import element_mul, liger_cross_entropy_kernel
+from liger_kernel.ops.cross_entropy import (
+    element_mul_kernel,
+    liger_cross_entropy_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
@@ -9,153 +12,227 @@ from liger_kernel.ops.cross_entropy import element_mul, liger_cross_entropy_kern
 MAX_FUSED_SIZE = 65536 // 2
 
 
-class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, _input, linear, target, ignore_index):
-        """
-        Fusing the last linear layer with cross-entropy loss
-            Reference: https://github.com/mgmalek/efficient_cross_entropy
+def fused_linear_cross_entropy_forward(
+    _input,
+    weight,
+    target,
+    bias=None,
+    ignore_index=-100,
+    label_smoothing=0.0,
+    reduction="mean",
+):
+    dtype = (
+        torch.get_autocast_gpu_dtype() if torch.is_autocast_enabled() else _input.dtype
+    )
+    device = _input.device
+
+    # inputs have shape: BT x H
+    # materialized activations will have shape: BT x V
+    # the increase in memory = BT x V
+    # reduction can be achieved by partitioning the number of tokens BT into smaller chunks.
+    # for ex: if we were to achieve the same memory consumption as BT x H, then the chunk size should be:
+    # inc_factor = (V+H-1)//H, chunk_size = (BT + inc_factor - 1)//inc_factor
+    # for ex: BT = 4096*4, V = 32000, H = 4096 ==> inc_factor = 8, chunk_size = 2048
+    BT, H = _input.shape
+    V = weight.shape[0]
+    BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
+
+    inc_factor = triton.cdiv(V, H)  # (V + H - 1) // H
+    chunk_size = triton.next_power_of_2(
+        triton.cdiv(BT, inc_factor)
+    )  # (BT + inc_factor - 1) // inc_factor
+    num_chunks = triton.cdiv(BT, chunk_size)  # (BT + chunk_size - 1) // chunk_size
+
+    grad_weight = (
+        torch.zeros_like(weight, device=device) if weight.requires_grad else None
+    )
+    grad_input = torch.zeros_like(_input, device=device)
+    grad_bias = torch.zeros_like(bias, device=device) if bias is not None else None
+    # we use fp32 for loss accumulator
+    loss_1d = torch.zeros(BT, dtype=torch.float32, device=device)
+
+    # NOTE: skip .item() here to avoid CUDA synchronization
+    total_n_non_ignore = (target != ignore_index).sum()
+
+    for chunk_id in range(num_chunks):
+        start_idx = chunk_id * chunk_size
+        end_idx = min((chunk_id + 1) * chunk_size, BT)
+        _input_chunk = _input[start_idx:end_idx]  # chunk_size x H
+
+        # when doing matmul, use the original precision
+        logits_chunk = _input_chunk @ weight.t()  # chunk_size x V
+        if bias is not None:
+            logits_chunk = logits_chunk + bias
+        target_chunk = target[start_idx:end_idx]  # chunk_size,
+
+        n_rows = logits_chunk.shape[0]
+
+        # unreduced loss
+        loss_1d_slice = loss_1d[start_idx:end_idx]  # chunk_size,
+        n_non_ignore = (target_chunk != ignore_index).sum().item()
+
+        # when doing CE, use the upcasted precision
+        logits_chunk = logits_chunk.float()
+
+        # ensure _input and target are contiguous
+        logits_chunk = logits_chunk.contiguous()
+        target_chunk = target_chunk.contiguous()
+
+        # Here we calculate the gradient of logits_chunk in place so we can save memory.
+        liger_cross_entropy_kernel[(n_rows,)](
+            X_ptr=logits_chunk,
+            X_stride=logits_chunk.stride(-2),
+            Y_ptr=target_chunk,
+            Y_stride=target_chunk.stride(-1),  # always 1
+            loss_ptr=loss_1d_slice,
+            loss_stride=loss_1d_slice.stride(-1),  # always 1
+            n_cols=V,
+            n_non_ignore=n_non_ignore,
+            ignore_index=ignore_index,
+            label_smoothing=label_smoothing,
+            reduction=reduction,
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=32,
+        )
 
-        Handle the forward and backward pass of the final linear layer via cross-entropy loss by avoiding
-        the materialization of the large logits tensor. Since Cross Entropy Loss is the last layer, we can
-        compute the gradient at the forward pass. By doing so, we don't have to store the _input and target
-        for the backward pass.
+        # gradient of logits_chunk is computed in-place by the above triton kernel.
+        # Following HuggingFace model source code, we do the forward and backward
+        # w.r.t. logits in fp32 for numerical stability especially as the num classes (vocab size) os huge.
+        # (reference: https://github.com/huggingface/transformers/blob/v4.42.4/src/transformers/models/llama/modeling_llama.py#L1194)
+        # Propagating to lm_head's backward, we'll switch back to the original dtype.
+        logits_chunk = logits_chunk.to(dtype)
 
-        _input: (B*T, H) where B is batch size, T is sequence length, H is hidden dimension.
-        target: (B*T) where each value is in [0, V-1]
-        linear: linear projection matrix of shape V x H.
-        ignore_index: the index to ignore in the target
-        """
-        dtype = (
-            torch.get_autocast_gpu_dtype()
-            if torch.is_autocast_enabled()
-            else _input.dtype
-        )
-        device = _input.device
-
-        # inputs have shape: BT x H
-        # materialized activations will have shape: BT x V
-        # the increase in memory = BT x V
-        # reduction can be achieved by paritioning the number of tokens BT into smaller chunks.
-        # for ex: if we were to achieve the same memory consumption as BT x H, then the chunk size should be:
-        # inc_factor = (V+H-1)//H, chunk_size = (BT + inc_factor - 1)//inc_factor
-        # for ex: BT = 4096*4, V = 32000, H = 4096 ==> inc_factor = 8, chunk_size = 2048
-        BT, H = _input.shape
-        V = linear.shape[0]
-        BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
-
-        inc_factor = triton.cdiv(V, H)  # (V + H - 1) // H
-        chunk_size = triton.next_power_of_2(
-            triton.cdiv(BT, inc_factor)
-        )  # (BT + inc_factor - 1) // inc_factor
-        num_chunks = triton.cdiv(BT, chunk_size)  # (BT + chunk_size - 1) // chunk_size
-
-        grad_linear = torch.zeros_like(linear, device=device)
-        grad_input = torch.zeros_like(_input, device=device)
-
-        # we use fp32 for loss accumulator
-        loss_1d = torch.zeros(BT, dtype=torch.float32, device=device)
-
-        total_n_non_ignore = (target != ignore_index).sum().item()
-
-        for chunk_id in range(num_chunks):
-            start_idx = chunk_id * chunk_size
-            end_idx = min((chunk_id + 1) * chunk_size, BT)
-            _input_chunk = _input[start_idx:end_idx]  # chunk_size x H
-
-            # when doing matmul, use the original precision
-            logits_chunk = _input_chunk @ linear.t()  # chunk_size x V
-            target_chunk = target[start_idx:end_idx]  # chunk_size,
-
-            n_rows = logits_chunk.shape[0]
-
-            # unreduced loss
-            loss_1d_slice = loss_1d[start_idx:end_idx]  # chunk_size,
-            n_non_ignore = (target_chunk != ignore_index).sum().item()
-
-            # when doing CE, use the upcasted precision
-            logits_chunk = logits_chunk.float()
-
-            # ensure _input and target are contiguous
-            logits_chunk = logits_chunk.contiguous()
-            target_chunk = target_chunk.contiguous()
-
-            # Here we calculate the gradient of logits_chunk in place so we can save memory.
-            liger_cross_entropy_kernel[(n_rows,)](
-                X_ptr=logits_chunk,
-                X_stride=logits_chunk.stride(-2),
-                Y_ptr=target_chunk,
-                Y_stride=target_chunk.stride(-1),  # always 1
-                loss_ptr=loss_1d_slice,
-                loss_stride=loss_1d_slice.stride(-1),  # always 1
-                n_cols=V,
-                n_non_ignore=n_non_ignore,
-                ignore_index=ignore_index,
-                BLOCK_SIZE=BLOCK_SIZE,
-                num_warps=32,
-            )
+        # gradient of logits_chunk is computed in-place by the above triton kernel and is of shape: chunk_size x V
+        # thus grad_input[start_idx: end_idx] should be of shape: chunk_size x H
+        # additionally, since we are chunking the inputs, observe that the loss and gradients are calculated only
+        # on `n_non_ignore` tokens. However, the gradient of the input should be calculated for all tokens.
+        # Thus, we need an additional scaling factor of (n_non_ignore/total_n_non_ignore) to scale the gradients.
 
-            # gradient of logits_chunk is computed inplace by the above triton kernel.
-            # Following HuggingFace model source code, we do the forward and backward
-            # w.r.t. logits in fp32 for numerical stability especially as the num classes (vocab size) os huge.
-            # (reference: https://github.com/huggingface/transformers/blob/v4.42.4/src/transformers/models/llama/modeling_llama.py#L1194)
-            # Propagating to lm_head's backward, we'll switch back to the original dtype.
-            logits_chunk = logits_chunk.to(dtype)
-
-            # gradient of logits_chunk is computed inplace by the above triton kernel and is of shape: chunk_size x V
-            # thus grad_input[start_idx: end_idx] should be of shape: chunk_size x H
-            # additionally, since we are chunking the inputs, observe that the loss and gradients are calculated only
-            # on `n_non_ignore` tokens. However, the gradient of the input should be calculated for all tokens.
-            # Thus, we need an additional scaling factor of (n_non_ignore/total_n_non_ignore) to scale the gradients.
-            grad_logits_chunk = logits_chunk * (n_non_ignore / total_n_non_ignore)
-            grad_input[start_idx:end_idx] = grad_logits_chunk @ linear
+        if reduction == "mean":
+            alpha = n_non_ignore / total_n_non_ignore if total_n_non_ignore > 0 else 0.0
+        else:
+            alpha = 1.0
 
+        loss_1d[start_idx:end_idx] = loss_1d_slice * alpha
+        grad_logits_chunk = logits_chunk * alpha  # chunk_size x V
+
+        grad_input[start_idx:end_idx] = grad_logits_chunk @ weight
+
+        if grad_weight is not None:
             torch.addmm(
-                input=grad_linear,
+                input=grad_weight,
                 mat1=logits_chunk.t(),
                 mat2=_input_chunk,
-                out=grad_linear,
-                alpha=n_non_ignore / total_n_non_ignore,
+                out=grad_weight,
+                alpha=alpha,
                 beta=1.0,
             )
 
-        loss = torch.sum(loss_1d) / total_n_non_ignore
+        if bias is not None:
+            torch.add(
+                input=grad_bias,
+                other=logits_chunk.sum(dim=0),
+                out=grad_bias,
+                alpha=alpha,
+            )
 
-        # downcast to dtype and store for backward
-        ctx.save_for_backward(grad_input.detach(), grad_linear.detach())
-        return loss
+    loss = torch.sum(loss_1d)
+    return loss, grad_input, grad_weight, grad_bias
+
+
+def fused_linear_cross_entropy_backward(
+    grad_output, grad_input, grad_weight, grad_bias
+):
+    # If cross entropy is the last layer, grad_output is 1.0. Skip the mul to save time
+    if torch.ne(grad_output, torch.tensor(1.0, device=grad_output.device)):
+        # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place
+        # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton.
+        BT, H = grad_input.shape
+        n_rows = BT
+        BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(H))
+
+        element_mul_kernel[(n_rows,)](
+            grad_input,
+            grad_input.stride(-2),
+            grad_output,
+            H,
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=32,
+        )
 
-    @staticmethod
-    def backward(ctx, grad_output):
-        (grad_input, grad_linear) = ctx.saved_tensors
-        # If cross entropy is the last layer, grad_output is 1.0. Skip the mul to save time
-        if torch.ne(grad_output, torch.tensor(1.0, device=grad_output.device)):
-            # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place
-            # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton.
-            BT, H = grad_input.shape
-            n_rows = BT
-            BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(H))
-
-            element_mul[(n_rows,)](
-                grad_input,
-                grad_input.stride(-2),
+        # handle grad_weight
+        if grad_weight is not None:
+            V, H = grad_weight.shape
+            n_rows = V
+
+            element_mul_kernel[(n_rows,)](
+                grad_weight,
+                grad_weight.stride(-2),
                 grad_output,
                 H,
                 BLOCK_SIZE=BLOCK_SIZE,
                 num_warps=32,
             )
 
-            # handle grad_linear
-            V, H = grad_linear.shape
+        if grad_bias is not None:
+            V = grad_bias.shape[0]
             n_rows = V
 
-            element_mul[(n_rows,)](
-                grad_linear,
-                grad_linear.stride(-2),
+            element_mul_kernel[(n_rows,)](
+                grad_bias,
+                grad_bias.stride(-1),
                 grad_output,
-                H,
+                1,
                 BLOCK_SIZE=BLOCK_SIZE,
                 num_warps=32,
             )
+    return grad_input, grad_weight, grad_bias
+
+
+class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        _input,
+        weight,
+        target,
+        bias=None,
+        ignore_index=-100,
+        label_smoothing=0.0,
+        reduction="mean",
+    ):
+        """
+        Fusing the last linear layer with cross-entropy loss
+            Reference: https://github.com/mgmalek/efficient_cross_entropy
+
+        Handle the forward and backward pass of the final linear layer via cross-entropy loss by avoiding
+        the materialization of the large logits tensor. Since Cross Entropy Loss is the last layer, we can
+        compute the gradient at the forward pass. By doing so, we don't have to store the _input and target
+        for the backward pass.
 
-        return (grad_input, grad_linear, None, None)
+        _input: (B*T, H) where B is batch size, T is sequence length, H is hidden dimension.
+        target: (B*T) where each value is in [0, V-1]
+        weight: (V, H) where V is the number of classes
+        bias: (V) where V is the number of classes
+        ignore_index: the index to ignore in the target
+        label_smoothing (float): The amount of smoothing when computing the loss, where 0.0 means no smoothing.
+        reduction: reduction to apply
+        """
+        loss, grad_input, grad_weight, grad_bias = fused_linear_cross_entropy_forward(
+            _input, weight, target, bias, ignore_index, label_smoothing, reduction
+        )
+        # downcast to dtype and store for backward
+        ctx.save_for_backward(
+            grad_input.detach(),
+            grad_weight.detach() if grad_weight is not None else None,
+            grad_bias.detach() if bias is not None else None,
+        )
+        return loss
+
+    @staticmethod
+    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(
+            grad_output, grad_input, grad_weight, grad_bias
+        )
+        return (grad_input, grad_weight, None, grad_bias, None, None, None)

liger_kernel/ops/geglu.py

@@ -11,7 +11,12 @@ from liger_kernel.ops.utils import (
 )
 
 if compare_version("triton", operator.ge, "3.0.0"):
-    from triton.language.extra.libdevice import tanh
+    try:
+        # typical import path with dispatch available
+        from triton.language.extra.libdevice import tanh
+    except ModuleNotFoundError:
+        # for working with NGC containers
+        from triton.language.extra.cuda.libdevice import tanh
 else:
     from triton.language.math import tanh
 
@@ -87,54 +92,61 @@ def _geglu_tanh_backward_kernel(
     tl.store(b + col_offsets, db_row, mask=mask)
 
 
+def geglu_forward(a, b):
+    ori_shape = a.shape
+
+    n_cols = ori_shape[-1]
+    a = a.view(-1, n_cols)
+    b = b.view(-1, n_cols)
+    c = torch.empty_like(a)
+    n_rows = a.shape[0]
+
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+
+    _geglu_tanh_forward_kernel[(n_rows,)](
+        a,
+        b,
+        c,
+        c.stride(-2),
+        n_cols=n_cols,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+    return a, b, c.view(*ori_shape)
+
+
+def geglu_backward(a, b, dc):
+    ori_shape = dc.shape
+    n_cols = ori_shape[-1]
+    dc = dc.view(-1, n_cols)
+    n_rows = dc.shape[0]
+
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+
+    _geglu_tanh_backward_kernel[(n_rows,)](
+        dc,
+        a,
+        b,
+        dc.stride(-2),
+        n_cols=n_cols,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+
+    return a.view(*ori_shape), b.view(*ori_shape)
+
+
 class LigerGELUMulFunction(torch.autograd.Function):
     @staticmethod
     @ensure_contiguous
     def forward(ctx, a, b):
-        ori_shape = a.shape
-
-        n_cols = ori_shape[-1]
-        a = a.view(-1, n_cols)
-        b = b.view(-1, n_cols)
-        c = torch.zeros_like(a)
-        n_rows = a.shape[0]
-
-        BLOCK_SIZE, num_warps = calculate_settings(n_cols)
-
-        _geglu_tanh_forward_kernel[(n_rows,)](
-            a,
-            b,
-            c,
-            c.stride(-2),
-            n_cols=n_cols,
-            BLOCK_SIZE=BLOCK_SIZE,
-            num_warps=num_warps,
-        )
-
+        a, b, c = geglu_forward(a, b)
         ctx.save_for_backward(a, b)
-
-        return c.view(*ori_shape)
+        return c
 
     @staticmethod
     @ensure_contiguous
     def backward(ctx, dc):
-
-        ori_shape = dc.shape
-        n_cols = ori_shape[-1]
-        dc = dc.view(-1, n_cols)
         a, b = ctx.saved_tensors
-        n_rows = dc.shape[0]
-
-        BLOCK_SIZE, num_warps = calculate_settings(n_cols)
-
-        _geglu_tanh_backward_kernel[(n_rows,)](
-            dc,
-            a,
-            b,
-            dc.stride(-2),
-            n_cols=n_cols,
-            BLOCK_SIZE=BLOCK_SIZE,
-            num_warps=num_warps,
-        )
-
-        return a.view(*ori_shape), b.view(*ori_shape)
+        a, b = geglu_backward(a, b, dc)
+        return a, b