use torchtitan moe impl

modeling_deepseek.py

@@ -59,6 +59,8 @@ from .configuration_deepseek import DeepseekV3Config
 import torch.distributed as dist
 import numpy as np
 
+from torchtitan.models.moe import MoE, MoEArgs
+
 if is_flash_attn_2_available():
     from flash_attn import flash_attn_func, flash_attn_varlen_func
     from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
@@ -1150,8 +1152,20 @@ class DeepseekV3DecoderLayer(nn.Module):
             config=config, layer_idx=layer_idx
         )
 
+        moe_args = MoEArgs(
+            num_experts=config.n_routed_experts,
+            num_shared_experts=config.n_shared_experts,
+            score_func=config.scoring_func,
+            route_norm=config.norm_topk_prob,
+            route_scale=config.routed_scaling_factor,
+            score_before_experts=False,
+            top_k=config.num_experts_per_tok,
+            use_grouped_mm=True,
+            load_balance_coeff=1e-3,
+        )
+
         self.mlp = (
-            DeepseekV3MoE(config)
+            MoE(moe_args, dim=config.hidden_size, hidden_dim=config.moe_intermediate_size)
             if (
                 config.n_routed_experts is not None
                 and layer_idx >= config.first_k_dense_replace