lmstudio-ai / lmstudio-ai/mlx-engine
BatchedModelKit does not respect multiple EOS token IDs from model config
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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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