liger-kernel-nightly 0.4.0.dev20241112204954__py3-none-any.whl → 0.4.1.dev20241112233746__py3-none-any.whl

liger_kernel/ops/rms_norm.py

@@ -116,6 +116,8 @@ def _rms_norm_forward_kernel(
 def _rms_norm_backward_kernel(
     dY_ptr,
     dY_row_stride,
+    dX_ptr,
+    dX_row_stride,
     X_ptr,
     X_row_stride,
     X_dtype: tl.constexpr,
@@ -146,6 +148,8 @@ def _rms_norm_backward_kernel(
     dW_row = tl.zeros((BLOCK_SIZE,), dtype=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
 
@@ -184,9 +188,10 @@ def _rms_norm_backward_kernel(
             # here X_row is already in fp32 (see previous if block)
             dW_row += dY_row * (X_row * rstd_row)
 
-        tl.store(dY_ptr + col_offsets, dX_row.to(X_dtype), mask=mask)
+        tl.store(dX_ptr + col_offsets, dX_row.to(X_dtype), mask=mask)
 
         dY_ptr += dY_row_stride
+        dX_ptr += dX_row_stride
         X_ptr += X_row_stride
         RSTD_ptr += RSTD_row_stride
 
@@ -251,7 +256,9 @@ def rms_norm_forward(X, W, eps, offset, casting_mode):
     return Y.view(*shape), X, RSTD, BLOCK_SIZE, num_warps, casting_mode
 
 
-def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warps):
+def rms_norm_backward(
+    dY, X, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warps, in_place
+):
     shape = dY.shape
     dim = shape[-1]
     dY = dY.view(-1, dim)
@@ -265,10 +272,17 @@ def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warp
         raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
     rows_per_program = math.ceil(n_rows / sm_count)
     grid = (sm_count,)
-    # Here we use dY to store the value of dX to save memory
+
+    if in_place is True:
+        dX = dY
+    else:
+        dX = torch.zeros_like(dY)
+
     _rms_norm_backward_kernel[grid](
         dY,
         dY.stride(0),
+        dX,
+        dX.stride(0),
         X,
         X.stride(0),
         torch_to_triton_dtype[X.dtype],
@@ -286,8 +300,9 @@ def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warp
         BLOCK_SIZE=BLOCK_SIZE,
         num_warps=num_warps,
     )
-    dX = dY.view(*shape)
+    dX = dX.view(*shape)
     dW = _dW.sum(dim=0).to(W.dtype)
+
     return dX, dW
 
 
@@ -307,11 +322,15 @@ class LigerRMSNormFunction(torch.autograd.Function):
     - 'llama': matches the Llama implementation, where only the inverse RMS is computed on fp32.
     - 'gemma': matches the Gemma implementation, where everything is cast to fp32, then computed, then cast back to the original dtype.
     - 'none': no casting is done. The computation is done in the original dtype. This saves memory and is slightly faster, but has more error w.r.t. the original implementation.
+
+    `in_place` option means whether to in_place modify dY to store dX. This is default to `True` to save memory. However, under certain cases, it can produce incorrect inputs.
+        For example, gemma2 uses two rmsnorm sequentially with residual in between. The resesidual part needs dY so it cannot be modified in-place.
+        Therefore, for the patching of RMSNorm in gemma2, we set `in_place` to `False`
     """
 
     @staticmethod
     @ensure_contiguous
-    def forward(ctx, X, W, eps, offset=0.0, casting_mode="llama"):
+    def forward(ctx, X, W, eps, offset=0.0, casting_mode="llama", in_place=True):
         """
         X: (B, T, H) or (BxT, H)
         W: (H,)
@@ -321,6 +340,7 @@ class LigerRMSNormFunction(torch.autograd.Function):
         )
         ctx.offset = offset
         ctx.casting_mode = casting_mode
+        ctx.in_place = in_place
         ctx.BLOCK_SIZE = BLOCK_SIZE
         ctx.num_warps = num_warps
         ctx.save_for_backward(X, W, RSTD)
@@ -342,5 +362,6 @@ class LigerRMSNormFunction(torch.autograd.Function):
             ctx.casting_mode,
             ctx.BLOCK_SIZE,
             ctx.num_warps,
+            ctx.in_place,
         )
-        return dX, dW, None, None, None
+        return dX, dW, None, None, None, None

liger_kernel/transformers/monkey_patch.py

@@ -507,7 +507,7 @@ def apply_liger_kernel_to_gemma2(
     from transformers.models.gemma2.modeling_gemma2 import Gemma2Model
 
     LigerRMSNormForGemma2 = partial(
-        LigerRMSNorm, offset=1.0, casting_mode="gemma", init_fn="zeros"
+        LigerRMSNorm, offset=1.0, casting_mode="gemma", init_fn="zeros", in_place=False
     )
     _patch_rms_norm_module_for_gemma2 = partial(
         _patch_rms_norm_module, offset=1.0, casting_mode="gemma"

liger_kernel/transformers/rms_norm.py

@@ -6,7 +6,13 @@ from liger_kernel.ops.rms_norm import LigerRMSNormFunction
 
 class LigerRMSNorm(nn.Module):
     def __init__(
-        self, hidden_size, eps=1e-6, offset=0.0, casting_mode="llama", init_fn="ones"
+        self,
+        hidden_size,
+        eps=1e-6,
+        offset=0.0,
+        casting_mode="llama",
+        init_fn="ones",
+        in_place=True,
     ):
         super().__init__()
         assert init_fn in [
@@ -16,10 +22,11 @@ class LigerRMSNorm(nn.Module):
         self.weight = nn.Parameter(
             torch.ones(hidden_size) if init_fn == "ones" else torch.zeros(hidden_size)
         )
-        self.variance_epsilon, self.offset, self.casting_mode = (
+        self.variance_epsilon, self.offset, self.casting_mode, self.in_place = (
             eps,
             offset,
             casting_mode,
+            in_place,
         )
 
     def forward(self, hidden_states):
@@ -29,7 +36,8 @@ class LigerRMSNorm(nn.Module):
             self.variance_epsilon,
             self.offset,
             self.casting_mode,
+            self.in_place,
         )
 
     def extra_repr(self):
-        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}, offset={self.offset}"
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}, offset={self.offset}, in_place={self.in_place}"

{liger_kernel_nightly-0.4.0.dev20241112204954.dist-info → liger_kernel_nightly-0.4.1.dev20241112233746.dist-info}/METADATA

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

{liger_kernel_nightly-0.4.0.dev20241112204954.dist-info → liger_kernel_nightly-0.4.1.dev20241112233746.dist-info}/RECORD

@@ -8,7 +8,7 @@ liger_kernel/ops/group_norm.py,sha256=VaRErVJGR4JqgXXvuIjNGTn3E2egjLtU1y3ymwIf4d
 liger_kernel/ops/jsd.py,sha256=anWfdioucxZy4JQfTvbHBR-IQrZKeH-gBF1MHwwTuTQ,5781
 liger_kernel/ops/kl_div.py,sha256=03FNXfvCb6M-56hhFepAFV9p6brArPR6KOKkdGD34mw,8374
 liger_kernel/ops/layer_norm.py,sha256=unGMYMOPqtkM9aTrokhcqgPmsV2AUN7Yzv86isVB9OI,7422
-liger_kernel/ops/rms_norm.py,sha256=9S9wyZLmzNyJlBxV4vbv4p5es7bGP-m_5wK9JC6JIdA,10911
+liger_kernel/ops/rms_norm.py,sha256=LAxCiFjpBbb7TDh9pOzsVmDGAR7eEbTDnEhjSd6TX_M,11583
 liger_kernel/ops/rope.py,sha256=jrzaA9-6Orn44y_IIam9_YNPQxOFK2FrIRNfFea4EtU,8513
 liger_kernel/ops/swiglu.py,sha256=Fwxtd76rhHKT9ShQAGca9RsnASplAVxtYKHmiT73_yA,2994
 liger_kernel/ops/utils.py,sha256=3JSF--O7KT5Wa5BuO70M4h0XetxoZ_e9IoW9GRlxlBg,3777
@@ -25,8 +25,8 @@ liger_kernel/transformers/group_norm.py,sha256=FJ9R7mS9G1wO-GRIQ6QKSmIhnZ6nQ6GIk
 liger_kernel/transformers/jsd.py,sha256=W-5CypO2mx4-bUWOxq1KScfCdoXlLoYbtt5xBnRzMs4,3056
 liger_kernel/transformers/kl_div.py,sha256=qVhjBg6tjRyue5iZ3NFxo8uySY4JuIFJyv0IM_50F24,431
 liger_kernel/transformers/layer_norm.py,sha256=fd6o4kSHJWolQMWxh-l1qObfgL08ruNbUoBiANKX1ow,972
-liger_kernel/transformers/monkey_patch.py,sha256=RSnHwStyOi5C4xBbHrKvo32X2hUU7InjSyBzb4iu9T0,38184
-liger_kernel/transformers/rms_norm.py,sha256=4XfMQI6dORF7s_5qUqVHKWv-3IUomaimU2dg-NwnpoM,1035
+liger_kernel/transformers/monkey_patch.py,sha256=rhLJlc-Dpdo7kVYssjw5sEEbPuG18ZxIIxKf8AgrLjg,38200
+liger_kernel/transformers/rms_norm.py,sha256=AHstklNIO1PLHjjCBU-TPuUD-Fl_pycJUTLlJNojbV8,1189
 liger_kernel/transformers/rope.py,sha256=m-ah8vZBYW8tfplTXCiAPMHJWlB1tdp_JPXJeWE-Boo,943
 liger_kernel/transformers/swiglu.py,sha256=0-tVJ8xEYfhxnduc16PflXFj8sZPxdx9sHUn3hfwCI4,2468
 liger_kernel/transformers/trainer_integration.py,sha256=W3ON51O5GkyzNJsItz0y5rKx-uy2f2cFfveZpqbUdhw,123
@@ -43,9 +43,9 @@ liger_kernel/transformers/model/qwen2.py,sha256=EyhSSzQOskGjSnCsKMZpd1s5IAIlHd5P
 liger_kernel/transformers/model/qwen2_vl.py,sha256=j6xAhp9AG195dsZK5f8dFYVM9uKtWApZrggT5Y08jn4,7055
 liger_kernel/triton/__init__.py,sha256=yfRe0zMb47QnqjecZWG7LnanfCTzeku7SgWRAwNVmzU,101
 liger_kernel/triton/monkey_patch.py,sha256=5BcGKTtdqeYchypBIBopGIWPx1-cFALz7sOKoEsqXJ0,1584
-liger_kernel_nightly-0.4.0.dev20241112204954.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
-liger_kernel_nightly-0.4.0.dev20241112204954.dist-info/METADATA,sha256=oxSZh5RfYz52ri89yD3KJGjE1lX_Y1d_Dvi-j5aq12Y,28096
-liger_kernel_nightly-0.4.0.dev20241112204954.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
-liger_kernel_nightly-0.4.0.dev20241112204954.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-liger_kernel_nightly-0.4.0.dev20241112204954.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
-liger_kernel_nightly-0.4.0.dev20241112204954.dist-info/RECORD,,
+liger_kernel_nightly-0.4.1.dev20241112233746.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
+liger_kernel_nightly-0.4.1.dev20241112233746.dist-info/METADATA,sha256=OIhhLRs8Y5yM27YaSiYDrdh_56NE4pqHdax-sVuy8n8,28096
+liger_kernel_nightly-0.4.1.dev20241112233746.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
+liger_kernel_nightly-0.4.1.dev20241112233746.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+liger_kernel_nightly-0.4.1.dev20241112233746.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
+liger_kernel_nightly-0.4.1.dev20241112233746.dist-info/RECORD,,