Varying experts per layer

README.md

@@ -91,6 +91,12 @@ A sample plot for `top2` routing mode (obtained from a tiny toy model) is:
 
 <img src="images/token_count.png" alt="Token Counts" width="70%">
 
+## Varying expert number across layers
+
+To set different number of experts across layers use the flag `--moe-layers` followed by a sequence of numbers corresponding to the number of experts per layer. For example, in a model with 5 layers, one can write `--moe-layers 1 8 16 8 1`.
+
+This flag does not currently support pipeline parallelism. Also, for MoE layers, each of these numbers should be multiple of `--expert-model-parallel-size` and greater than 2. For a dense layer, the number should be set to 1.
+
 # NVIDIA Megatron-LM (copied from upstream)
 
 Megatron ([1](https://arxiv.org/pdf/1909.08053.pdf), [2](https://arxiv.org/pdf/2104.04473.pdf), and [3](https://arxiv.org/pdf/2205.05198)) is a large, powerful transformer developed by the Applied Deep Learning Research team at NVIDIA. This repository is for ongoing research on training large transformer language models at scale. We developed efficient, model-parallel ([tensor](https://arxiv.org/pdf/1909.08053.pdf), [sequence](https://arxiv.org/pdf/2205.05198), and [pipeline](https://arxiv.org/pdf/2104.04473.pdf)), and multi-node pre-training of transformer based models such as [GPT](https://arxiv.org/abs/2005.14165), [BERT](https://arxiv.org/pdf/1810.04805.pdf), and [T5](https://arxiv.org/abs/1910.10683) using mixed precision.

megatron/arguments.py

@@ -390,8 +390,13 @@ def validate_args(args, defaults={}):
     # MoE Spec check
     if args.num_experts is not None:
         assert args.model_spec is None, "Model Spec must be None when using MoEs"
+        assert args.num_experts > 1, "--num-experts should be greater than 2."
     if args.use_balancing_loss is not None:
         assert (args.routing_mode == 'top1' or args.routing_mode == 'top2'), "Need --routing-mode = 'top1' or 'top2' if setting --use-balancing-loss."
+    if args.moe_layers is not None:
+        assert len(args.moe_layers) == args.num_layers, "length of --moe-layers should equal --num-layers."
+        assert min(x for x in args.moe_layers if x != 1) > 2, "Experts per layer should be greater than 2."
+        assert args.use_mcore_models == True, "--moe-layers supported only with --use-mcore-models."
 
     # Expert parallelism check
     if args.expert_model_parallel_size  > 1:
@@ -401,6 +406,8 @@ def validate_args(args, defaults={}):
         if args.tensor_model_parallel_size > 1:
             assert args.sequence_parallel, \
                 "When using expert parallelism and tensor parallelism, sequence parallelism must be used."
+    if args.moe_layers is not None:
+        assert all(x % args.expert_model_parallel_size == 0 for x in args.moe_layers if x != 1), "Experts per layer should be multiple of --expert-model-parallel-size."
 
     # Print arguments.
     _print_args("arguments", args)
@@ -628,6 +635,9 @@ def _add_network_size_args(parser):
                        dest='bert_binary_head')
     group.add_argument('--num-experts', type=int, default=None,
                        help='Number of Experts in Switch Transformer (None means no Switch)')
+    group.add_argument('--moe-layers', nargs='+', type=int, default=None,
+                       help='Number of experts for each layer (`1` means dense layer). '
+                       'Does not support pipeline parallelism.')
     group.add_argument('--routing-mode', type=str, default='sinkhorn',
                        choices=['sinkhorn', 'top1', 'top2', 'sinkhorn_top2'],
                        help='Mode of the expert routing.')

megatron/core/transformer/switch_mlp.py

@@ -50,8 +50,11 @@ def __init__(self, config: TransformerConfig, submodules: MLPSubmodules, layer=N
         args = get_args()
 
         self.config: TransformerConfig = config
-
-        self.router = torch.nn.Linear(self.config.hidden_size, self.config.num_moe_experts)
+        if args.moe_layers:
+            self.num_moe_experts = args.moe_layers[layer-1]
+        else:
+            self.num_moe_experts = self.config.num_moe_experts
+        self.router = torch.nn.Linear(self.config.hidden_size, self.num_moe_experts)
         self.add_bias = config.add_bias_linear
         self.routing = args.routing_mode # 'sinkhorn', 'top1', 'top2', 'sinkhorn_top2'
         self.layer = layer
@@ -61,8 +64,8 @@ def __init__(self, config: TransformerConfig, submodules: MLPSubmodules, layer=N
         self.router_activation = torch.sigmoid
         self.expert_parallel_size = parallel_state.get_expert_model_parallel_world_size()
 
-        assert self.config.num_moe_experts % self.expert_parallel_size == 0
-        self.num_local_experts = self.config.num_moe_experts // self.expert_parallel_size
+        assert self.num_moe_experts % self.expert_parallel_size == 0
+        self.num_local_experts = self.num_moe_experts // self.expert_parallel_size
         local_expert_indices_offset = (
             parallel_state.get_expert_model_parallel_rank() * self.num_local_experts
         )
@@ -97,7 +100,7 @@ def forward(self, hidden_states):
         args = get_args()
         hidden_shape = hidden_states.shape
         route = self.router(hidden_states)
-        route = route.view(-1, self.config.num_moe_experts)
+        route = route.view(-1, self.num_moe_experts)
 
         if self.config.timers is not None:
             self.config.timers('routing_block1', log_level=2).start()
@@ -169,14 +172,14 @@ def forward(self, hidden_states):
         if (args.use_balancing_loss is not None) and self.training:
             if hasattr(args, 'l_aux'):
                 me = torch.mean(route, dim=0)
-                mask1 = F.one_hot(global_indices, num_classes=self.config.num_moe_experts)
+                mask1 = F.one_hot(global_indices, num_classes=self.num_moe_experts)
                 ce = torch.mean(mask1.float(), dim=0)
-                args.l_aux += torch.sum(me * ce) * self.config.num_moe_experts
+                args.l_aux += torch.sum(me * ce) * self.num_moe_experts
                 if self.routing == 'top2':
                     me_2 = torch.mean(masked_route, dim=0)
-                    mask1 = F.one_hot(global_indices_2, num_classes=self.config.num_moe_experts)
+                    mask1 = F.one_hot(global_indices_2, num_classes=self.num_moe_experts)
                     ce_2 = torch.mean(mask1.float(), dim=0)
-                    args.l_aux += torch.sum(me_2 * ce_2) * self.config.num_moe_experts
+                    args.l_aux += torch.sum(me_2 * ce_2) * self.num_moe_experts
 
         # Collect token count for each expert and save to file
         if self.router_profiling_interval and (args.curr_iteration % self.router_profiling_interval == 0) and args.curr_iteration > 0:        
@@ -223,7 +226,6 @@ def forward(self, hidden_states):
         if  self.config.timers is not None:
             self.config.timers('routing_loop').stop()
 
-
         if self.config.timers is not None:
             self.config.timers('ep_scatter', log_level=2).start()
         if self.sequence_parallel or (self.expert_parallel_size > 1):
@@ -262,16 +264,12 @@ def forward(self, hidden_states):
 
         if self.config.timers is not None:
             self.config.timers('final_route', log_level=2).start()
-        output_total = output_total * max_prob
         if self.routing == 'top2' or self.routing == 'sinkhorn_top2':
-            output_total_2 = output_total_2 * max_prob_2
-            output_total = output_total + output_total_2
+            output_total = (output_total * max_prob + output_total_2 * max_prob_2) / (max_prob + max_prob_2)
         output_total = output_total.view(hidden_shape)
         if self.add_bias:
-            output_bias_total = output_bias_total * max_prob
             if self.routing == 'top2' or self.routing == 'sinkhorn_top2':
-                output_bias_total_2 = output_bias_total_2 * max_prob_2
-                output_bias_total = output_bias_total + output_bias_total_2
+                output_bias_total = (output_bias_total * max_prob + output_bias_total_2 * max_prob_2) / (max_prob + max_prob_2)
             output_bias_total = output_bias_total.view(hidden_shape)
         else:
             output_bias_total = None

megatron/core/transformer/transformer_block.py

@@ -14,6 +14,8 @@
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
 from megatron.core.utils import make_sharded_tensor_for_checkpoint, make_viewless_tensor
+from megatron.core.models.gpt.gpt_layer_specs import gpt_layer_with_transformer_engine_spec
+from megatron import get_args
 
 
 class TransformerBlock(MegatronModule):
@@ -57,12 +59,29 @@ def _build_layers(self, transformer_layer_spec):
         #     coeff = self.layer_number
         #     self.norm_factor *= coeff
         def build_layer(layer_number):
-            layer = TransformerLayer(
-                config=self.config,
-                submodules=transformer_layer_spec.submodules,
-                layer_number=layer_number,
-                self_attn_mask_type=self.self_attn_mask_type,
-            )
+            args = get_args()
+            if args.moe_layers:
+                if args.moe_layers[layer_number-1] == 1:
+                    layer = TransformerLayer(
+                        config=self.config,
+                        submodules=gpt_layer_with_transformer_engine_spec.submodules,
+                        layer_number=layer_number,
+                        self_attn_mask_type=self.self_attn_mask_type,
+                    )
+                else:
+                    layer = TransformerLayer(
+                        config=self.config,
+                        submodules=transformer_layer_spec.submodules,
+                        layer_number=layer_number,
+                        self_attn_mask_type=self.self_attn_mask_type,
+                    )
+            else:
+                layer = TransformerLayer(
+                        config=self.config,
+                        submodules=transformer_layer_spec.submodules,
+                        layer_number=layer_number,
+                        self_attn_mask_type=self.self_attn_mask_type,
+                    )
             return layer
 
         if parallel_state.get_virtual_pipeline_model_parallel_world_size() is not None:

megatron/initialize.py

@@ -95,9 +95,9 @@ def finish_mpu_init():
         dir_path = os.path.join(args.router_profiling_path)
         if not os.path.exists(dir_path):
             os.makedirs(dir_path)
-
-        # No continuation function
-        return None
+            
+    # No continuation function
+    return None
 
 
 def _compile_dependencies():

pretrain_gpt.py

@@ -52,7 +52,7 @@ def model_provider(pre_process=True, post_process=True) -> Union[GPTModel, megat
         if args.model_spec is not None:
             transformer_layer_spec = import_module(args.model_spec)
         else:
-            if args.num_experts is None:
+            if (args.num_experts is None) and (args.moe_layers is None):
                 transformer_layer_spec = gpt_layer_with_transformer_engine_spec
             else:
                 transformer_layer_spec = gpt_layer_with_transformer_engine_spec_moe