lmstudio-ai / lmstudio-ai/mlx-engine
Double RMSNorm in PatchedQwen3_5TextModel
- 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
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