liger-kernel-nightly 0.5.2.dev20250115025212__py3-none-any.whl → 0.5.2.dev20250118115043__py3-none-any.whl

liger_kernel/ops/fused_linear_cross_entropy.py

@@ -24,7 +24,9 @@ def fused_linear_cross_entropy_forward(
     label_smoothing=0.0,
     reduction="mean",
     softcap=None,
+    return_z_loss=False,
 ):
+    assert isinstance(return_z_loss, bool), f"return_z_loss must be True or False. Got: {return_z_loss}"
     device = _input.device
 
     # inputs have shape: BT x H
@@ -47,6 +49,7 @@ def fused_linear_cross_entropy_forward(
     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)
+    z_loss_1d = torch.zeros(BT, dtype=_input.dtype, device=_input.device) if return_z_loss else None
 
     # TODO: evaluate how CUDA synchronization caused by .item() affects the speed
     target_mask = target != ignore_index
@@ -81,6 +84,7 @@ def fused_linear_cross_entropy_forward(
 
         # unreduced loss
         loss_1d_slice = loss_1d[start_idx:end_idx]  # chunk_size,
+        z_loss_1d_slice = z_loss_1d[start_idx:end_idx] if return_z_loss else None
 
         # ensure _input and target are contiguous
         logits_chunk = logits_chunk.contiguous()
@@ -94,7 +98,7 @@ def fused_linear_cross_entropy_forward(
             Y_stride=target_chunk.stride(-1),  # always 1
             weight_ptr=ce_weight,
             loss_ptr=loss_1d_slice,
-            z_loss_ptr=None,
+            z_loss_ptr=z_loss_1d_slice,
             loss_stride=loss_1d_slice.stride(-1),  # always 1
             n_cols=V,
             n_non_ignore=total_n_non_ignore,
@@ -105,7 +109,7 @@ def fused_linear_cross_entropy_forward(
             label_smoothing=label_smoothing,
             reduction=reduction,
             softcap=softcap,
-            RETURN_Z_LOSS=False,
+            RETURN_Z_LOSS=return_z_loss,
             HAS_WEIGHT=True if ce_weight is not None else False,
             HAS_SOFTCAPPING=True if softcap is not None else False,
             BLOCK_SIZE=BLOCK_SIZE,
@@ -113,6 +117,8 @@ def fused_linear_cross_entropy_forward(
         )
 
         loss_1d[start_idx:end_idx] = loss_1d_slice
+        if return_z_loss:
+            z_loss_1d[start_idx:end_idx] = z_loss_1d_slice
         grad_logits_chunk = logits_chunk  # chunk_size x V
 
         grad_input[start_idx:end_idx] = grad_logits_chunk @ weight
@@ -139,9 +145,13 @@ def fused_linear_cross_entropy_forward(
 
     if reduction == "none":
         loss = loss_1d
+        z_loss = z_loss_1d if return_z_loss else None
+
     else:
         loss = torch.sum(loss_1d)
-    return loss, grad_input, grad_weight, grad_bias
+        z_loss = torch.sum(z_loss_1d) if return_z_loss else None
+    print(f"{z_loss=}")
+    return loss, z_loss, grad_input, grad_weight, grad_bias
 
 
 def fused_linear_cross_entropy_backward(grad_output, grad_input, grad_weight, grad_bias):
@@ -206,6 +216,7 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
         label_smoothing=0.0,
         reduction="mean",
         softcap=None,
+        return_z_loss: bool = False,
     ):
         """
         Fusing the last linear layer with cross-entropy loss
@@ -226,7 +237,7 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
         reduction: reduction to apply
         """
 
-        loss, grad_input, grad_weight, grad_bias = fused_linear_cross_entropy_forward(
+        loss, z_loss, grad_input, grad_weight, grad_bias = fused_linear_cross_entropy_forward(
             _input=_input,
             weight=weight,
             target=target,
@@ -237,6 +248,7 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
             label_smoothing=label_smoothing,
             reduction=reduction,
             softcap=softcap,
+            return_z_loss=return_z_loss,
         )
         # downcast to dtype and store for backward
         ctx.save_for_backward(
@@ -244,11 +256,14 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
             grad_weight.detach() if grad_weight is not None else None,
             grad_bias.detach() if bias is not None else None,
         )
-        return loss
+        ctx.return_z_loss = return_z_loss
+        return loss, z_loss
 
     @staticmethod
     @amp_custom_bwd
-    def backward(ctx, grad_output):
+    def backward(ctx, grad_output, grad_output2):
+        if ctx.return_z_loss:
+            del grad_output2  # z_loss is only for logging
         (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
@@ -264,4 +279,5 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
             None,
             None,
             None,
+            None,
         )

liger_kernel/transformers/functional.py

@@ -56,8 +56,9 @@ def liger_fused_linear_cross_entropy(
     label_smoothing: float = 0.0,
     reduction: str = "mean",
     softcap: Optional[float] = None,
+    return_z_loss: bool = False,
 ):
-    return LigerFusedLinearCrossEntropyFunction.apply(
+    loss, z_loss = LigerFusedLinearCrossEntropyFunction.apply(
         input,
         weight,
         target,
@@ -68,7 +69,11 @@ def liger_fused_linear_cross_entropy(
         label_smoothing,
         reduction,
         softcap,
+        return_z_loss,
     )
+    if not return_z_loss:
+        return loss
+    return loss, z_loss
 
 
 def liger_fused_linear_jsd(

liger_kernel/transformers/fused_linear_cross_entropy.py

@@ -14,6 +14,7 @@ class LigerFusedLinearCrossEntropyLoss(torch.nn.Module):
         label_smoothing: float = 0.0,
         reduction: str = "mean",
         softcap: Optional[float] = None,
+        return_z_loss: bool = False,
     ):
         super().__init__()
         assert (label_smoothing >= 0) and (
@@ -31,9 +32,10 @@ class LigerFusedLinearCrossEntropyLoss(torch.nn.Module):
         self.label_smoothing = label_smoothing
         self.reduction = reduction
         self.softcap = softcap
+        self.return_z_loss = return_z_loss
 
     def forward(self, lin_weight, _input, target, bias=None):
-        return LigerFusedLinearCrossEntropyFunction.apply(
+        loss, z_loss = LigerFusedLinearCrossEntropyFunction.apply(
             _input,
             lin_weight,
             target,
@@ -44,4 +46,8 @@ class LigerFusedLinearCrossEntropyLoss(torch.nn.Module):
             self.label_smoothing,
             self.reduction,
             self.softcap,
+            self.return_z_loss,
         )
+        if not self.return_z_loss:
+            return loss
+        return loss, z_loss

{liger_kernel_nightly-0.5.2.dev20250115025212.dist-info → liger_kernel_nightly-0.5.2.dev20250118115043.dist-info}/METADATA

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

{liger_kernel_nightly-0.5.2.dev20250115025212.dist-info → liger_kernel_nightly-0.5.2.dev20250118115043.dist-info}/RECORD

@@ -12,7 +12,7 @@ liger_kernel/chunked_loss/orpo_loss.py,sha256=yjcrrbVeemLYodoSKT-FMSnaPtyKAZ3aOr
 liger_kernel/chunked_loss/simpo_loss.py,sha256=3TTc7U79Orjgi-Wu81WZkWk5MgsdqKXIOBHgIvDazPw,3865
 liger_kernel/ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 liger_kernel/ops/cross_entropy.py,sha256=SRzAF9Ek84pBVFy3wqQZs7AhRoorKRIgQ-Td_rtl1Kk,18606
-liger_kernel/ops/fused_linear_cross_entropy.py,sha256=hezFRwbcPc-HNGZUFqUn5AYUqUpboPpFh4MNqEW4WgU,10108
+liger_kernel/ops/fused_linear_cross_entropy.py,sha256=Ad8mhD1LZSX-kNSi8xY_zeM-Gq-BzrleyDdf8E0c3Rs,10934
 liger_kernel/ops/fused_linear_jsd.py,sha256=eKqaADj7LgWfoYqyH03tjrmhNTfJOF1Dhx_bWzBTnTU,9600
 liger_kernel/ops/geglu.py,sha256=axGvCIvlBzuluoAIrWTsp2iZM4BFKNInkPov8YVvH9E,4126
 liger_kernel/ops/group_norm.py,sha256=qD4D4lSjSgVtO52EBNLC2iTseALRgPgqXE50U2woggk,10837
@@ -29,8 +29,8 @@ liger_kernel/ops/experimental/mm_int8int2.py,sha256=TrS9lpwekrik_w5qE7AhMJD1bcq-
 liger_kernel/transformers/__init__.py,sha256=QPmYkL6hosBPpPqCUGqvIvAtD9XzLgvZqZxUyYMZeVk,2008
 liger_kernel/transformers/auto_model.py,sha256=0qCTRZt280Bj_LcFdzo9hlaR-BWNazawXOGgoCZjgEg,1545
 liger_kernel/transformers/cross_entropy.py,sha256=LtiHlj_tK2YFpilwvbG_NEVzbf82zKRpWCZMjaFUd4M,1681
-liger_kernel/transformers/functional.py,sha256=B1wkHWLx-YNhxvXBEXB4Ch1yEwF3mjwTPCeXA5aCV_c,4490
-liger_kernel/transformers/fused_linear_cross_entropy.py,sha256=LAN8-pjUI2Erz_MnfMer-0ZmxJ0JlKxGzdZGJY-N65g,1569
+liger_kernel/transformers/functional.py,sha256=lDOjch622dJIc78K3ePFK_H1DX00GC5kKjodjcbEgbM,4624
+liger_kernel/transformers/fused_linear_cross_entropy.py,sha256=ygU7cycCHWvSrrTgn2TseN8t-Qwfer4V7ldwhZ1E6_g,1776
 liger_kernel/transformers/fused_linear_jsd.py,sha256=bZ4otCvWBuOnA5XdQL-FzZVItJlDt-ht9e_pG7PG93E,3999
 liger_kernel/transformers/geglu.py,sha256=mrgqzIUVd6lN7fkDKLkw5YaESDxDtFgbot430WwPVOQ,1107
 liger_kernel/transformers/group_norm.py,sha256=URmjkQFsrbMffzcJiGpX7ckxWlpL95AiJS-80hwAWPk,2173
@@ -58,9 +58,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.2.dev20250115025212.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
-liger_kernel_nightly-0.5.2.dev20250115025212.dist-info/METADATA,sha256=sDbWP1JYiBLWMiChKJMWMd5sVddlHdRk0Jw8oQnkEP8,21055
-liger_kernel_nightly-0.5.2.dev20250115025212.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
-liger_kernel_nightly-0.5.2.dev20250115025212.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-liger_kernel_nightly-0.5.2.dev20250115025212.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
-liger_kernel_nightly-0.5.2.dev20250115025212.dist-info/RECORD,,
+liger_kernel_nightly-0.5.2.dev20250118115043.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
+liger_kernel_nightly-0.5.2.dev20250118115043.dist-info/METADATA,sha256=-JNLpqcUGxtW2t-EcDWMEkSyjE5B1ZyyHzJJs3P_ZyI,21055
+liger_kernel_nightly-0.5.2.dev20250118115043.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
+liger_kernel_nightly-0.5.2.dev20250118115043.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+liger_kernel_nightly-0.5.2.dev20250118115043.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
+liger_kernel_nightly-0.5.2.dev20250118115043.dist-info/RECORD,,