Patch layernorm

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 (

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,