From a84b1dffc4a166c718048378fd2d35b76ab03635 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 24 Sep 2024 17:22:20 -0700
Subject: [PATCH] Patch layernorm

---
 unsloth/kernels/__init__.py  |  6 +++++-
 unsloth/kernels/layernorm.py | 23 +++++++++++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/unsloth/kernels/__init__.py b/unsloth/kernels/__init__.py
index 2fb9d11f99..841d0ce0f0 100644
--- a/unsloth/kernels/__init__.py
+++ b/unsloth/kernels/__init__.py
@@ -14,7 +14,11 @@
 
 from .cross_entropy_loss import fast_cross_entropy_loss
 from .rms_layernorm import fast_rms_layernorm
-from .layernorm import fast_layernorm
+from .layernorm import (
+    fast_layernorm,
+    patch_layernorm,
+    unpatch_layernorm,
+)
 from .rope_embedding import fast_rope_embedding, inplace_rope_embedding
 from .swiglu import swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel
 from .geglu import (
diff --git a/unsloth/kernels/layernorm.py b/unsloth/kernels/layernorm.py
index c0ff4d7440..0d456109ea 100644
--- a/unsloth/kernels/layernorm.py
+++ b/unsloth/kernels/layernorm.py
@@ -150,6 +150,7 @@ pass
 
 
 def fast_layernorm(layernorm, X):
+    assert(layernorm.elementwise_affine is True)
     W    = layernorm.weight
     bias = layernorm.bias
     eps = layernorm.variance_epsilon if \
@@ -160,6 +161,28 @@ def fast_layernorm(layernorm, X):
 pass
 
 
+from torch.nn import LayerNorm
+class Fast_LayerNorm_Module(LayerNorm):
+    def forward(self, X):
+        return fast_layernorm(self, X)
+    pass
+pass
+
+
+def patch_layernorm():
+    import torch.nn
+    torch.nn.LayerNorm = Fast_LayerNorm_Module
+    return
+pass
+
+
+def unpatch_layernorm():
+    import torch.nn
+    torch.nn.LayerNorm = LayerNorm
+    return
+pass
+
+
 def test_layernorm(
     dim = 1024, eps = 1e-5, dtype = torch.float16,
     bsz = 21, random_state = 3407, seqlen = 3341,