← back to Exo

Fix GPT OSS tensor sharding with upstream MLX LM (#1223)

758464703de2e76908d3a60e7a03245f655c1132 · 2026-01-20 18:24:54 +0000 · rltakashige

## Motivation
MLX LM has given GPT OSS a shard method, but MLX does not have an update
to match.

## Changes

<!-- Describe what you changed in detail -->

## Why It Works

<!-- Explain why your approach solves the problem -->

## Test Plan

### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
<!-- What you did: -->
<!-- - -->

### Automated Testing
<!-- Describe changes to automated tests, or how existing tests cover
this change -->
<!-- - -->

Files touched

• M src/exo/worker/engines/mlx/auto_parallel.py

Diff

commit 758464703de2e76908d3a60e7a03245f655c1132
Author: rltakashige <rl.takashige@gmail.com>
Date:   Tue Jan 20 18:24:54 2026 +0000

    Fix GPT OSS tensor sharding with upstream MLX LM (#1223)
    
    ## Motivation
    MLX LM has given GPT OSS a shard method, but MLX does not have an update
    to match.
    
    ## Changes
    
    <!-- Describe what you changed in detail -->
    
    ## Why It Works
    
    <!-- Explain why your approach solves the problem -->
    
    ## Test Plan
    
    ### Manual Testing
    <!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
    connected via Thunderbolt 4) -->
    <!-- What you did: -->
    <!-- - -->
    
    ### Automated Testing
    <!-- Describe changes to automated tests, or how existing tests cover
    this change -->
    <!-- - -->
---
 src/exo/worker/engines/mlx/auto_parallel.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py
index cb0ea2a4..159c2a45 100644
--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -334,7 +334,7 @@ def tensor_auto_parallel(
         group=group,
     )
 
-    if hasattr(model, "shard"):
+    if hasattr(model, "shard") and not isinstance(model, GptOssModel):
         try:
             model.shard(group)  # type: ignore
             return patch_tensor_model(model)
@@ -383,7 +383,6 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
         )
-
     else:
         raise ValueError(f"Unsupported model type: {type(model)}")
 

← 9e2179c8 Register original layer in CustomMlxLayer (#1229)  ·  back to Exo  ·  Add mflux type stubs (#1234) 6a9251b9 →