← back to Exo

fix qwen moe tensor sharding (#1604)

12c3015f52bfa9ff6e3a64951a7820ed48903483 · 2026-02-23 21:23:27 +0000 · rltakashige

## Motivation

<!-- Why is this change needed? What problem does it solve? -->
<!-- If it fixes an open issue, please link to the issue here -->

## 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 12c3015f52bfa9ff6e3a64951a7820ed48903483
Author: rltakashige <rl.takashige@gmail.com>
Date:   Mon Feb 23 21:23:27 2026 +0000

    fix qwen moe tensor sharding (#1604)
    
    ## Motivation
    
    <!-- Why is this change needed? What problem does it solve? -->
    <!-- If it fixes an open issue, please link to the issue here -->
    
    ## 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 | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py
index f2490006..9acd3163 100644
--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -852,6 +852,8 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                 layer.self_attn.o_proj = self.sharded_to_all_linear(
                     layer.self_attn.o_proj
                 )
+                layer.self_attn.n_heads //= self.N
+                layer.self_attn.n_kv_heads //= self.N
             else:
                 assert isinstance(layer, Qwen3NextDecoderLayer)
                 if hasattr(layer, "linear_attn"):

← 365dd68d Final fixes for release (#1603)  ·  back to Exo  ·  fix: improve text contrast on downloads page (#1601) 644c5573 →