Move MoE Implementation into src/, add Load Balancing Losses

run_generate.py

@@ -59,6 +59,7 @@ def get_args():
     parser.add_argument("--dp", type=int, default=1)
     parser.add_argument("--pp", type=int, default=0)
     parser.add_argument("--tp", type=int, default=0)
+    parser.add_argument("--ep", type=int, default=0)
     parser.add_argument("--max-new-tokens", type=int, default=128, help="Maximum number of new tokens to generate")
     return parser.parse_args()
 
@@ -79,13 +80,15 @@ def main():
         pp_engine=OneForwardOneBackwardPipelineEngine(),
         tp_mode=TensorParallelLinearMode.ALL_REDUCE,
         tp_linear_async_communication=False,
+        expert_parallel_size=args.ep or config.parallelism.expert_parallel_size,
     )
 
     # Initialise all process groups
     parallel_context = ParallelContext(
         data_parallel_size=parallel_config.dp,
         pipeline_parallel_size=parallel_config.pp,
         tensor_parallel_size=parallel_config.tp,
+        expert_parallel_size=parallel_config.expert_parallel_size,
     )
 
     # Set log levels

src/nanotron/config/config.py

@@ -161,13 +161,15 @@ class GeneralArgs:
 
     Args:
         project: Name of the project (a project gather several runs in common tensorboard/hub-folders)
+        entity: Weights and bias entity name (optional)
         run: Name of the run
         step: Global step (updated when we save the checkpoint)
         consumed_train_samples: Number of samples consumed during training (should be actually just step*batch_size)
         ignore_sanity_checks: Whether to ignore sanity checks
     """
 
     project: str
+    entity: Optional[str] = None
     run: Optional[str] = None
     seed: Optional[int] = None
     step: Optional[int] = None

src/nanotron/config/models_config.py

@@ -51,6 +51,11 @@ class LlamaConfig:
     tie_word_embeddings: bool = False
     use_cache: bool = True
     vocab_size: int = 32000
+    # MoE specific
+    moe_num_experts: int = 1
+    num_experts_per_tok: int = 1
+    moe_loss_weight: float = 0.01
+    moe_z_loss_weight: float = 0.001
 
     def __post_init__(self):
         # NOTE: user don't set self._init_method, ModelArgs will set it

src/nanotron/generation/decode.py

@@ -257,7 +257,12 @@ def decode_text(
                             sharded_logits = model(
                                 input_ids=state.new_input_ids,
                                 input_mask=state.new_input_mask,
-                            )
+                                aux_loss=(
+                                    torch.zeros(1, device=state.new_input_ids.device)
+                                    if is_decoder_input_rank
+                                    else TensorPointer(decoder_input_rank)
+                                ),
+                            )["sharded_logits"]
                     else:
                         if isinstance(state.new_input_ids, torch.Tensor):
                             batch_generated_ids = torch.cat(state.generation_ids, dim=-1)
@@ -268,7 +273,12 @@ def decode_text(
                         sharded_logits = model(
                             input_ids=batch_generated_ids,
                             input_mask=batch_generated_mask,
-                        )
+                            aux_loss=(
+                                torch.zeros(1, device=state.new_input_ids.device)
+                                if is_decoder_input_rank
+                                else TensorPointer(decoder_input_rank)
+                            ),
+                        )["sharded_logits"]
 
                     if isinstance(sharded_logits, torch.Tensor) and logits_are_batch_first:
                         sharded_logits = sharded_logits.transpose(0, 1)