From 80fb156fce0962eec12086ae76118451deffe933 Mon Sep 17 00:00:00 2001 From: luca huang Date: Tue, 6 Oct 2026 03:48:09 +0000 Subject: [PATCH] fix(monkey_patch): patch Q/K norms on existing Qwen models (#1470) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Patch `self_attn.q_norm` and `k_norm` when applying Liger to existing Qwen3, Qwen3 MoE, Qwen3 Next, Qwen3.5 and Qwen3.5 MoE instances. Reuse each model's RMSNorm helper and respect `rms_norm`. Extend the existing instance tests following the Qwen3 VL/VL MoE and Gemma3 pattern: construct native HF instances first, then check Q/K norm forward bindings before and after patching. Also verify preserved weights and epsilon, casting settings, and unchanged Linear Attention norms. ## Testing Done - Hardware Type: H100 80GB - Instance tests: 21 passed. All 13 new Q/K regression assertions fail before the fix and pass after it. - RMSNorm numerical tests: 64 passed. - Convergence: 16 passed, 1 known baseline failure, 5 existing skips. The BF16 Qwen3 MoE with-logits failure also occurs on the base commit. These results are from prior validation; tests were not rerun after switching to explicit model classes. PyTorch 2.9.1 / CUDA 12.8, Triton 3.5.1, Transformers 5.15.1. Hybrid convergence uses FLA 0.5.2 and TileLang 0.1.14. Before uses the existing Liger instance patch with native Q/K norms; After adds this fix. Both use matched random weights and inputs, BF16, SDPA, identical other kernel settings, and no `torch.compile`. Each runs in a separate process with warmup and five measurement rounds. Speedup is Before / After. For these configurations, the patch covers all attention layers in Qwen3-8B (36/36) and Qwen3-30B-A3B (48/48). In Qwen3.5-9B and Qwen3.5-35B-A3B, it covers only the Full Attention layers (8/32 and 10/40); Linear Attention norms are unchanged. Q/K norms and a single Full Attention module on one H100 80GB, forward + backward median time in ms: | Model | Batch × seq | Q+K Before → After | Speedup | Attention Before → After | Speedup | | --- | ---: | ---: | ---: | ---: | ---: | | Qwen3-8B | 4 × 2048 | 2.439 → 0.562 | 4.34× | 8.883 → 6.705 | 1.32× | | Qwen3-8B | 1 × 8192 | 2.435 → 0.542 | 4.50× | 13.552 → 11.269 | 1.20× | | Qwen3.5-9B | 4 × 2048 | 2.560 → 0.739 | 3.46× | 10.763 → 8.802 | 1.22× | | Qwen3.5-9B | 1 × 8192 | 2.564 → 0.729 | 3.52× | 16.063 → 14.108 | 1.14× | | Qwen3-30B-A3B | 4 × 2048 | 2.219 → 0.643 | 3.45× | 7.071 → 5.280 | 1.34× | | Qwen3-30B-A3B | 1 × 8192 | 2.212 → 0.659 | 3.36× | 11.597 → 9.858 | 1.18× | | Qwen3.5-35B-A3B | 4 × 2048 | 2.292 → 0.740 | 3.10× | 8.248 → 6.324 | 1.30× | | Qwen3.5-35B-A3B | 1 × 8192 | 2.288 → 0.691 | 3.31× | 13.475 → 11.539 | 1.17× | - [ ] run `make test` (targeted tests above; full suite not run) - [x] run `make checkstyle` (prior validation) - [ ] run `make test-convergence` (targeted tests above; full suite not run) Co-authored-by: Arup De --- src/liger_kernel/transformers/monkey_patch.py | 13 ++ test/transformers/test_monkey_patch.py | 186 +++++++++++++----- 2 files changed, 151 insertions(+), 48 deletions(-) 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)