RMS Layernorm

unsloth/kernels/__init__.py

@@ -13,7 +13,11 @@
 # limitations under the License.
 
 from .cross_entropy_loss import fast_cross_entropy_loss
-from .rms_layernorm import fast_rms_layernorm
+from .rms_layernorm import (
+    fast_rms_layernorm,
+    patch_rms_layernorm,
+    unpatch_rms_layernorm,
+)
 from .layernorm import (
     fast_layernorm,
     patch_layernorm,

unsloth/kernels/rms_layernorm.py

@@ -192,3 +192,25 @@ def fast_rms_layernorm(layernorm, X, gemma = False):
     out = Fast_RMS_Layernorm.apply(X, W, eps, gemma)
     return out
 pass
+
+
+from transformers.models.llama.modeling_llama import LlamaRMSNorm
+class Unsloth_LlamaRMSNorm(LlamaRMSNorm):
+    def forward(self, X):
+        return fast_rms_layernorm(self, X, gemma = False)
+    pass
+pass
+
+
+def patch_rms_layernorm():
+    import transformers.models.llama.modeling_llama
+    transformers.models.llama.modeling_llama.LlamaRMSNorm = Unsloth_LlamaRMSNorm
+    return
+pass
+
+
+def unpatch_rms_layernorm():
+    import transformers.models.llama.modeling_llama
+    transformers.models.llama.modeling_llama.LlamaRMSNorm = LlamaRMSNorm
+    return
+pass