oops

lucidrains · Sep 10, 2023 · 52b5c8a · 52b5c8a
1 parent 83d75b8
commit 52b5c8a
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/assert.py b/assert.py
@@ -76,7 +76,7 @@ def start(
     cleanup()
 
 if __name__ == '__main__':
-    world_size = 5
+    world_size = 8
     num_experts = 8
     batch_size = 2
     batch_size_var_len = False

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'st-moe-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.27',
+  version = '0.0.28',
   license='MIT',
   description = 'ST - Mixture of Experts - Pytorch',
   author = 'Phil Wang',

diff --git a/st_moe_pytorch/st_moe_pytorch.py b/st_moe_pytorch/st_moe_pytorch.py
@@ -227,7 +227,7 @@ def forward(
 
                 assert divisible_by(world_size, num_experts), 'if number of machines is greater than number of experts, machines must be divisible by number of experts, so experts are evenly distributed'
                 num_experts_per_rank = 1
-                expert_start_index = rank // num_experts
+                expert_start_index = rank // (world_size // num_experts)
 
             expert_slice = slice(expert_start_index, expert_start_index + num_experts_per_rank)
         else: