lmstudio-ai / lmstudio-ai/mlx-engine

BatchedModelKit does not respect multiple EOS token IDs from model config

Open
#374 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
1.2k
Forks
133
Avg merge
21h 6m
Merged PRs (30d)
1

Description

### Summary

`BatchedModelKit` does not currently forward the model's configured `eos_token_id` values when loading the tokenizer.

This breaks models that define multiple EOS / EOG tokens. In such cases, batched generation may continue past a valid end-of-generation token into the next serialized conversation turn.

I reproduced this with `lmstudio-community/GLM-4.7-Flash-MLX-8bit`.

### Reproduction

Using the LM Studio OpenAI-compatible API:

```bash
curl http://localhost:1234/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "zai-org/glm-4.7-flash",
"messages": [
{
"role": "user",
"content": "Answer in one word: what is the capital of France?"
}
],
"temperature": 0.2,
"max_tokens": 300
}'
```

Without a manually configured stop string, the response continues after the assistant answer:

```text
Paris<|user|>...
```

and eventually finishes with:

```text
finish_reason: length
```

Adding:

```json
"stop": ["<|user|>"]
```

makes generation terminate correctly with:

```text
Paris
```

and:

```text
finish_reason: stop
```

Direct generation through `mlx_lm.generate` also terminates correctly, so the issue appears specific to the batched engine path.

### Root cause

The model defines multiple EOS token IDs:

```json
"eos_token_id": [
154820,
154827,
154829
]
```

where one of the additional EOS/EOG tokens corresponds to `<|user|>`.

`BatchedModelKit` loads the tokenizer without forwarding these configured EOS IDs:

```python
self.tokenizer = mlx_lm.tokenizer_utils.load(self._model_path)
```

As a result, the tokenizer can fall back to its default/single EOS token instead of preserving the complete model-defined EOS set.

There is also a related configuration difference: `BatchedModelKit` reads `config.json` directly, while `mlx_lm` uses `load_config()`, which also applies relevant values from `generation_config.json`.

### Proposed fix

Use the same configuration path as `mlx_lm` and explicitly forward the configured EOS token IDs:

```python
config_json = mlx_lm.utils.load_config(model_path)

self.tokenizer = mlx_lm.tokenizer_utils.load(
self._model_path,
eos_token_ids=config_json.get("eos_token_id", None),
)
```

This preserves multiple EOS/EOG tokens and respects `generation_config.json` overrides.

### Verification

With this change applied locally, GLM-4.7-Flash stops correctly at `<|user|>` without any manually configured LM Studio Stop Strings.

I also added regression tests covering:

* multiple `eos_token_id` values from the model configuration;
* `eos_token_id` overrides from `generation_config.json`.

The targeted `BatchedModelKit` scheduler tests pass after the change.

### Fix / reference implementation

I prepared the fix and regression tests in my fork:

https://github.com/iphizic/mlx-engine

I cannot currently open a pull request against this repository, so I am providing the implementation here for reference. I would be happy to adjust the patch if a maintainer prefers a different approach.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Begin with BatchedModelKit's tokenizer initialization and compare it with mlx_lm.utils.load_config and tokenizer_utils.load. Run the targeted BatchedModelKit scheduler tests, including cases for model-config and generation_config.json EOS IDs. Done means multiple EOS/EOG IDs are preserved and the GLM reproduction stops at the user token without a manual stop string.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend, machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.