linkedin/Liger-Kernel
1
0
Fork 0
mirror of https://github.com/linkedin/Liger-Kernel.git synced 2026-10-12 02:17:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Bug Fix: name patching for modules

Browse Source
This commit is contained in:
Vaibhav Jindal
2025-07-24 00:25:49 +00:00
parent 8cd359894f
commit 67a5439ebf
1 changed files with 12 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/liger_kernel/transformers/monkey_patch.py
+12 -12
View File
@@ -78,8 +78,8 @@ def _patch_rms_norm_module(module, offset=0.0, eps=1e-6, casting_mode="llama", i
_bind_method_to_module(module.modules_to_save.default, "extra_repr", LigerRMSNorm.extra_repr)
_bind_method_to_module(module.original_module, "forward", LigerRMSNorm.forward)
_bind_method_to_module(module.original_module, "extra_repr", LigerRMSNorm.extra_repr)
module.modules_to_save.default.__class__.__name__ = LigerRMSNorm.__name__
module.original_module.__class__.__name__ = LigerRMSNorm.__name__
_bind_method_to_module(module.modules_to_save.default, "_get_name", lambda self: LigerRMSNorm.__name__)
_bind_method_to_module(module.original_module, "_get_name", lambda self: LigerRMSNorm.__name__)
else:
module.offset = offset
module.casting_mode = casting_mode
@@ -88,7 +88,7 @@ def _patch_rms_norm_module(module, offset=0.0, eps=1e-6, casting_mode="llama", i
module.row_mode = row_mode
_bind_method_to_module(module, "forward", LigerRMSNorm.forward)
_bind_method_to_module(module, "extra_repr", LigerRMSNorm.extra_repr)
module.__class__.__name__ = LigerRMSNorm.__name__
_bind_method_to_module(module, "_get_name", lambda self: LigerRMSNorm.__name__)
def _patch_layer_norm_module(module, eps=1e-6):
@@ -110,28 +110,28 @@ def _patch_layer_norm_module(module, eps=1e-6):
module.original_module.hidden_size = getattr(module, "hidden_size", None) or getattr(
module, "normalized_shape", None
)
_bind_method_to_module(module.modules_to_save.default, "forward", LigerRMSNorm.forward)
_bind_method_to_module(module.modules_to_save.default, "extra_repr", LigerRMSNorm.extra_repr)
_bind_method_to_module(module.original_module, "forward", LigerRMSNorm.forward)
_bind_method_to_module(module.original_module, "extra_repr", LigerRMSNorm.extra_repr)
module.modules_to_save.default.__class__.__name__ = LigerLayerNorm.__name__
module.original_module.__class__.__name__ = LigerLayerNorm.__name__
_bind_method_to_module(module.modules_to_save.default, "forward", LigerLayerNorm.forward)
_bind_method_to_module(module.modules_to_save.default, "extra_repr", LigerLayerNorm.extra_repr)
_bind_method_to_module(module.original_module, "forward", LigerLayerNorm.forward)
_bind_method_to_module(module.original_module, "extra_repr", LigerLayerNorm.extra_repr)
_bind_method_to_module(module.modules_to_save.default, "_get_name", lambda self: LigerLayerNorm.__name__)
_bind_method_to_module(module.original_module, "_get_name", lambda self: LigerLayerNorm.__name__)
else:
module.variance_epsilon = getattr(module, "variance_epsilon", None) or getattr(module, "eps", None) or eps
module.hidden_size = getattr(module, "hidden_size", None) or getattr(module, "normalized_shape", None)
_bind_method_to_module(module, "forward", LigerLayerNorm.forward)
_bind_method_to_module(module, "extra_repr", LigerLayerNorm.extra_repr)
module.__class__.__name__ = LigerLayerNorm.__name__
_bind_method_to_module(module, "_get_name", lambda self: LigerLayerNorm.__name__)
def _patch_swiglu_module(module, liger_module):
_bind_method_to_module(module, "forward", liger_module.forward)
module.__class__.__name__ = liger_module.__name__
_bind_method_to_module(module, "_get_name", lambda self: liger_module.__name__)
def _patch_geglu_module(module):
_bind_method_to_module(module, "forward", LigerGEGLUMLP.forward)
module.__class__.__name__ = LigerGEGLUMLP.__name__
_bind_method_to_module(module, "_get_name", lambda self: LigerGEGLUMLP.__name__)
def apply_liger_kernel_to_granite(