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 <[email protected]>

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)

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