update common perf arguments - ce fusion - moe gemms

Author: vidushi8

examples/megatron/configs/deepseek_v2_lite-pretrain.yaml

@@ -57,7 +57,7 @@ modules:
       # fused wgrad gemm and accumulation
       gradient_accumulation_fusion: false
       # recommend set `false` in fp8
-      moe_use_legacy_grouped_gemm: true
+      moe_use_legacy_grouped_gemm: false
       # fused topk router with aux score
       moe_use_fused_router_with_aux_score: false
       # pad 192/128 for deepseek attention
@@ -82,4 +82,8 @@ modules:
       # Turbo
       enable_primus_turbo: true
       use_turbo_attention: true
-      use_turbo_grouped_mlp: true
+      use_turbo_grouped_mlp: true
+
+      # Cross entropy flags
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/deepseek_v3-pretrain.yaml

@@ -57,7 +57,7 @@ modules:
       # fused wgrad gemm and accumulation
       gradient_accumulation_fusion: false
       # recommend set `false` in fp8
-      moe_use_legacy_grouped_gemm: true
+      moe_use_legacy_grouped_gemm: false
       # fused topk router with aux score
       moe_use_fused_router_with_aux_score: false
       # pad 192/128 for deepseek attention
@@ -85,3 +85,7 @@ modules:
       enable_primus_turbo: true
       use_turbo_attention: true
       use_turbo_grouped_mlp: true
+
+      # Cross entropy flags
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/llama2_70B-pretrain.yaml

@@ -75,3 +75,7 @@ modules:
       enable_primus_turbo: true
       use_turbo_attention: true
       use_turbo_grouped_mlp: true
+
+      # Cross entropy flags
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/llama2_7B-pretrain.yaml

@@ -78,3 +78,7 @@ modules:
       # overlap_param_gather: false
       # ckpt_format: torch
       # sequence_parallel: 1
+
+      # Cross entropy flags
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/llama3.1_70B-pretrain.yaml

@@ -74,3 +74,7 @@ modules:
       # Turbo
       enable_primus_turbo: true
       use_turbo_attention: true
+
+      # Cross entropy flags
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/llama3.1_8B-pretrain.yaml

@@ -65,3 +65,7 @@ modules:
       no_save_rng: null
       disable_last_saving: true
       ckpt_format: torch
+
+      # Cross entropy flags
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/llama3.3_70B-pretrain.yaml

@@ -75,3 +75,7 @@ modules:
       enable_primus_turbo: true
       use_turbo_attention: true
       use_turbo_grouped_mlp: true
+
+      # Cross entropy flags
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/llama3_70B-pretrain.yaml

@@ -75,3 +75,7 @@ modules:
       enable_primus_turbo: true
       use_turbo_attention: true
       use_turbo_grouped_mlp: true
+
+      # Cross entropy flags
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/llama3_8B-pretrain.yaml

@@ -71,3 +71,7 @@ modules:
       enable_primus_turbo: true
       use_turbo_attention: true
       use_turbo_grouped_mlp: true
+
+      # Cross entropy flags
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/mixtral_8x22B_v0.1-pretrain.yaml

@@ -59,7 +59,7 @@ modules:
       # fusion
       moe_permute_fusion: false
       gradient_accumulation_fusion: false
-      moe_use_legacy_grouped_gemm: true
+      moe_use_legacy_grouped_gemm: false
 
       # ckpt
       finetune: false
@@ -73,3 +73,7 @@ modules:
       no_save_rng: null
       disable_last_saving: true
       ckpt_format: torch
+
+      # Cross entropy flags
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true