1) extending _Linear's input as 4-dimension tensor: [outer_batch, expert_dim, -1, hidden], and _Liner's BatchMatMul becomes BatchMatMul(4_dim_tensor, 3_dim_tensor);
2) configuring the _Linear's BatchMatMul sharding strategy as [[dp, ep, 1, 1], [ep, 1, mp]];
3) introducing a new parameter 'expert_parallel' in TransformerOpParallelConfig, creating a new class MoEParallelConfig to include 'data_parallel', 'model_parallel' and 'expert_parallel';
4) changing parallel config for FeedForward, TransformerEncoderLayer, TransformerDecoderLayer.
1) extend the Liner cell for including BatchMatMul implementation, in
which the first dimension indicates the expert number;
2) implement a Switch (top1) router;
3) implement a MoE cell, which extends the FeedForward cell.