From 488acbd9c0776042cc82337c667157b105394cce Mon Sep 17 00:00:00 2001 From: Edson <113824357+Zhuoxi2000@users.noreply.github.com> Date: Mon, 5 Oct 2026 21:26:33 +0000 Subject: [PATCH] Fix llama4 class-level SwiGLU patch for configs with dense layers (#1507) ## Summary `apply_liger_kernel_to_llama4(swiglu=True)` (the default) sets `modeling_llama4.Llama4TextMLP = LigerSwiGLUMLP`. Transformers builds dense (non-MoE) Llama4 decoder layers as `Llama4TextMLP(config, intermediate_size=config.intermediate_size_mlp)`, but `LigerSwiGLUMLP.__init__` only accepted `config`. Once the class-level patch was applied before model init (`apply_liger_kernel_to_llama4()`, `_apply_liger_kernel("llama4")`, `AutoLigerKernelForCausalLM`), building any Llama4 model with a dense layer raised: ``` TypeError: LigerSwiGLUMLP.__init__() got an unexpected keyword argument 'intermediate_size' ``` This affects Llama-4-Maverick (`interleave_moe_layer_step=2`) and any config whose `moe_layers` does not cover every layer. The `mini_llama4` convergence config uses the default `interleave_moe_layer_step=1`, so every layer there is MoE and the path was never exercised. #907 fixed the instance-level path only. Changes: - `swiglu.py`: `LigerSwiGLUMLP.__init__(self, config, intermediate_size=None)` now falls back to `config.intermediate_size`, which is the same pattern as `LigerQwen3MoeSwiGLUMLP` and HF's `Llama4TextMLP`. Existing callers that pass only `config` behave exactly as before. - `test_monkey_patch.py`: adds `test_apply_liger_kernel_to_llama4_class_patch_with_dense_layers`. The test applies the class-level patch (swiglu only), builds a 2-layer `Llama4ForCausalLM` with `moe_layers=[1]`, and checks two things: - the dense layer's MLP is a `LigerSwiGLUMLP` of width `intermediate_size_mlp`; - the MoE `shared_expert` keeps width `intermediate_size`. It restores `modeling_llama4.Llama4TextMLP` via `monkeypatch`. ## Testing Done The new test fails without the fix and passes with it: ``` # main @ 43293d8 + new test only $ python -m pytest test/transformers/test_monkey_patch.py -k test_apply_liger_kernel_to_llama4_class_patch_with_dense_layers FAILED ... TypeError: LigerSwiGLUMLP.__init__() got an unexpected keyword argument 'intermediate_size' 1 failed # with the fix $ python -m pytest test/transformers/test_monkey_patch.py -k llama4 3 passed $ python -m pytest test/transformers/test_monkey_patch.py 67 passed ``` `ruff check` and `ruff format --check` (ruff 0.15.22, within the pinned range) pass on the changed files. - Hardware Type: Apple M4 (CPU only, no CUDA). transformers 5.18.0, torch 2.14.1. - [ ] run `make test` to ensure correctness. I ran only `test/transformers/test_monkey_patch.py`, because the kernel tests need a GPU. The change only touches module construction (layer widths), not the kernel. - [x] run `make checkstyle` to ensure code style. I ran the same `ruff check` / `ruff format --check` on the two changed files. - [ ] run `make test-convergence` to ensure convergence. Not run because I have no GPU. Default-width behaviour is unchanged for every existing caller. AI assistance: this change was drafted with an AI coding assistant (Claude) and verified locally with the tests above. Co-authored-by: Vaibhav Jindal --- src/liger_kernel/transformers/swiglu.py | 4 +-- test/transformers/test_monkey_patch.py | 35 +++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/src/liger_kernel/transformers/swiglu.py b/src/liger_kernel/transformers/swiglu.py index 7acbd89c..83195bca 100644 --- a/src/liger_kernel/transformers/swiglu.py +++ b/src/liger_kernel/transformers/swiglu.py @@ -20,11 +20,11 @@ def _swiglu_dispatch(a, b, gate_multiplier=1.0, down_multiplier=1.0): class LigerSwiGLUMLP(nn.Module): - def __init__(self, config): + def __init__(self, config, intermediate_size=None): super().__init__() self.config = config self.hidden_size = config.hidden_size - self.intermediate_size = config.intermediate_size + self.intermediate_size = intermediate_size if intermediate_size is not None else config.intermediate_size self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) diff --git a/test/transformers/test_monkey_patch.py b/test/transformers/test_monkey_patch.py index 1ac81c31..cf07478a 100755 --- a/test/transformers/test_monkey_patch.py +++ b/test/transformers/test_monkey_patch.py @@ -1557,6 +1557,41 @@ def test_apply_liger_kernel_to_instance_for_llama4_for_causal_lm(): pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}") +@pytest.mark.skipif(not is_llama4_available(), reason="llama4 module not available") +def test_apply_liger_kernel_to_llama4_class_patch_with_dense_layers(monkeypatch): + # Llama4 checkpoints that interleave dense and MoE layers (e.g. Maverick, interleave_moe_layer_step=2) + # build their dense feed-forward as ``Llama4TextMLP(config, intermediate_size=config.intermediate_size_mlp)``. + # The class-level patch (used before model init, e.g. by AutoLigerKernelForCausalLM) must keep that working. + from transformers.models.llama4 import modeling_llama4 + from transformers.models.llama4.modeling_llama4 import Llama4ForCausalLM + + # Restore the module-level symbols swapped by the class-level patch once the test is done + monkeypatch.setattr(modeling_llama4, "Llama4TextMLP", modeling_llama4.Llama4TextMLP) + + monkey_patch.apply_liger_kernel_to_llama4( + rope=False, rms_norm=False, swiglu=True, fused_linear_cross_entropy=False, layer_norm=False + ) + + config = transformers.models.llama4.configuration_llama4.Llama4TextConfig( + hidden_size=32, + intermediate_size=64, + intermediate_size_mlp=128, + hidden_act="silu", + num_hidden_layers=2, + moe_layers=[1], + ) + model = Llama4ForCausalLM._from_config(config) + + dense_mlp = model.model.layers[0].feed_forward + assert isinstance(dense_mlp, LigerSwiGLUMLP) + assert dense_mlp.gate_proj.out_features == config.intermediate_size_mlp + assert dense_mlp.down_proj.in_features == config.intermediate_size_mlp + + shared_expert = model.model.layers[1].feed_forward.shared_expert + assert isinstance(shared_expert, LigerSwiGLUMLP) + assert shared_expert.gate_proj.out_features == config.intermediate_size + + @pytest.mark.skipif(not is_llama4_available(), reason="llama4 module not available") def test_apply_liger_kernel_to_instance_for_llama4_for_conditional_generation(): # Ensure any monkey patching is cleaned up for subsequent tests