liger-kernel-nightly 0.3.1.dev20241031232423__py3-none-any.whl → 0.3.1.dev20241101201851__py3-none-any.whl

liger_kernel/ops/fused_linear_jsd.py

@@ -4,7 +4,7 @@ import torch
 import triton
 
 from liger_kernel.ops.jsd import _jsd_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
@@ -24,6 +24,7 @@ def fused_linear_jsd_forward(
     temperature,
 ):
     device = student_input.device
+    dtype = student_input.dtype
 
     # inputs have shape: BT x H
     # materialized activations will have shape: BT x V
@@ -64,9 +65,15 @@ def fused_linear_jsd_forward(
         student_input_chunk = student_input[start_idx:end_idx]
         teacher_input_chunk = teacher_input[start_idx:end_idx]
 
-        # when doing matmul, use the original precision, shape: chunk_size x V
-        student_logits_chunk = student_input_chunk @ student_weight.t()
-        teacher_logits_chunk = teacher_input_chunk @ teacher_weight.t()
+        # shape: chunk_size x V
+        # For anything starting from logits to the final JSD loss, we do computation
+        # in FP32 to avoid losing numerical stability.
+        student_logits_chunk = (student_input_chunk @ student_weight.t()).to(
+            torch.float32
+        )
+        teacher_logits_chunk = (teacher_input_chunk @ teacher_weight.t()).to(
+            torch.float32
+        )
         chunk_n_rows = student_logits_chunk.shape[0]
 
         # unreduced loss
@@ -113,18 +120,16 @@ def fused_linear_jsd_forward(
                 student_prob_chunk.shape
             )
         ) / temperature
+        # now we traverse back to grad w.r.t. input to `lm_head` and grad
+        # w.r.t. `lm_head` which should be computed in original dtype
+        student_logits_chunk = student_logits_chunk.to(dtype)
         grad_input[start_idx:end_idx] = student_logits_chunk @ student_weight
 
         if grad_weight is not None:
-            torch.addmm(
-                input=grad_weight,
-                mat1=student_logits_chunk.t(),  # gradients of logits_chunk
-                mat2=student_input_chunk,
-                out=grad_weight,
-            )
+            grad_weight.add_(student_logits_chunk.t() @ student_input_chunk)
 
     loss = torch.sum(loss_1d)
-    return loss.to(student_input.dtype), grad_input, grad_weight
+    return loss, grad_input, grad_weight
 
 
 def fused_linear_jsd_backward(grad_output, grad_input, grad_weight):
@@ -172,6 +177,7 @@ class LigerFusedLinearJSDFunction(torch.autograd.Function):
     """
 
     @staticmethod
+    @amp_custom_fwd
     def forward(
         ctx,
         student_input: torch.Tensor,
@@ -225,6 +231,7 @@ class LigerFusedLinearJSDFunction(torch.autograd.Function):
         return loss
 
     @staticmethod
+    @amp_custom_bwd
     def backward(ctx, grad_output):
         (grad_input, grad_weight) = ctx.saved_tensors
         grad_input, grad_weight = fused_linear_jsd_backward(

{liger_kernel_nightly-0.3.1.dev20241031232423.dist-info → liger_kernel_nightly-0.3.1.dev20241101201851.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.3.1.dev20241031232423
+Version: 0.3.1.dev20241101201851
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 
@@ -36,7 +36,7 @@ License-File: LICENSE-MIT-llmc
 License-File: LICENSE-MIT-triton
 License-File: NOTICE
 Requires-Dist: torch>=2.1.2
-Requires-Dist: triton>=2.3.0
+Requires-Dist: triton>=2.3.1
 Provides-Extra: dev
 Requires-Dist: transformers>=4.44.2; extra == "dev"
 Requires-Dist: matplotlib>=3.7.2; extra == "dev"

{liger_kernel_nightly-0.3.1.dev20241031232423.dist-info → liger_kernel_nightly-0.3.1.dev20241101201851.dist-info}/RECORD

@@ -2,7 +2,7 @@ liger_kernel/env_report.py,sha256=LFUJ6UMkFFGPBYXBlqHFGy4bhsemEpSI-_1edSazlHI,11
 liger_kernel/ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 liger_kernel/ops/cross_entropy.py,sha256=OB3nvIONLB_sj9LO6UQv1qLnf861k-pR58RtwgoiyYA,11192
 liger_kernel/ops/fused_linear_cross_entropy.py,sha256=qg7qBQFLDJClnkUOGhFFHPSW_x7rPvQekbm_4OOYxys,9331
-liger_kernel/ops/fused_linear_jsd.py,sha256=eZ8y4GPtPRE0QcNNMLX8l4gSEMPzA3ZuzknfbAbiREA,9234
+liger_kernel/ops/fused_linear_jsd.py,sha256=ZQUxNqm3yOokZhUId0sIfob_3e43rGbMTDxeCk9A92o,9549
 liger_kernel/ops/geglu.py,sha256=MQL4zyzneZqZYUGPvb1QjI_EYT9_pKfSDgR25WD9jrI,4127
 liger_kernel/ops/jsd.py,sha256=anWfdioucxZy4JQfTvbHBR-IQrZKeH-gBF1MHwwTuTQ,5781
 liger_kernel/ops/kl_div.py,sha256=qnmtFQwuO3FR7Ovup_DDzpkD1A1LpwOaWlcO6K9ysHk,8342
@@ -40,14 +40,14 @@ liger_kernel/transformers/model/qwen2.py,sha256=3inWFXGHYT7wA10OR6bq3mDUBrr10AS5
 liger_kernel/transformers/model/qwen2_vl.py,sha256=ymsm9aQpSUiSU12GY8FO608p9dSHOz4TCnNI1htX5bk,6975
 liger_kernel/triton/__init__.py,sha256=yfRe0zMb47QnqjecZWG7LnanfCTzeku7SgWRAwNVmzU,101
 liger_kernel/triton/monkey_patch.py,sha256=5BcGKTtdqeYchypBIBopGIWPx1-cFALz7sOKoEsqXJ0,1584
-liger_kernel_nightly-0.3.1.dev20241031232423.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
-liger_kernel_nightly-0.3.1.dev20241031232423.dist-info/LICENSE-Apache-2.0,sha256=NRaCIsL9eblGS35gk4WKTC0usNYnR_mgRHJTKqz2_UE,11348
-liger_kernel_nightly-0.3.1.dev20241031232423.dist-info/LICENSE-MIT-AutoAWQ,sha256=pfiOyInrAPY3xQbvV1i-gOqNZK7QEyIepT1IbqOYYYo,1067
-liger_kernel_nightly-0.3.1.dev20241031232423.dist-info/LICENSE-MIT-Efficient-Cross-Entropy,sha256=PaC9HqyFYTy-ClS0H8Zfa2motJuTppjECXmjHwJcaOk,1063
-liger_kernel_nightly-0.3.1.dev20241031232423.dist-info/LICENSE-MIT-llmc,sha256=kyFLt_XUcXS88CuxQt5-PjOcLjpJP2m-T4gtqZf3GLc,1071
-liger_kernel_nightly-0.3.1.dev20241031232423.dist-info/LICENSE-MIT-triton,sha256=wL6W8IwsKiyHtzXubg8TCXhRZuo8S83EPdqXffYtqWg,1131
-liger_kernel_nightly-0.3.1.dev20241031232423.dist-info/METADATA,sha256=LPXFEjmB3kvrLlkGnL0jwPSUNPXavpb5c_LLwSpXZ78,27717
-liger_kernel_nightly-0.3.1.dev20241031232423.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
-liger_kernel_nightly-0.3.1.dev20241031232423.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-liger_kernel_nightly-0.3.1.dev20241031232423.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
-liger_kernel_nightly-0.3.1.dev20241031232423.dist-info/RECORD,,
+liger_kernel_nightly-0.3.1.dev20241101201851.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
+liger_kernel_nightly-0.3.1.dev20241101201851.dist-info/LICENSE-Apache-2.0,sha256=NRaCIsL9eblGS35gk4WKTC0usNYnR_mgRHJTKqz2_UE,11348
+liger_kernel_nightly-0.3.1.dev20241101201851.dist-info/LICENSE-MIT-AutoAWQ,sha256=pfiOyInrAPY3xQbvV1i-gOqNZK7QEyIepT1IbqOYYYo,1067
+liger_kernel_nightly-0.3.1.dev20241101201851.dist-info/LICENSE-MIT-Efficient-Cross-Entropy,sha256=PaC9HqyFYTy-ClS0H8Zfa2motJuTppjECXmjHwJcaOk,1063
+liger_kernel_nightly-0.3.1.dev20241101201851.dist-info/LICENSE-MIT-llmc,sha256=kyFLt_XUcXS88CuxQt5-PjOcLjpJP2m-T4gtqZf3GLc,1071
+liger_kernel_nightly-0.3.1.dev20241101201851.dist-info/LICENSE-MIT-triton,sha256=wL6W8IwsKiyHtzXubg8TCXhRZuo8S83EPdqXffYtqWg,1131
+liger_kernel_nightly-0.3.1.dev20241101201851.dist-info/METADATA,sha256=qx1HwgDXyy5RdAAuvpVrmpnLA4CJ6b6A7zGVQrU-rpg,27717
+liger_kernel_nightly-0.3.1.dev20241101201851.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
+liger_kernel_nightly-0.3.1.dev20241101201851.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+liger_kernel_nightly-0.3.1.dev20241101201851.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
+liger_kernel_nightly-0.3.1.dev20241101201851.dist-info/RECORD,,