liger-kernel-nightly 0.5.10.dev20250602014906__tar.gz → 0.5.10.dev20250602134913__tar.gz

{liger_kernel_nightly-0.5.10.dev20250602014906 → liger_kernel_nightly-0.5.10.dev20250602134913}/PKG-INFO

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

{liger_kernel_nightly-0.5.10.dev20250602014906 → liger_kernel_nightly-0.5.10.dev20250602134913}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.5.10.dev20250602014906"
+version = "0.5.10.dev20250602134913"
 description = "Efficient Triton kernels for LLM Training"
 urls = { "Homepage" = "https://github.com/linkedin/Liger-Kernel" }
 readme = { file = "README.md", content-type = "text/markdown" }

{liger_kernel_nightly-0.5.10.dev20250602014906 → liger_kernel_nightly-0.5.10.dev20250602134913}/src/liger_kernel/transformers/monkey_patch.py

@@ -627,8 +627,8 @@ def apply_liger_kernel_to_gemma(
     from transformers.models.gemma import modeling_gemma
     from transformers.models.gemma.modeling_gemma import GemmaModel
 
-    # https://github.com/huggingface/transformers/blob/v4.44.2/src/transformers/models/gemma/modeling_gemma.py#L109
-    LigerRMSNormForGemma = partial(LigerRMSNorm, offset=1.0, init_fn="zeros", casting_mode="gemma")
+    from liger_kernel.transformers.rms_norm import LigerRMSNormForGemma
+
     _patch_rms_norm_module_for_gemma = partial(_patch_rms_norm_module, casting_mode="gemma", offset=1.0)
 
     if rope:
@@ -701,7 +701,8 @@ def apply_liger_kernel_to_gemma2(
     from transformers.models.gemma2 import modeling_gemma2
     from transformers.models.gemma2.modeling_gemma2 import Gemma2Model
 
-    LigerRMSNormForGemma2 = partial(LigerRMSNorm, offset=1.0, casting_mode="gemma", init_fn="zeros", in_place=False)
+    from liger_kernel.transformers.rms_norm import LigerRMSNormForGemma2
+
     _patch_rms_norm_module_for_gemma2 = partial(
         _patch_rms_norm_module, offset=1.0, casting_mode="gemma", in_place=False
     )
@@ -780,8 +781,8 @@ def apply_liger_kernel_to_gemma3_text(
     from transformers.models.gemma3.modeling_gemma3 import Gemma3ForCausalLM
     from transformers.models.gemma3.modeling_gemma3 import Gemma3TextModel
 
-    from liger_kernel.transformers.gema3_rms import LigerRMSNormForGemma3
     from liger_kernel.transformers.model.gemma3 import causal_forward
+    from liger_kernel.transformers.rms_norm import LigerRMSNormForGemma3
 
     _patch_rms_norm_module_for_gemma3 = partial(
         _patch_rms_norm_module, offset=1.0, casting_mode="gemma", in_place=False
@@ -1451,11 +1452,12 @@ def apply_liger_kernel_to_olmo2(
     from transformers.models.olmo2.modeling_olmo2 import Olmo2Model
 
     from liger_kernel.transformers.model.olmo2 import lce_forward as olmo2_lce_forward
+    from liger_kernel.transformers.rms_norm import LigerRMSNormForOlmo2
 
     if rope:
         modeling_olmo2.apply_rotary_pos_emb = liger_rotary_pos_emb
     if rms_norm:
-        modeling_olmo2.Olmo2RMSNorm = partial(LigerRMSNorm, in_place=False)
+        modeling_olmo2.Olmo2RMSNorm = LigerRMSNormForOlmo2
     if swiglu:
         modeling_olmo2.Olmo2MLP = LigerSwiGLUMLP
     if cross_entropy:
@@ -1514,11 +1516,12 @@ def apply_liger_kernel_to_glm4(
     from transformers.models.glm4.modeling_glm4 import Glm4Model
 
     from liger_kernel.transformers.model.glm4 import lce_forward as glm4_lce_forward
+    from liger_kernel.transformers.rms_norm import LigerRMSNormForGlm4
 
     if rope:
         raise NotImplementedError("liger_rotary_pos_emb is not available for Glm4 models.")
     if rms_norm:
-        modeling_glm4.Glm4RMSNorm = partial(LigerRMSNorm, in_place=False)
+        modeling_glm4.Glm4RMSNorm = LigerRMSNormForGlm4
     if swiglu:
         modeling_glm4.Glm4MLP = LigerPhi3SwiGLUMLP
     if cross_entropy:

liger_kernel_nightly-0.5.10.dev20250602134913/src/liger_kernel/transformers/rms_norm.py

@@ -0,0 +1,81 @@
+import torch
+import torch.nn as nn
+
+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",
+        in_place=True,
+        row_mode=None,
+    ):
+        super().__init__()
+        assert init_fn in [
+            "ones",
+            "zeros",
+        ], f"init_fn must be either 'ones' or 'zeros', got {init_fn}"
+        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.in_place, self.row_mode = (
+            eps,
+            offset,
+            casting_mode,
+            in_place,
+            row_mode,
+        )
+
+    def forward(self, hidden_states):
+        return LigerRMSNormFunction.apply(
+            hidden_states,
+            self.weight,
+            self.variance_epsilon,
+            self.offset,
+            self.casting_mode,
+            self.in_place,
+            self.row_mode,
+        )
+
+    def extra_repr(self):
+        return (
+            f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}, offset={self.offset}, in_place={self.in_place}"
+        )
+
+
+class LigerRMSNormForGemma(LigerRMSNorm):
+    def __init__(
+        self, hidden_size, eps=1e-6, offset=1.0, casting_mode="gemma", init_fn="zeros", in_place=True, row_mode=None
+    ):
+        super().__init__(hidden_size, eps, offset, casting_mode, init_fn, in_place, row_mode)
+
+
+class LigerRMSNormForGemma2(LigerRMSNorm):
+    def __init__(
+        self, hidden_size, eps=1e-6, offset=1.0, casting_mode="gemma", init_fn="zeros", in_place=False, row_mode=None
+    ):
+        super().__init__(hidden_size, eps, offset, casting_mode, init_fn, in_place, row_mode)
+
+
+class LigerRMSNormForGemma3(LigerRMSNorm):
+    """Gemma3RMSNorm has a dim argument not hidden_size used in q_norm and k_norm."""
+
+    def __init__(self, dim, eps=0.000001, offset=1.0, casting_mode="gemma", init_fn="zeros", in_place=False):
+        super().__init__(dim, eps, offset, casting_mode, init_fn, in_place)
+
+
+class LigerRMSNormForOlmo2(LigerRMSNorm):
+    def __init__(
+        self, hidden_size, eps=1e-6, offset=0.0, casting_mode="llama", init_fn="ones", in_place=False, row_mode=None
+    ):
+        super().__init__(hidden_size, eps, offset, casting_mode, init_fn, in_place, row_mode)
+
+
+class LigerRMSNormForGlm4(LigerRMSNorm):
+    def __init__(
+        self, hidden_size, eps=1e-6, offset=0.0, casting_mode="llama", init_fn="ones", in_place=False, row_mode=None
+    ):
+        super().__init__(hidden_size, eps, offset, casting_mode, init_fn, in_place, row_mode)

{liger_kernel_nightly-0.5.10.dev20250602014906 → liger_kernel_nightly-0.5.10.dev20250602134913}/src/liger_kernel_nightly.egg-info/PKG-INFO

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

{liger_kernel_nightly-0.5.10.dev20250602014906 → liger_kernel_nightly-0.5.10.dev20250602134913}/src/liger_kernel_nightly.egg-info/SOURCES.txt

@@ -156,7 +156,6 @@ src/liger_kernel/transformers/functional.py
 src/liger_kernel/transformers/fused_linear_cross_entropy.py
 src/liger_kernel/transformers/fused_linear_jsd.py
 src/liger_kernel/transformers/geglu.py
-src/liger_kernel/transformers/gema3_rms.py
 src/liger_kernel/transformers/group_norm.py
 src/liger_kernel/transformers/grpo_loss.py
 src/liger_kernel/transformers/jsd.py

liger_kernel_nightly-0.5.10.dev20250602014906/src/liger_kernel/transformers/gema3_rms.py

@@ -1,8 +0,0 @@
-from .rms_norm import LigerRMSNorm
-
-
-class LigerRMSNormForGemma3(LigerRMSNorm):
-    """Gemma3RMSNorm has a dim argument not hidden_size used in q_norm and k_norm."""
-
-    def __init__(self, dim, eps=0.000001, offset=1.0, casting_mode="gemma", init_fn="zeros", in_place=False):
-        super().__init__(dim, eps, offset, casting_mode, init_fn, in_place)

liger_kernel_nightly-0.5.10.dev20250602014906/src/liger_kernel/transformers/rms_norm.py

@@ -1,46 +0,0 @@
-import torch
-import torch.nn as nn
-
-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",
-        in_place=True,
-        row_mode=None,
-    ):
-        super().__init__()
-        assert init_fn in [
-            "ones",
-            "zeros",
-        ], f"init_fn must be either 'ones' or 'zeros', got {init_fn}"
-        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.in_place, self.row_mode = (
-            eps,
-            offset,
-            casting_mode,
-            in_place,
-            row_mode,
-        )
-
-    def forward(self, hidden_states):
-        return LigerRMSNormFunction.apply(
-            hidden_states,
-            self.weight,
-            self.variance_epsilon,
-            self.offset,
-            self.casting_mode,
-            self.in_place,
-            self.row_mode,
-        )
-
-    def extra_repr(self):
-        return (
-            f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}, offset={self.offset}, in_place={self.in_place}"
-        )