[object Object]

← back to Exo

FIX GPT OSS (#1165)

aaf4e36bc35677d9457cc9a4583d7b8f2bf7c248 · 2026-01-15 19:20:17 +0000 · rltakashige

## Motivation

Adds several unmerged fixes for GPT OSS.
Also adds GPT OSS 20B MXFP4 Q8 instead of Q4 for numerical stability (as
this is unstable for MLX LM too)
<!-- Why is this change needed? What problem does it solve? -->
<!-- If it fixes an open issue, please link to the issue here -->


## Test Plan

### Manual Testing
Manually tested. No further gibberish responses.

### Automated Testing
Ran EXO Bench - pipeline, tensor and single node work on both 20B and
120B models

Files touched

Diff

commit aaf4e36bc35677d9457cc9a4583d7b8f2bf7c248
Author: rltakashige <rl.takashige@gmail.com>
Date:   Thu Jan 15 19:20:17 2026 +0000

    FIX GPT OSS (#1165)
    
    ## Motivation
    
    Adds several unmerged fixes for GPT OSS.
    Also adds GPT OSS 20B MXFP4 Q8 instead of Q4 for numerical stability (as
    this is unstable for MLX LM too)
    <!-- Why is this change needed? What problem does it solve? -->
    <!-- If it fixes an open issue, please link to the issue here -->
    
    
    ## Test Plan
    
    ### Manual Testing
    Manually tested. No further gibberish responses.
    
    ### Automated Testing
    Ran EXO Bench - pipeline, tensor and single node work on both 20B and
    120B models
---
 src/exo/shared/models/model_cards.py    | 14 +++++++-------
 src/exo/worker/engines/mlx/utils_mlx.py |  6 ++++++
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/src/exo/shared/models/model_cards.py b/src/exo/shared/models/model_cards.py
index 18768135..56d83113 100644
--- a/src/exo/shared/models/model_cards.py
+++ b/src/exo/shared/models/model_cards.py
@@ -425,15 +425,15 @@ MODEL_CARDS: dict[str, ModelCard] = {
             supports_tensor=True,
         ),
     ),
-    "gpt-oss-20b-4bit": ModelCard(
-        short_id="gpt-oss-20b-4bit",
-        model_id=ModelId("mlx-community/gpt-oss-20b-MXFP4-Q4"),
-        name="GPT-OSS 20B (MXFP4-Q4, MLX)",
-        description="""OpenAI's GPT-OSS 20B is a medium-sized MoE model for lower-latency and local or specialized use cases; this MLX variant uses MXFP4 4-bit quantization.""",
+    "gpt-oss-20b-MXFP4-Q8": ModelCard(
+        short_id="gpt-oss-20b-MXFP4-Q8",
+        model_id=ModelId("mlx-community/gpt-oss-20b-MXFP4-Q8"),
+        name="GPT-OSS 20B (MXFP4-Q8, MLX)",
+        description="""OpenAI's GPT-OSS 20B is a medium-sized MoE model for lower-latency and local or specialized use cases; this variant is a 4-bit MLX conversion for Apple Silicon.""",
         tags=[],
         metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/gpt-oss-20b-MXFP4-Q4"),
-            pretty_name="GPT-OSS 20B (MXFP4-Q4, MLX)",
+            model_id=ModelId("mlx-community/gpt-oss-20b-MXFP4-Q8"),
+            pretty_name="GPT-OSS 20B (MXFP4-Q8, MLX)",
             storage_size=Memory.from_kb(11_744_051),
             n_layers=24,
             hidden_size=2880,
diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py
index 4a4e45cf..64807eb6 100644
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -20,6 +20,7 @@ except ImportError:
 
 from mlx_lm.models.cache import KVCache, QuantizedKVCache, RotatingKVCache
 from mlx_lm.models.deepseek_v3 import DeepseekV3Model
+from mlx_lm.models.gpt_oss import Model as GptOssModel
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
 from exo.worker.engines.mlx.constants import (
@@ -396,6 +397,11 @@ def make_kv_cache(
 ) -> list[KVCache | RotatingKVCache | QuantizedKVCache]:
     assert hasattr(model, "layers")
 
+    # TODO: Do this for all models
+    if hasattr(model, "make_cache") and isinstance(model, GptOssModel):
+        logger.info("Using MLX LM's make cache")
+        return model.make_cache()  # type: ignore
+
     if max_kv_size is None:
         if KV_CACHE_BITS is None:
             logger.info("Using default KV cache")

← 3e623ccf up http timeout to 3 seconds and retry on BadStatusLine (#11  ·  back to Exo  ·  Parse GPT OSS in runner (#1160) a735dad6 →