step9: add args to dpa3

Author: iProzd

deepmd/dpmodel/descriptor/dpa3.py

@@ -170,6 +170,14 @@ class RepFlowArgs:
         In the dynamic selection case, neighbor-scale normalization will use `e_sel / sel_reduce_factor`
         or `a_sel / sel_reduce_factor` instead of the raw `e_sel` or `a_sel` values,
         accommodating larger selection numbers.
+    use_moe : bool, optional
+        Whether to use Mixture-of-Experts for the MLP layers in each RepFlowLayer.
+    n_routing_experts : int, optional
+        Total number of routing experts across all GPUs.
+    moe_topk : int, optional
+        Number of experts selected per token.
+    n_shared_experts : int, optional
+        Number of shared experts (replicated on every GPU).
     """
 
     def __init__(
@@ -201,6 +209,10 @@ def __init__(
         use_exp_switch: bool = False,
         use_dynamic_sel: bool = False,
         sel_reduce_factor: float = 10.0,
+        use_moe: bool = False,
+        n_routing_experts: int = 0,
+        moe_topk: int = 0,
+        n_shared_experts: int = 0,
     ) -> None:
         self.n_dim = n_dim
         self.e_dim = e_dim
@@ -231,6 +243,10 @@ def __init__(
         self.use_exp_switch = use_exp_switch
         self.use_dynamic_sel = use_dynamic_sel
         self.sel_reduce_factor = sel_reduce_factor
+        self.use_moe = use_moe
+        self.n_routing_experts = n_routing_experts
+        self.moe_topk = moe_topk
+        self.n_shared_experts = n_shared_experts
 
     def __getitem__(self, key: str) -> Any:
         if hasattr(self, key):
@@ -266,6 +282,10 @@ def serialize(self) -> dict:
             "use_exp_switch": self.use_exp_switch,
             "use_dynamic_sel": self.use_dynamic_sel,
             "sel_reduce_factor": self.sel_reduce_factor,
+            "use_moe": self.use_moe,
+            "n_routing_experts": self.n_routing_experts,
+            "moe_topk": self.moe_topk,
+            "n_shared_experts": self.n_shared_experts,
         }
 
     @classmethod

deepmd/pt/model/descriptor/dpa3.py

@@ -122,6 +122,10 @@ def __init__(
         use_loc_mapping: bool = True,
         type_map: list[str] | None = None,
         add_chg_spin_ebd: bool = False,
+        # MoE EP params (not part of RepFlowArgs, set at runtime).
+        ep_group=None,
+        ep_rank: int = 0,
+        ep_size: int = 1,
     ) -> None:
         super().__init__()
 
@@ -173,6 +177,13 @@ def init_subclass_params(sub_data: Any, sub_class: Any) -> Any:
             precision=precision,
             seed=child_seed(seed, 1),
             trainable=trainable,
+            use_moe=self.repflow_args.use_moe,
+            n_routing_experts=self.repflow_args.n_routing_experts,
+            moe_topk=self.repflow_args.moe_topk,
+            n_shared_experts=self.repflow_args.n_shared_experts,
+            ep_group=ep_group,
+            ep_rank=ep_rank,
+            ep_size=ep_size,
         )
 
         self.use_econf_tebd = use_econf_tebd

deepmd/pt/model/descriptor/repflows.py

@@ -226,6 +226,13 @@ def __init__(
         optim_update: bool = True,
         seed: int | list[int] | None = None,
         trainable: bool = True,
+        use_moe: bool = False,
+        n_routing_experts: int = 0,
+        moe_topk: int = 0,
+        n_shared_experts: int = 0,
+        ep_group=None,
+        ep_rank: int = 0,
+        ep_size: int = 1,
     ) -> None:
         super().__init__()
         self.e_rcut = float(e_rcut)
@@ -256,6 +263,7 @@ def __init__(
         self.a_compress_use_split = a_compress_use_split
         self.use_loc_mapping = use_loc_mapping
         self.optim_update = optim_update
+        self.use_moe = use_moe
         self.smooth_edge_update = smooth_edge_update
         self.edge_init_use_dist = edge_init_use_dist
         self.use_exp_switch = use_exp_switch
@@ -335,6 +343,13 @@ def __init__(
                     smooth_edge_update=self.smooth_edge_update,
                     seed=child_seed(child_seed(seed, 1), ii),
                     trainable=trainable,
+                    use_moe=use_moe,
+                    n_routing_experts=n_routing_experts,
+                    moe_topk=moe_topk,
+                    n_shared_experts=n_shared_experts,
+                    ep_group=ep_group,
+                    ep_rank=ep_rank,
+                    ep_size=ep_size,
                 )
             )
         self.layers = torch.nn.ModuleList(layers)
@@ -656,6 +671,7 @@ def forward(
                 a_sw,
                 edge_index=edge_index,
                 angle_index=angle_index,
+                type_embedding=atype_embd if self.use_moe else None,
             )
 
         # nb x nloc x 3 x e_dim