diff --git a/src/liger_kernel/transformers/monkey_patch.py b/src/liger_kernel/transformers/monkey_patch.py index dd8cf3a8..758e7f44 100755 --- a/src/liger_kernel/transformers/monkey_patch.py +++ b/src/liger_kernel/transformers/monkey_patch.py @@ -1871,6 +1871,8 @@ def apply_liger_kernel_to_qwen3( if rms_norm: _patch_rms_norm_module(decoder_layer.input_layernorm) _patch_rms_norm_module(decoder_layer.post_attention_layernorm) + _patch_rms_norm_module(decoder_layer.self_attn.q_norm) + _patch_rms_norm_module(decoder_layer.self_attn.k_norm) def apply_liger_kernel_to_qwen3_moe( @@ -1936,6 +1938,8 @@ def apply_liger_kernel_to_qwen3_moe( if rms_norm: _patch_rms_norm_module(decoder_layer.input_layernorm) _patch_rms_norm_module(decoder_layer.post_attention_layernorm) + _patch_rms_norm_module(decoder_layer.self_attn.q_norm) + _patch_rms_norm_module(decoder_layer.self_attn.k_norm) def apply_liger_kernel_to_gpt_oss( @@ -3194,6 +3198,9 @@ def apply_liger_kernel_to_qwen3_next( if rms_norm: _patch_rms_norm_module_for_qwen3_next(decoder_layer.input_layernorm) _patch_rms_norm_module_for_qwen3_next(decoder_layer.post_attention_layernorm) + if hasattr(decoder_layer, "self_attn"): + _patch_rms_norm_module_for_qwen3_next(decoder_layer.self_attn.q_norm) + _patch_rms_norm_module_for_qwen3_next(decoder_layer.self_attn.k_norm) # Qwen3MoeMLP and Qwen3NextMLP are identical, hence we reuse LigerQwen3MoeSwiGLUMLP if swiglu: @@ -3302,6 +3309,9 @@ def apply_liger_kernel_to_qwen3_5( if rms_norm: _patch_rms_norm_module_for_qwen3_5(decoder_layer.input_layernorm) _patch_rms_norm_module_for_qwen3_5(decoder_layer.post_attention_layernorm) + if hasattr(decoder_layer, "self_attn"): + _patch_rms_norm_module_for_qwen3_5(decoder_layer.self_attn.q_norm) + _patch_rms_norm_module_for_qwen3_5(decoder_layer.self_attn.k_norm) if swiglu: _patch_swiglu_module(decoder_layer.mlp, LigerQwen3MoeSwiGLUMLP) @@ -3402,6 +3412,9 @@ def apply_liger_kernel_to_qwen3_5_moe( if rms_norm: _patch_rms_norm_module_for_qwen3_5_moe(decoder_layer.input_layernorm) _patch_rms_norm_module_for_qwen3_5_moe(decoder_layer.post_attention_layernorm) + if hasattr(decoder_layer, "self_attn"): + _patch_rms_norm_module_for_qwen3_5_moe(decoder_layer.self_attn.q_norm) + _patch_rms_norm_module_for_qwen3_5_moe(decoder_layer.self_attn.k_norm) if swiglu: _patch_swiglu_module(decoder_layer.mlp.shared_expert, LigerQwen3MoeSwiGLUMLP) diff --git a/test/transformers/test_monkey_patch.py b/test/transformers/test_monkey_patch.py index cf07478a..ae59729b 100755 --- a/test/transformers/test_monkey_patch.py +++ b/test/transformers/test_monkey_patch.py @@ -2526,10 +2526,49 @@ def test_apply_liger_kernel_to_instance_for_qwen2(): pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}") +def _check_qwen_qk_norm_patch(model, text_model, gemma=False, **kwargs): + qk_norms = [ + norm + for layer in text_model.layers + if hasattr(layer, "self_attn") + for norm in (layer.self_attn.q_norm, layer.self_attn.k_norm) + ] + assert qk_norms, "The test must include full attention Q/K norms" + linear_norms = [layer.linear_attn.norm for layer in text_model.layers if hasattr(layer, "linear_attn")] + if gemma: + assert linear_norms, "The hybrid model test must also include linear attention" + linear_forwards = [norm.forward for norm in linear_norms] + original = [ + (norm.weight, norm.weight.detach().clone(), norm.eps if gemma else norm.variance_epsilon) for norm in qk_norms + ] + native_forwards = [norm.forward for norm in qk_norms] + for norm in qk_norms: + assert inspect.getsource(norm.forward) != inspect.getsource(LigerRMSNorm.forward) + + _apply_liger_kernel_to_instance( + model=model, rms_norm=False, rope=False, swiglu=False, fused_linear_cross_entropy=False + ) + assert [norm.forward for norm in qk_norms] == native_forwards + + _apply_liger_kernel_to_instance(model=model, **kwargs) + for norm, (weight, values, epsilon) in zip(qk_norms, original): + assert inspect.getsource(norm.forward) == inspect.getsource(LigerRMSNorm.forward) + assert norm.weight is weight + assert torch.equal(norm.weight, values) + assert norm.variance_epsilon == epsilon + assert norm.offset == (1.0 if gemma else 0.0) + assert norm.casting_mode == ("gemma" if gemma else "llama") + assert norm.in_place is not gemma + assert [norm.forward for norm in linear_norms] == linear_forwards + + @pytest.mark.skipif(not is_qwen3_available(), reason="qwen3 module not available") -def test_apply_liger_kernel_to_instance_for_qwen3(): +@pytest.mark.parametrize("base_model", [False, True]) +def test_apply_liger_kernel_to_instance_for_qwen3(base_model): # Ensure any monkey patching is cleaned up for subsequent tests with patch("transformers.models.qwen3.modeling_qwen3"): + from transformers.models.qwen3.modeling_qwen3 import Qwen3Model + from liger_kernel.transformers.model.qwen3 import lce_forward as qwen3_lce_forward # Instantiate a dummy model @@ -2541,23 +2580,29 @@ def test_apply_liger_kernel_to_instance_for_qwen3(): hidden_act="silu", num_hidden_layers=2, ) - dummy_model_instance = AutoModelForCausalLM.from_config(config) + if base_model: + dummy_model_instance = Qwen3Model._from_config(config) + else: + dummy_model_instance = AutoModelForCausalLM.from_config(config) + text_model = dummy_model_instance if base_model else dummy_model_instance.model # Check that model instance variables are not yet patched with Liger modules - assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(qwen3_lce_forward) - assert inspect.getsource(dummy_model_instance.model.norm.forward) != inspect.getsource(LigerRMSNorm.forward) - for layer in dummy_model_instance.model.layers: + if not base_model: + assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(qwen3_lce_forward) + assert inspect.getsource(text_model.norm.forward) != inspect.getsource(LigerRMSNorm.forward) + for layer in text_model.layers: assert inspect.getsource(layer.mlp.forward) != inspect.getsource(LigerSwiGLUMLP.forward) assert inspect.getsource(layer.input_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward) assert inspect.getsource(layer.post_attention_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward) # Test applying kernels to the model instance - _apply_liger_kernel_to_instance(model=dummy_model_instance) + _check_qwen_qk_norm_patch(dummy_model_instance, text_model, fused_linear_cross_entropy=not base_model) # Check that the model's instance variables were correctly patched with Liger modules - assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(qwen3_lce_forward) - assert inspect.getsource(dummy_model_instance.model.norm.forward) == inspect.getsource(LigerRMSNorm.forward) - for layer in dummy_model_instance.model.layers: + if not base_model: + assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(qwen3_lce_forward) + assert inspect.getsource(text_model.norm.forward) == inspect.getsource(LigerRMSNorm.forward) + for layer in text_model.layers: assert inspect.getsource(layer.mlp.forward) == inspect.getsource(LigerSwiGLUMLP.forward) assert inspect.getsource(layer.input_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward) assert inspect.getsource(layer.post_attention_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward) @@ -2569,9 +2614,12 @@ def test_apply_liger_kernel_to_instance_for_qwen3(): @pytest.mark.skipif(not is_qwen3_available(), reason="qwen3 module not available") -def test_apply_liger_kernel_to_instance_for_qwen3_moe(): +@pytest.mark.parametrize("base_model", [False, True]) +def test_apply_liger_kernel_to_instance_for_qwen3_moe(base_model): # Ensure any monkey patching is cleaned up for subsequent tests with patch("transformers.models.qwen3_moe.modeling_qwen3_moe"): + from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeModel + from liger_kernel.transformers.model.qwen3_moe import lce_forward as qwen3_moe_lce_forward # Instantiate a dummy model @@ -2583,12 +2631,17 @@ def test_apply_liger_kernel_to_instance_for_qwen3_moe(): hidden_act="silu", num_hidden_layers=2, ) - dummy_model_instance = AutoModelForCausalLM.from_config(config) + if base_model: + dummy_model_instance = Qwen3MoeModel._from_config(config) + else: + dummy_model_instance = AutoModelForCausalLM.from_config(config) + text_model = dummy_model_instance if base_model else dummy_model_instance.model # Check that model instance variables are not yet patched with Liger modules - assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(qwen3_moe_lce_forward) - assert inspect.getsource(dummy_model_instance.model.norm.forward) != inspect.getsource(LigerRMSNorm.forward) - for layer in dummy_model_instance.model.layers: + if not base_model: + assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(qwen3_moe_lce_forward) + assert inspect.getsource(text_model.norm.forward) != inspect.getsource(LigerRMSNorm.forward) + for layer in text_model.layers: if IS_TRANSFORMERS_V5_OR_LATER: assert inspect.getsource(layer.mlp.experts.forward) != inspect.getsource(LigerExperts.forward) else: @@ -2598,12 +2651,13 @@ def test_apply_liger_kernel_to_instance_for_qwen3_moe(): assert inspect.getsource(layer.post_attention_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward) # Test applying kernels to the model instance - _apply_liger_kernel_to_instance(model=dummy_model_instance) + _check_qwen_qk_norm_patch(dummy_model_instance, text_model, fused_linear_cross_entropy=not base_model) # Check that the model's instance variables were correctly patched with Liger modules - assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(qwen3_moe_lce_forward) - assert inspect.getsource(dummy_model_instance.model.norm.forward) == inspect.getsource(LigerRMSNorm.forward) - for layer in dummy_model_instance.model.layers: + if not base_model: + assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(qwen3_moe_lce_forward) + assert inspect.getsource(text_model.norm.forward) == inspect.getsource(LigerRMSNorm.forward) + for layer in text_model.layers: if IS_TRANSFORMERS_V5_OR_LATER: assert inspect.getsource(layer.mlp.experts.forward) == inspect.getsource(LigerExperts.forward) else: @@ -3588,9 +3642,12 @@ def test_apply_liger_kernel_to_instance_for_smollm3(): @pytest.mark.skipif(not is_qwen3_next_available(), reason="qwen3_next module not available") -def test_apply_liger_kernel_to_instance_for_qwen3_next(): +@pytest.mark.parametrize("base_model", [False, True]) +def test_apply_liger_kernel_to_instance_for_qwen3_next(base_model): # Ensure any monkey patching is cleaned up for subsequent tests with patch("transformers.models.qwen3_next.modeling_qwen3_next"): + from transformers.models.qwen3_next.modeling_qwen3_next import Qwen3NextModel + # Instantiate a dummy model config = transformers.models.qwen3_next.configuration_qwen3_next.Qwen3NextConfig( dtype=torch.bfloat16, @@ -3601,16 +3658,22 @@ def test_apply_liger_kernel_to_instance_for_qwen3_next(): shared_expert_intermediate_size=16, hidden_act="silu", num_hidden_layers=2, + layer_types=["linear_attention", "full_attention"], num_experts=2, num_experts_per_tok=1, mlp_only_layers=[1], ) - dummy_model_instance = AutoModelForCausalLM.from_config(config) + if base_model: + dummy_model_instance = Qwen3NextModel._from_config(config) + else: + dummy_model_instance = AutoModelForCausalLM.from_config(config) + text_model = dummy_model_instance if base_model else dummy_model_instance.model # Check that model instance variables are not yet patched with Liger modules - assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(qwen3_next_lce_forward) - assert inspect.getsource(dummy_model_instance.model.norm.forward) != inspect.getsource(LigerRMSNorm.forward) - for layer in dummy_model_instance.model.layers: + if not base_model: + assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(qwen3_next_lce_forward) + assert inspect.getsource(text_model.norm.forward) != inspect.getsource(LigerRMSNorm.forward) + for layer in text_model.layers: if hasattr(layer, "mlp") and hasattr(layer.mlp, "experts"): if IS_TRANSFORMERS_V5_OR_LATER: assert inspect.getsource(layer.mlp.experts.forward) != inspect.getsource(LigerExperts.forward) @@ -3628,12 +3691,15 @@ def test_apply_liger_kernel_to_instance_for_qwen3_next(): assert inspect.getsource(layer.post_attention_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward) # Test applying kernels to the model instance - _apply_liger_kernel_to_instance(model=dummy_model_instance) + _check_qwen_qk_norm_patch( + dummy_model_instance, text_model, gemma=True, fused_linear_cross_entropy=not base_model + ) # Check that the model's instance variables were correctly patched with Liger modules - assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(qwen3_next_lce_forward) - assert inspect.getsource(dummy_model_instance.model.norm.forward) == inspect.getsource(LigerRMSNorm.forward) - for layer in dummy_model_instance.model.layers: + if not base_model: + assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(qwen3_next_lce_forward) + assert inspect.getsource(text_model.norm.forward) == inspect.getsource(LigerRMSNorm.forward) + for layer in text_model.layers: if hasattr(layer, "mlp") and hasattr(layer.mlp, "experts"): if IS_TRANSFORMERS_V5_OR_LATER: assert inspect.getsource(layer.mlp.experts.forward) == inspect.getsource(LigerExperts.forward) @@ -3657,9 +3723,12 @@ def test_apply_liger_kernel_to_instance_for_qwen3_next(): @pytest.mark.skipif(not is_qwen3_5_moe_available(), reason="qwen3_5_moe module not available") -def test_apply_liger_kernel_to_instance_for_qwen3_5_moe(): +@pytest.mark.parametrize("base_model", [False, True]) +def test_apply_liger_kernel_to_instance_for_qwen3_5_moe(base_model): # Ensure any monkey patching is cleaned up for subsequent tests with patch("transformers.models.qwen3_5_moe.modeling_qwen3_5_moe"): + from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeTextModel + from liger_kernel.transformers.model.qwen3_5_moe import lce_forward as qwen3_5_moe_lce_forward # Instantiate a dummy model @@ -3671,6 +3740,7 @@ def test_apply_liger_kernel_to_instance_for_qwen3_5_moe(): shared_expert_intermediate_size=16, hidden_act="silu", num_hidden_layers=2, + layer_types=["linear_attention", "full_attention"], num_attention_heads=2, num_key_value_heads=1, head_dim=16, @@ -3682,12 +3752,17 @@ def test_apply_liger_kernel_to_instance_for_qwen3_5_moe(): num_experts=2, num_experts_per_tok=1, ) - dummy_model_instance = AutoModelForCausalLM.from_config(config) + if base_model: + dummy_model_instance = Qwen3_5MoeTextModel._from_config(config) + else: + dummy_model_instance = AutoModelForCausalLM.from_config(config) + text_model = dummy_model_instance if base_model else dummy_model_instance.model # Check that model instance variables are not yet patched with Liger modules - assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(qwen3_5_moe_lce_forward) - assert inspect.getsource(dummy_model_instance.model.norm.forward) != inspect.getsource(LigerRMSNorm.forward) - for layer in dummy_model_instance.model.layers: + if not base_model: + assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(qwen3_5_moe_lce_forward) + assert inspect.getsource(text_model.norm.forward) != inspect.getsource(LigerRMSNorm.forward) + for layer in text_model.layers: if IS_TRANSFORMERS_V5_OR_LATER: assert inspect.getsource(layer.mlp.experts.forward) != inspect.getsource(LigerExperts.forward) else: @@ -3700,12 +3775,15 @@ def test_apply_liger_kernel_to_instance_for_qwen3_5_moe(): assert inspect.getsource(layer.post_attention_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward) # Test applying kernels to the model instance - _apply_liger_kernel_to_instance(model=dummy_model_instance) + _check_qwen_qk_norm_patch( + dummy_model_instance, text_model, gemma=True, fused_linear_cross_entropy=not base_model + ) # Check that the model's instance variables were correctly patched with Liger modules - assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(qwen3_5_moe_lce_forward) - assert inspect.getsource(dummy_model_instance.model.norm.forward) == inspect.getsource(LigerRMSNorm.forward) - for layer in dummy_model_instance.model.layers: + if not base_model: + assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(qwen3_5_moe_lce_forward) + assert inspect.getsource(text_model.norm.forward) == inspect.getsource(LigerRMSNorm.forward) + for layer in text_model.layers: if IS_TRANSFORMERS_V5_OR_LATER: assert inspect.getsource(layer.mlp.experts.forward) == inspect.getsource(LigerExperts.forward) else: @@ -3732,6 +3810,7 @@ def _build_qwen3_5_moe_multimodal_config(): shared_expert_intermediate_size=16, hidden_act="silu", num_hidden_layers=2, + layer_types=["linear_attention", "full_attention"], num_attention_heads=2, num_key_value_heads=1, head_dim=16, @@ -3804,7 +3883,7 @@ def test_apply_liger_kernel_to_instance_for_qwen3_5_moe_for_conditional_generati assert inspect.getsource(layer.post_attention_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward) # Test applying kernels to the model instance - _apply_liger_kernel_to_instance(model=dummy_model_instance) + _check_qwen_qk_norm_patch(dummy_model_instance, dummy_model_instance.model.language_model, gemma=True) # Check that the model's instance variables were correctly patched with Liger modules assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource( @@ -3859,7 +3938,7 @@ def test_apply_liger_kernel_to_instance_for_qwen3_5_moe_model(): assert inspect.getsource(layer.post_attention_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward) # Test applying kernels to the model instance - _apply_liger_kernel_to_instance(model=dummy_model_instance) + _check_qwen_qk_norm_patch(dummy_model_instance, dummy_model_instance.language_model, gemma=True) # Check that the model's instance variables were correctly patched with Liger modules assert inspect.getsource(dummy_model_instance.language_model.norm.forward) == inspect.getsource( @@ -3884,9 +3963,12 @@ def test_apply_liger_kernel_to_instance_for_qwen3_5_moe_model(): @pytest.mark.skipif(not is_qwen3_5_available(), reason="qwen3_5 module not available") -def test_apply_liger_kernel_to_instance_for_qwen3_5(): +@pytest.mark.parametrize("base_model", [False, True]) +def test_apply_liger_kernel_to_instance_for_qwen3_5(base_model): # Ensure any monkey patching is cleaned up for subsequent tests with patch("transformers.models.qwen3_5.modeling_qwen3_5"): + from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5TextModel + # Instantiate a dummy model config = transformers.models.qwen3_5.configuration_qwen3_5.Qwen3_5TextConfig( dtype=torch.bfloat16, @@ -3905,23 +3987,31 @@ def test_apply_liger_kernel_to_instance_for_qwen3_5(): linear_num_value_heads=2, layer_types=["linear_attention", "linear_attention", "linear_attention", "full_attention"], ) - dummy_model_instance = AutoModelForCausalLM.from_config(config) + if base_model: + dummy_model_instance = Qwen3_5TextModel._from_config(config) + else: + dummy_model_instance = AutoModelForCausalLM.from_config(config) + text_model = dummy_model_instance if base_model else dummy_model_instance.model # Check that model instance variables are not yet patched with Liger modules - assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(qwen3_5_lce_forward) - assert inspect.getsource(dummy_model_instance.model.norm.forward) != inspect.getsource(LigerRMSNorm.forward) - for layer in dummy_model_instance.model.layers: + if not base_model: + assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(qwen3_5_lce_forward) + assert inspect.getsource(text_model.norm.forward) != inspect.getsource(LigerRMSNorm.forward) + for layer in text_model.layers: assert inspect.getsource(layer.mlp.forward) != inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward) assert inspect.getsource(layer.input_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward) assert inspect.getsource(layer.post_attention_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward) # Test applying kernels to the model instance - _apply_liger_kernel_to_instance(model=dummy_model_instance) + _check_qwen_qk_norm_patch( + dummy_model_instance, text_model, gemma=True, fused_linear_cross_entropy=not base_model + ) # Check that the model's instance variables were correctly patched with Liger modules - assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(qwen3_5_lce_forward) - assert inspect.getsource(dummy_model_instance.model.norm.forward) == inspect.getsource(LigerRMSNorm.forward) - for layer in dummy_model_instance.model.layers: + if not base_model: + assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(qwen3_5_lce_forward) + assert inspect.getsource(text_model.norm.forward) == inspect.getsource(LigerRMSNorm.forward) + for layer in text_model.layers: assert inspect.getsource(layer.mlp.forward) == inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward) assert inspect.getsource(layer.input_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward) assert inspect.getsource(layer.post_attention_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward) @@ -3993,7 +4083,7 @@ def test_apply_liger_kernel_to_instance_for_qwen3_5_for_conditional_generation() assert inspect.getsource(layer.post_attention_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward) # Test applying kernels to the model instance - _apply_liger_kernel_to_instance(model=dummy_model_instance) + _check_qwen_qk_norm_patch(dummy_model_instance, dummy_model_instance.model.language_model, gemma=True) # Check that the model's instance variables were correctly patched with Liger modules assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(qwen3_5_lce_forward_for_multimodal)