lmstudio-ai / lmstudio-ai/mlx-engine

Bug: Batched VLM loader treats mtp.safetensors as base weights for Qwen3.5-4B-OptiQ-4bit

Open
#343 3 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

## 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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.