Fix the incorrect logic when loading Mixtral series model weights. (#1064)

Author: flyinglandlord

lightllm/models/mixtral/layer_infer/transformer_layer_infer.py

@@ -32,8 +32,8 @@ def _ffn(self, input, infer_state: InferStateInfo, layer_weight: MixtralTransfor
 
         return fused_experts_impl(
             hidden_states=hidden_states,
-            w1=layer_weight.experts.w1,
-            w2=layer_weight.experts.w2,
+            w1=layer_weight.experts.w1[0],
+            w2=layer_weight.experts.w2[0],
             topk_weights=topk_weights,
             topk_ids=topk_ids,
             inplace=True,

lightllm/models/mixtral/layer_weights/transformer_layer_weight.py

@@ -1,3 +1,4 @@
+import os
 from lightllm.utils.log_utils import init_logger
 from lightllm.utils.envs_utils import enable_env_vars
 from lightllm.models.llama.layer_weights.transformer_layer_weight import LlamaTransformerLayerWeight
@@ -30,10 +31,10 @@ def _init_weight_names(self):
         self.moe_gate_weight_name = f"model.layers.{self.layer_num_}.block_sparse_moe.gate.weight"
         self.moe_gate_bias_name = None
 
-    def _init_ffn(self, weights):
-        self._init_moe(weights)
+    def _init_ffn(self):
+        self._init_moe()
 
-    def _init_moe(self, weights):
+    def _init_moe(self):
         inter_size = self.network_config_["intermediate_size"]
         split_inter_size = inter_size // self.tp_world_size_
 
@@ -45,16 +46,26 @@ def _init_moe(self, weights):
             layer_num=self.layer_num_,
             name="moe_gate",
             tp_rank=0,
-            tp_size=1,  # no tensor parallelism
+            tp_world_size=1,  # no tensor parallelism
         )
 
-        load_func = FusedMoeWeightEP if enable_env_vars("ETP_MODE_ENABLED") else FusedMoeWeightTP
-        self.experts = load_func(
-            gate_proj_name="w1",
-            down_proj_name="w2",
-            up_proj_name="w3",
-            weight_prefix=f"model.layers.{self.layer_num_}.block_sparse_moe.experts",
-            n_routed_experts=self.n_routed_experts,
-            split_inter_size=split_inter_size,
-            data_type=self.data_type_,
-        )
+        moe_mode = os.getenv("MOE_MODE", "TP")
+        assert moe_mode in ["TP"], f"Unsupported moe mode: {moe_mode}"
+
+        if moe_mode == "TP":
+            self.experts = FusedMoeWeightTP(
+                gate_proj_name="w1",
+                down_proj_name="w2",
+                up_proj_name="w3",
+                e_score_correction_bias_name="",
+                weight_prefix=f"model.layers.{self.layer_num_}.block_sparse_moe.experts",
+                n_routed_experts=self.n_routed_experts,
+                split_inter_size=split_inter_size,
+                data_type=self.data_type_,
+                network_config=self.network_config_,
+                layer_num=self.layer_num_,
+                quant_cfg=self.quant_cfg,
+                num_fused_shared_experts=0,
+            )
+        else:
+            raise ValueError(f"Unsupported moe mode: {moe_mode}")