add prmoe to cifar10 example (deepspeedai#163)

Author: yaozhewei

cifar/cifar10_deepspeed.py

@@ -51,9 +51,18 @@ def add_argument():
                         type=int,
                         help='(moe) expert parallel world size')
     parser.add_argument('--num-experts',
-                        default=1,
                         type=int,
-                        help='(moe) number of total experts')
+                        nargs='+',
+                        default=[
+                            1,
+                        ],
+                        help='number of experts list, MoE related.')
+    parser.add_argument(
+        '--mlp-type',
+        type=str,
+        default='standard',
+        help=
+        'Only applicable when num-experts > 1, accepts [standard, residual]')
     parser.add_argument('--top-k',
                         default=1,
                         type=int,
@@ -168,9 +177,6 @@ def imshow(img):
 
 args = add_argument()
 
-if args.moe:
-    deepspeed.utils.groups.initialize(ep_size=args.ep_world_size)
-
 
 class Net(nn.Module):
     def __init__(self):
@@ -181,14 +187,21 @@ def __init__(self):
         self.fc1 = nn.Linear(16 * 5 * 5, 120)
         self.fc2 = nn.Linear(120, 84)
         if args.moe:
-            self.fc3 = nn.Linear(84, 84)
-            self.fc3 = deepspeed.moe.layer.MoE(
-                hidden_size=84,
-                expert=self.fc3,
-                num_experts=args.num_experts,
-                k=args.top_k,
-                min_capacity=args.min_capacity,
-                noisy_gate_policy=args.noisy_gate_policy)
+            fc3 = nn.Linear(84, 84)
+            self.moe_layer_list = []
+            for n_e in args.num_experts:
+                # create moe layers based on the number of experts
+                self.moe_layer_list.append(
+                    deepspeed.moe.layer.MoE(
+                        hidden_size=84,
+                        expert=fc3,
+                        num_experts=n_e,
+                        ep_size=args.ep_world_size,
+                        use_residual=args.mlp_type == 'residual',
+                        k=args.top_k,
+                        min_capacity=args.min_capacity,
+                        noisy_gate_policy=args.noisy_gate_policy))
+            self.moe_layer_list = nn.ModuleList(self.moe_layer_list)
             self.fc4 = nn.Linear(84, 10)
         else:
             self.fc3 = nn.Linear(84, 10)
@@ -200,7 +213,8 @@ def forward(self, x):
         x = F.relu(self.fc1(x))
         x = F.relu(self.fc2(x))
         if args.moe:
-            x, _, _ = self.fc3(x)
+            for layer in self.moe_layer_list:
+                x, _, _ = layer(x)
             x = self.fc4(x)
         else:
             x = self.fc3(x)
@@ -213,7 +227,10 @@ def forward(self, x):
 def create_moe_param_groups(model):
     from deepspeed.moe.utils import split_params_into_different_moe_groups_for_optimizer
 
-    parameters = {'params': model.parameters(), 'name': 'parameters'}
+    parameters = {
+        'params': [p for p in model.parameters()],
+        'name': 'parameters'
+    }
 
     return split_params_into_different_moe_groups_for_optimizer(parameters)
 

cifar/run_ds_prmoe.sh

@@ -0,0 +1,22 @@
+#!/bin/bash
+
+# Number of nodes
+NUM_NODES=1
+# Number of GPUs per node
+NUM_GPUS=2
+# Size of expert parallel world (should be less than total world size)
+EP_SIZE=2
+# Number of total experts, note here we need to pass >= two numbers (numbers can be different)
+EXPERTS='2 4'
+
+deepspeed --num_nodes=${NUM_NODES} --num_gpus=${NUM_GPUS} cifar10_deepspeed.py \
+	--log-interval 100 \
+	--deepspeed \
+	--deepspeed_config ds_config.json \
+	--moe \
+	--ep-world-size ${EP_SIZE} \
+	--num-experts ${EXPERTS} \
+	--top-k 1 \
+	--mlp-type 'residual' \
+	--noisy-gate-policy 'RSample' \
+	--moe-param-group