pass test

.gitignore

@@ -7,3 +7,6 @@ exp/
 a.out
 *.egg-info
 *.egg
+build
+*swp
+logs

fmoe/moe.py

@@ -8,7 +8,7 @@ from .moe_function import moe
 class FMoE(nn.Module):
     def __init__(self, num_expert=32, in_feat=1024, out_feat=1024,
             world_size=None):
-        super(MOELayer, self).__init__()
+        super(FMoE, self).__init__()
         self.num_expert = num_expert
         self.in_feat = in_feat
         self.out_feat = out_feat
@@ -29,7 +29,7 @@ class FMoE(nn.Module):
 class BruteForceMoE(nn.Module):
     def __init__(self, num_expert=32, in_feat=1024, out_feat=1024, 
             world_size=0):
-        super(MOELayer_raw, self).__init__()
+        super(BruteForceMoE, self).__init__()
         self.num_expert = num_expert
         self.in_feat = in_feat
         self.out_feat = out_feat

tests/dev_test.sh

@@ -4,11 +4,21 @@ then
 	export CUDA_VISIBLE_DEVICES=$OMPI_COMM_WORLD_LOCAL_RANK
 fi
 
-export PYTHONPATH=$PWD/build/lib.linux-x86_64-3.7
+if [ -z $OMPI_COMM_WORLD_RANK ]
+then
+	RANK=single
+else
+	RANK=$OMPI_COMM_WORLD_RANK 
+fi
+
+mkdir -p logs
+
+SCRIPT_PATH=$(dirname $(dirname $(realpath $0)))
+export PYTHONPATH=$SCRIPT_PATH:$SCRIPT_PATH/build/lib.linux-x86_64-3.7:$PYTHONPATH
 export LD_LIBRARY_PATH=/home/laekov/.local/lib/python3.7/site-packages/torch/lib:$LD_LIBRARY_PATH
 if [ -z $1 ]
 then
-	python3 moe_test.py 2>logs/$OMPI_COMM_WORLD_RANK.log
+	python3 tests/moe_test.py 2>logs/$RANK.log
 else
-	python3 $@ 2>logs/$OMPI_COMM_WORLD_RANK.log
+	python3 $@ 2>logs/$RANK.log
 fi

tests/moe_test.py

-from moe import MOELayer, MOELayer_raw
+from fmoe import FMoE as MOELayer 
+from fmoe import BruteForceMoE as MOELayer_raw
 import torch
 from torch import nn
 import time