From 815700b5fd9d660630a086a4271a870ee68db44c Mon Sep 17 00:00:00 2001 From: Vedavyas <128640011+vedavyas2727@users.noreply.github.com> Date: Mon, 10 Aug 2026 00:35:30 -0700 Subject: [PATCH] fix: avoid global LayerNorm monkey patches Clone model-local torch.nn namespaces before replacing LayerNorm, and update InternVL's cached norm factory so new model instances still use LigerLayerNorm without mutating torch.nn process-wide. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/liger_kernel/transformers/monkey_patch.py | 20 ++-- test/transformers/test_monkey_patch.py | 101 ++++++++++++++++++ 2 files changed, 115 insertions(+), 6 deletions(-) diff --git a/src/liger_kernel/transformers/monkey_patch.py b/src/liger_kernel/transformers/monkey_patch.py index 4d33d7e41..3109bb3c8 100755 --- a/src/liger_kernel/transformers/monkey_patch.py +++ b/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") diff --git a/test/transformers/test_monkey_patch.py b/test/transformers/test_monkey_patch.py index 25099f7be..77619c652 100755 --- a/test/transformers/test_monkey_patch.py +++ b/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"):