[object Object]

← back to Exo

Extend eos_token_id fix for other models (#938)

a1b14a272ea4451ad3a073c6477bf2a4b85afb7f · 2025-12-20 20:18:17 +0000 · rltakashige

## Motivation

<!-- Why is this change needed? What problem does it solve? -->
We currently use mlx_lm's load_tokenizer instead of load. This means
that some models are missing some configurations, such as eos_token_id.
This is clear for a model like GLM, which does not finish token
generation.

## Changes

<!-- Describe what you changed in detail -->
A small stopgap, to allow eos_token_ids to be added, and a TODO for us
to migrate to load. The reason we don't want to do this now is that a
solid testing framework is not configured in this repo yet.

## Why It Works

<!-- Explain why your approach solves the problem -->
It just uses the eos_token_ids I obtained from loading a tokenizer in
mlx_lm and calling `tokenizer.eos_token_ids` .

## Test Plan

### Manual Testing
Tested on several Macs.

### Automated Testing
None yet, as described.

---------

Co-authored-by: Evan <evanev7@gmail.com>

Files touched

Diff

commit a1b14a272ea4451ad3a073c6477bf2a4b85afb7f
Author: rltakashige <rl.takashige@gmail.com>
Date:   Sat Dec 20 20:18:17 2025 +0000

    Extend eos_token_id fix for other models (#938)
    
    ## Motivation
    
    <!-- Why is this change needed? What problem does it solve? -->
    We currently use mlx_lm's load_tokenizer instead of load. This means
    that some models are missing some configurations, such as eos_token_id.
    This is clear for a model like GLM, which does not finish token
    generation.
    
    ## Changes
    
    <!-- Describe what you changed in detail -->
    A small stopgap, to allow eos_token_ids to be added, and a TODO for us
    to migrate to load. The reason we don't want to do this now is that a
    solid testing framework is not configured in this repo yet.
    
    ## Why It Works
    
    <!-- Explain why your approach solves the problem -->
    It just uses the eos_token_ids I obtained from loading a tokenizer in
    mlx_lm and calling `tokenizer.eos_token_ids` .
    
    ## Test Plan
    
    ### Manual Testing
    Tested on several Macs.
    
    ### Automated Testing
    None yet, as described.
    
    ---------
    
    Co-authored-by: Evan <evanev7@gmail.com>
---
 src/exo/worker/engines/mlx/utils_mlx.py | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py
index 3606b90b..19d565ca 100644
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -252,15 +252,22 @@ def shard_and_load(
 
 
 def get_tokenizer(model_path: Path, shard_metadata: ShardMetadata):
+    # TODO: Let's move away from this custom logic to mlx_lm.load()
+    if "kimi-k2" in shard_metadata.model_meta.model_id.lower():
+        eos_token_ids = [163586]
+
+    elif "glm" in shard_metadata.model_meta.model_id.lower():
+        eos_token_ids = [151336, 151329, 151338]
+
+    else:
+        eos_token_ids = None
+
     tokenizer = cast(
         TokenizerWrapper,
         load_tokenizer(
             model_path,
             tokenizer_config_extra={"trust_remote_code": TRUST_REMOTE_CODE},
-            # TODO: HACK for Kimi K2 wrong eos token id
-            eos_token_ids=[163586]
-            if "kimi-k2" in shard_metadata.model_meta.model_id.lower()
-            else None,
+            eos_token_ids=eos_token_ids,
         ),
     )
     assert isinstance(tokenizer, TokenizerWrapper)

← f8483cfc Update README.md. (#932)  ·  back to Exo  ·  fix(downloads): use certifi for robust SSL certificate verif f4792dce →