electron-rare/exo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

remove moe sharding

Browse Source
This commit is contained in:
Ryuichi Leo Takashige
2026-01-16 19:24:06 +00:00
parent c613df4d8d
commit 9bf1bb3025
1 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/exo/worker/engines/mlx/auto_parallel.py
+6 -6
View File
@@ -501,12 +501,12 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
* (self.group.rank() + 1)
]
self.all_to_sharded_linear_in_place(layer.mlp.experts.gate_proj)
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.sharding_group = self.group
# self.all_to_sharded_linear_in_place(layer.mlp.experts.gate_proj)
# 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.sharding_group = self.group
return model