MOE: Support disable top2 2nd expert sampling (#362)

DeepSpeed's MoE top2 gating performs sampling to select 2nd expert.
Support configuration for disabling of sampling (i.e. using argmax).
New argument: --disable-moe-top2-2nd-expert-sampling.

Signed-off-by: Moshe Island <misland@habana.ai>
Co-authored-by: Moshe Island <misland@habana.ai>

Author: mosheisland

megatron/arguments.py

@@ -899,6 +899,9 @@ def _add_training_args(parser):
     group.add_argument('--create-moe-param-group', action='store_true',
                        help='Create separate groups for MoE params.'
                        'This is necessary for techniques like ZeRO.')
+    group.add_argument('--disable-moe-top2-2nd-expert-sampling', action='store_false',
+                       help='Disable MoE top2 sampling of the 2nd expert. Instead of sampling, use argmax.',
+                       dest='moe_top2_2nd_expert_sampling')
     group.add_argument('--use-flash-attn', '--use-flash-attn-v1', dest='use_flash_attn_v1', action='store_true',
                        help='use first version FlashAttention implementation of attention. '
                        'https://arxiv.org/abs/2205.14135')

megatron/model/transformer.py

@@ -952,18 +952,20 @@ def __init__(self, config,
             else: # DeepSpeed's MoE
                 enable_expert_tensor_parallelism = args.enable_expert_tensor_parallelism
                 self.mlp = MoE(args.hidden_size,
-                                ParallelMLP(config,
-                                    moe=True,
-                                    enable_expert_tensor_parallelism=enable_expert_tensor_parallelism),
-                                num_experts=self.num_experts,
-                                ep_size=args.moe_expert_parallel_size,
-                                k=args.topk,
-                                use_residual=(args.mlp_type == 'residual'),
-                                capacity_factor=args.moe_train_capacity_factor,
-                                eval_capacity_factor=args.moe_eval_capacity_factor,
-                                min_capacity=args.moe_min_capacity,
-                                drop_tokens=args.moe_token_dropping, use_tutel=args.use_tutel,
-                                enable_expert_tensor_parallelism=enable_expert_tensor_parallelism)
+                               ParallelMLP(config,
+                                           moe=True,
+                                           enable_expert_tensor_parallelism=enable_expert_tensor_parallelism),
+                               num_experts=self.num_experts,
+                               ep_size=args.moe_expert_parallel_size,
+                               k=args.topk,
+                               use_residual=(args.mlp_type == 'residual'),
+                               capacity_factor=args.moe_train_capacity_factor,
+                               eval_capacity_factor=args.moe_eval_capacity_factor,
+                               min_capacity=args.moe_min_capacity,
+                               drop_tokens=args.moe_token_dropping,
+                               use_tutel=args.use_tutel,
+                               enable_expert_tensor_parallelism=enable_expert_tensor_parallelism,
+                               top2_2nd_expert_sampling=args.moe_top2_2nd_expert_sampling)
 
         # Set bias+dropout+add fusion grad_enable execution handler.
         TORCH_MAJOR = int(torch.__version__.split('.')[0])