lmstudio-ai / lmstudio-ai/mlx-engine
Bug: Batched VLM loader treats mtp.safetensors as base weights for Qwen3.5-4B-OptiQ-4bit
- Dominant language
- Python
- Stars
- 1.2k
- Forks
- 133
- Avg merge
- 21h 6m
- Merged PRs (30d)
- 1
Description
## Summary
[`mlx-community/Qwen3.5-4B-OptiQ-4bit`](https://huggingface.co/mlx-community/Qwen3.5-4B-OptiQ-4bit/tree/main) fails to load in LM Studio's MLX runtime because the batched VLM loading path appears to pass `mtp.safetensors` into the base model weight load.
The immediate bug is that `mtp.*` sidecar weights are being strict-loaded as ordinary base model parameters.
Related issues:
- https://github.com/lmstudio-ai/mlx-engine/issues/323
- https://github.com/lmstudio-ai/lmstudio-bug-tracker/issues/2103
## Environment
- LM Studio: `0.4.19+2`
- `lms` CLI commit: `9902c3a`
- Selected MLX runtime: `mlx-llm-mac-arm64-apple-metal-nax-advsimd@1.10.0`
- macOS: `26.5.1`, `arm64`
- Model: `mlx-community/Qwen3.5-4B-OptiQ-4bit`
Model files include:
```text
model.safetensors
optiq_vision.safetensors
mtp.safetensors
model.safetensors.index.json
config.json
```
Relevant config fields:
```json
{
"architectures": ["Qwen3_5ForConditionalGeneration"],
"model_type": "qwen3_5",
"mtp_file": "mtp.safetensors",
"mtp_tensor_count": 29,
"mtp_policy": "optiq-int4-prequantized-gs64",
"optiq_vision": {
"sidecar": "optiq_vision.safetensors"
}
}
```
`model.safetensors.index.json` references the base `model.safetensors`, but does not list the `mtp.safetensors` sidecar.
## Reproduction
After downloading `mlx-community/Qwen3.5-4B-OptiQ-4bit` in LM Studio, try loading it with the MLX runtime.
I reproduced the same failure with:
```bash
lms load qwen3.5-4b-optiq --identifier qwen35-optiq-test --yes
lms load qwen3.5-4b-optiq --identifier qwen35-optiq-test --speculative-draft-mtp --yes
lms load qwen3.5-4b-optiq --identifier qwen35-optiq-test --no-speculative-draft-mtp --yes
```
## Actual Behavior
Model loading fails before generation starts.
The traceback goes through the batched VLM path:
```text
mlx_engine/model_kit/batched_vision/model_kit.py
self.model = mlx_vlm.utils.load_model(...)
mlx_vlm/utils.py
model.load_weights(list(weights.items()), strict=strict)
mlx/nn/layers/base.py
raise ValueError(...)
```
The final error is:
```text
ValueError: Received 29 parameters not in model:
mtp.fc.weight,
mtp.layers.0.input_layernorm.weight,
mtp.layers.0.mlp.down_proj.biases,
mtp.layers.0.mlp.down_proj.scales,
mtp.layers.0.mlp.down_proj.weight,
mtp.layers.0.mlp.gate_proj.biases,
mtp.layers.0.mlp.gate_proj.scales,
mtp.layers.0.mlp.gate_proj.weight,
mtp.layers.0.mlp.up_proj.biases,
mtp.layers.0.mlp.up_proj.scales,
mtp.layers.0.mlp.up_proj.weight,
mtp.layers.0.post_attention_layernorm.weight,
mtp.layers.0.self_attn.k_norm.weight,
mtp.layers.0.self_attn.k_proj.biases,
mtp.layers.0.self_attn.k_proj.scales,
mtp.layers.0.self_attn.k_proj.weight,
mtp.layers.0.self_attn.o_proj.biases,
mtp.layers.0.self_attn.o_proj.scales,
mtp.layers.0.self_attn.o_proj.weight,
mtp.layers.0.self_attn.q_norm.weight,
mtp.layers.0.self_attn.q_proj.biases,
mtp.layers.0.self_attn.q_proj.scales,
mtp.layers.0.self_attn.q_proj.weight,
mtp.layers.0.self_attn.v_proj.biases,
mtp.layers.0.self_attn.v_proj.scales,
mtp.layers.0.self_attn.v_proj.weight,
mtp.norm.weight,
mtp.pre_fc_norm_embedding.weight,
mtp.pre_fc_norm_hidden.weight
```
## Expected Behavior
The model should at least load without MTP acceleration.
If MTP speculative decoding is not implemented for this LM Studio MLX VLM path yet, the runtime should either:
1. Ignore or exclude the configured `mtp_file` during base model loading.
2. Load it through an explicit MTP drafter/head path only when MTP is supported and enabled.
It should not pass `mtp.safetensors` into strict base model loading as normal model weights.
## Suspected Cause
The vendored `mlx_vlm.utils.load_model()` path appears to collect all `*.safetensors` files in the model directory and merge them into a single weight dictionary.
For this model, that means it includes:
```text
model.safetensors
optiq_vision.safetensors
mtp.safetensors
```
The `mtp.safetensors` file contains 29 tensors with the `mtp.*` prefix. Those are MTP sidecar weights, not base model parameters, so `model.load_weights(..., strict=True)` rejects them.
I tested the loading behavior locally without modifying LM Studio files:
- Loading only `model.safetensors` fails because `vision_tower.*` weights are missing.
- Loading `model.safetensors` plus `optiq_vision.safetensors`, while excluding `mtp.safetensors`, gets past the base model weight-load step.
So the minimal fix may be to select base plus configured vision sidecar weights, while excluding the configured `mtp_file` from the base weight load.
## Scope Clarification
This issue is intentionally about the base load failure.
There is a related but larger feature request: supporting MTP acceleration for Qwen3.5/Qwen3.6 OptiQ VLM models inside LM Studio. That seems related to https://github.com/lmstudio-ai/mlx-engine/issues/323, but this issue blocks even non-MTP loading.
## Contribution Notes Checked
I checked the repository guidance before filing this:
- `CONTRIBUTING.md` asks bug reports to include reproduction steps, actual behavior, and expected behavior.
- `CONTRIBUTING.md` also recommends opening or commenting on an issue before substantial changes.
- `README.md` indicates local development uses Python 3.11, `requirements.txt`, pre-commit, and `python -m pytest tests/`.
Contributor guide
Research direction
Start in mlx_engine/model_kit/batched_vision/model_kit.py and trace its call into mlx_vlm.utils.load_model(), focusing on how safetensors files are collected and passed to strict loading. Reproduce with the Qwen3.5-4B-OptiQ-4bit model and verify that base plus optiq_vision.safetensors loads successfully while the configured mtp_file is excluded from the base load.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100