From a412ec6d04b3b97518dcf9b036ed1370e778ee9b Mon Sep 17 00:00:00 2001
From: Ryuichi Leo Takashige <leo@exolabs.net>
Date: Fri, 16 Jan 2026 19:33:22 +0000
Subject: [PATCH] Comment out custom moe layer

---
 src/exo/worker/engines/mlx/auto_parallel.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py
index bbaf8ce4..f93224f2 100644
--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -505,7 +505,7 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
             self.sharded_to_all_linear_in_place(layer.mlp.experts.down_proj)
             self.all_to_sharded_linear_in_place(layer.mlp.experts.up_proj)
 
-            layer.mlp = ShardedGptOssMoE(layer.mlp)  # type: ignore
+            # layer.mlp = ShardedGptOssMoE(layer.mlp)  # type: ignore
             # layer.mlp.sharding_group = self.group
 
         return model