lmstudio-ai / lmstudio-ai/mlx-engine

Double RMSNorm in PatchedQwen3_5TextModel

Open Beginner friendly
#322 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1.2k
Forks
133
Avg merge
21h 6m
Merged PRs (30d)
1

Description

@will-lms Introduced in #298 (Qwen 3.5 Unified).

`PatchedQwen3_5TextModel.__call__` returns `self.norm(hidden_states)`, but the class it overrides (`Qwen3_5TextModel`) is expected to return un-normed hidden states. mlx-lm's `TextModel.__call__` applies `self.model.norm` itself before the lm_head:

```python
# mlx_lm/models/qwen3_5.py in TextModel.__call__
hidden = self.model(...) # expects un-normed hidden
normed = self.model.norm(hidden)
out = self.lm_head(normed)
```

Since `self.model` is now `PatchedQwen3_5TextModel`, `hidden` is already normed, and `self.model.norm` is applied a second time.

This causes a double RMSNorm before the lm_head on all Qwen3.5 text and vision inference.

## Fix

Do not call `self.norm` in `mlx_engine/model_kit/patches/qwen3_5.py`:

```diff
- return self.norm(hidden_states)
+ return hidden_states
```

Contributor guide

Open the contributing guide

Research direction

Start in mlx_engine/model_kit/patches/qwen3_5.py at PatchedQwen3_5TextModel.__call__, then compare its return value with mlx_lm/models/qwen3_5.py in TextModel.__call__. Confirm the model path applies normalization only once, and verify that Qwen3.5 text and vision inference no longer perform a double RMSNorm before the lm_head.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
1/5
Estimated time
Under an hour
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
88/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.