Fix global LayerNorm mutation in model patching

src/liger_kernel/transformers/monkey_patch.py

@@ -3,6 +3,7 @@
 
 from functools import partial
 from types import MethodType
+from types import ModuleType
 from typing import Callable
 from typing import Optional
 
@@ -133,6 +134,13 @@ def _patch_layer_norm_module(module, eps=1e-6):
         _bind_method_to_module(module, "_get_name", lambda self: LigerLayerNorm.__name__)
 
 
+def _patch_layer_norm_class(modeling_module):
+    nn_namespace = ModuleType(modeling_module.nn.__name__)
+    nn_namespace.__dict__.update(modeling_module.nn.__dict__)
+    nn_namespace.LayerNorm = LigerLayerNorm
+    modeling_module.nn = nn_namespace
+
+
 def _patch_swiglu_module(module, liger_module):
     _bind_method_to_module(module, "forward", liger_module.forward)
     _bind_method_to_module(module, "_get_name", lambda self: liger_module.__name__)
@@ -561,7 +569,7 @@ def apply_liger_kernel_to_mllama(
     if rope:
         modeling_mllama.apply_rotary_pos_emb = liger_rotary_pos_emb
     if layer_norm and model is None:
-        modeling_mllama.nn.LayerNorm = LigerLayerNorm
+        _patch_layer_norm_class(modeling_mllama)
     if rms_norm:
         modeling_mllama.MllamaTextRMSNorm = LigerRMSNorm
     if swiglu:
@@ -1190,7 +1198,7 @@ def apply_liger_kernel_to_gemma3(
     )
 
     if layer_norm and model is None:
-        modeling_siglip.nn.LayerNorm = LigerLayerNorm
+        _patch_layer_norm_class(modeling_siglip)
 
     apply_liger_kernel_to_gemma3_text(
         rope=rope, cross_entropy=False, fused_linear_cross_entropy=False, rms_norm=rms_norm, geglu=geglu
@@ -1529,7 +1537,7 @@ def apply_liger_kernel_to_paligemma(
 
     # The vision_tower is a SiglipVisionModel
     if layer_norm and model is None:
-        modeling_siglip.nn.LayerNorm = LigerLayerNorm
+        _patch_layer_norm_class(modeling_siglip)
 
     # SiglipMLP is standard FFN so LigerGEGLUMLP is not compatible
     # The multi_modal_projector is Linear, nothing to do
@@ -2706,12 +2714,12 @@ def apply_liger_kernel_to_internvl(
     from transformers.models.internvl.modeling_internvl import InternVLVisionModel
     from transformers.models.internvl.modeling_internvl import InternVLVisionRMSNorm
 
-    from liger_kernel.transformers.layer_norm import LigerLayerNorm
     from liger_kernel.transformers.model.internvl import lce_forward as internvl_lce_forward
     from liger_kernel.transformers.rms_norm import LigerRMSNorm
 
     if layer_norm and model is None:
-        modeling_internvl.nn.LayerNorm = LigerLayerNorm
+        _patch_layer_norm_class(modeling_internvl)
+        modeling_internvl.NORM2FN["layer_norm"] = LigerLayerNorm
 
     if cross_entropy:
         logger.info("Apply liger cross entropy")
@@ -2820,7 +2828,7 @@ def apply_liger_kernel_to_smolvlm(
 
     # Patch LayerNorm for vision model if model is not provided (pre-initialization)
     if layer_norm and model is None:
-        modeling_smolvlm.nn.LayerNorm = LigerLayerNorm
+        _patch_layer_norm_class(modeling_smolvlm)
 
     if cross_entropy:
         logger.info("Apply liger cross entropy")

test/transformers/test_monkey_patch.py

@@ -12,6 +12,7 @@
 from packaging import version
 from test.utils import get_mllama_rope_config
 from test.utils import get_qwen3_vl_rope_config
+from torch.nn import LayerNorm as TorchLayerNorm
 from transformers import AutoModelForCausalLM
 from transformers import PretrainedConfig
 from transformers import PreTrainedModel
@@ -473,6 +474,106 @@ def test_patching_apis_support_patching_model_instance():
         )
 
 
+@pytest.mark.parametrize(
+    "apply_fn_name,modeling_module_path,kwargs",
+    [
+        (
+            "apply_liger_kernel_to_gemma3",
+            "transformers.models.siglip.modeling_siglip",
+            {
+                "rope": False,
+                "cross_entropy": False,
+                "fused_linear_cross_entropy": False,
+                "rms_norm": False,
+                "geglu": False,
+            },
+        ),
+        (
+            "apply_liger_kernel_to_internvl",
+            "transformers.models.internvl.modeling_internvl",
+            {
+                "cross_entropy": False,
+                "fused_linear_cross_entropy": False,
+                "rms_norm": False,
+            },
+        ),
+        (
+            "apply_liger_kernel_to_mllama",
+            "transformers.models.mllama.modeling_mllama",
+            {
+                "rope": False,
+                "cross_entropy": False,
+                "fused_linear_cross_entropy": False,
+                "rms_norm": False,
+                "swiglu": False,
+            },
+        ),
+        (
+            "apply_liger_kernel_to_paligemma",
+            "transformers.models.siglip.modeling_siglip",
+            {
+                "rope": False,
+                "cross_entropy": False,
+                "fused_linear_cross_entropy": False,
+                "rms_norm": False,
+                "geglu": False,
+            },
+        ),
+        (
+            "apply_liger_kernel_to_smolvlm",
+            "transformers.models.smolvlm.modeling_smolvlm",
+            {
+                "cross_entropy": False,
+                "fused_linear_cross_entropy": False,
+                "rms_norm": False,
+            },
+        ),
+    ],
+)
+def test_class_level_layer_norm_patch_does_not_modify_torch_nn(apply_fn_name, modeling_module_path, kwargs):
+    modeling_module = pytest.importorskip(modeling_module_path)
+    apply_fn = getattr(monkey_patch, apply_fn_name)
+
+    with (
+        patch.object(torch.nn, "LayerNorm", TorchLayerNorm),
+        patch.object(modeling_module, "nn", torch.nn),
+    ):
+        apply_fn(layer_norm=True, **kwargs)
+
+        assert torch.nn.LayerNorm is TorchLayerNorm
+        assert modeling_module.nn is not torch.nn
+        assert modeling_module.nn.LayerNorm is LigerLayerNorm
+        assert modeling_module.nn.Linear is torch.nn.Linear
+
+
+def test_internvl_class_level_layer_norm_patch_updates_norm_factory():
+    modeling_internvl = pytest.importorskip("transformers.models.internvl.modeling_internvl")
+    from transformers.models.internvl.configuration_internvl import InternVLVisionConfig
+
+    with (
+        patch.object(torch.nn, "LayerNorm", TorchLayerNorm),
+        patch.object(modeling_internvl, "nn", torch.nn),
+        patch.dict(modeling_internvl.NORM2FN, {"layer_norm": TorchLayerNorm}),
+    ):
+        monkey_patch.apply_liger_kernel_to_internvl(
+            cross_entropy=False,
+            fused_linear_cross_entropy=False,
+            rms_norm=False,
+            layer_norm=True,
+        )
+        layer = modeling_internvl.InternVLVisionLayer(
+            InternVLVisionConfig(
+                hidden_size=32,
+                intermediate_size=64,
+                num_attention_heads=4,
+                norm_type="layer_norm",
+            )
+        )
+
+        assert isinstance(layer.layernorm_before, LigerLayerNorm)
+        assert isinstance(layer.layernorm_after, LigerLayerNorm)
+
+
 def test_apply_liger_kernel_to_instance_for_llama():
     # Ensure any monkey patching is cleaned up for subsequent tests
     with patch("transformers.models.llama.modeling_llama"):