check moe input/output batch sizes are the same

fmoe/layers.py

@@ -179,6 +179,14 @@ class FMoE(nn.Module):
         according to the gate.  The score of the selected gate given by the
         expert is multiplied to the experts' output tensors as a weight.
         """
+
+        moe_inp_batch_size = tree.flatten(
+            tree.map_structure(lambda tensor: tensor.shape[0], moe_inp)
+        )
+        assert all(
+            [batch_size == moe_inp_batch_size[0] for batch_size in moe_inp_batch_size]
+        ), "MoE inputs must have the same batch size"
+
         if self.world_size > 1:
 
             def ensure_comm_func(tensor):
@@ -199,14 +207,14 @@ class FMoE(nn.Module):
         if self.gate_hook is not None:
             self.gate_hook(gate_top_k_idx, gate_score, None)
 
+        # delete masked tensors
+        if self.mask is not None and self.mask_dict is not None:
             # TODO: to fix
             def delete_mask_func(tensor):
                 # to: (BxL') x d_model
                 tensor = tensor[mask == 0, :]
                 return tensor
 
-        # delete masked tensors
-        if self.mask is not None and self.mask_dict is not None:
             mask = self.mask.view(-1)
             moe_inp = tree.map_structure(delete_mask_func, moe_inp)
             gate_top_k_idx = gate_top_k_idx[mask == 0, :]
@@ -263,4 +271,11 @@ class FMoE(nn.Module):
                 )
 
             moe_outp = tree.map_structure(all_gather_func, moe_outp)
+
+        moe_outp_batch_size = tree.flatten(
+            tree.map_structure(lambda tensor: tensor.shape[0], moe_outp)
+        )
+        assert all(
+            [batch_size == moe_outp_batch_size[0] for batch_size in moe_outp_batch_size]
+        ), "MoE outputs must have the same batch size"
         return moe_outp