Merge pull request #371 from malay-nagda/malay/token_drop

Malay/token drop
NVIDIA · Jul 9, 2024 · 74d8409 · 74d8409
2 parents 70278f9 + b885571
commit 74d8409
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/launcher_scripts/conf/training/mixtral/mixtral_8x3b.yaml b/launcher_scripts/conf/training/mixtral/mixtral_8x3b.yaml
@@ -47,6 +47,9 @@ exp_manager:
 model:
   mcore_gpt: true
   moe_grouped_gemm: true
+  moe_token_dispatcher_type: alltoall
+  moe_pad_expert_input_to_capacity: True
+  moe_expert_capacity_factor: 1.0
   micro_batch_size: 1
   global_batch_size: 128
   rampup_batch_size: null

diff --git a/launcher_scripts/conf/training/mixtral/mixtral_8x7b.yaml b/launcher_scripts/conf/training/mixtral/mixtral_8x7b.yaml
@@ -48,6 +48,8 @@ model:
   mcore_gpt: true
   moe_grouped_gemm: true
   moe_token_dispatcher_type: alltoall
+  moe_pad_expert_input_to_capacity: True
+  moe_expert_capacity_factor: 1.0
   moe_aux_loss_coeff: 0.01
   micro_batch_size: 1
   global_batch_size: 256