Merge pull request #2 from dsxsteven/model_register

add gpu_model_register

Author: 845473182

vllm/distributed/eplb/gpu_model_register.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import types
+import typing
+import torch
+from vllm.model_executor.layers.fused_moe import FusedMoE
+from vllm.model_executor.models.utils import is_pp_missing_parameter
+from typing import Callable
+
+def set_eplb_state(
+    self,
+    expert_load_view: torch.Tensor,
+    logical_to_physical_map: torch.Tensor,
+    logical_replica_count: torch.Tensor,
+) -> None:
+    for layer_idx, layer in enumerate(self.moe_layers):
+        # Register the expert weights.
+        self.expert_weights.append(layer.get_expert_weights())
+        layer.set_eplb_state(
+            moe_layer_idx=layer_idx,
+            expert_load_view=expert_load_view,
+            logical_to_physical_map=logical_to_physical_map,
+            logical_replica_count=logical_replica_count,
+        )
+
+def update_physical_experts_metadata(
+    self,
+    num_physical_experts: int,
+    num_local_physical_experts: int,
+) -> None:
+    assert self.num_local_physical_experts == num_local_physical_experts
+    self.num_physical_experts = num_physical_experts
+    self.num_local_physical_experts = num_local_physical_experts
+    self.num_redundant_experts = (num_physical_experts -
+                                    self.num_logical_experts)
+    for layer in self.model.layers:
+        if isinstance(layer.mlp, self.example_moe):
+            moe = layer.mlp
+            moe.n_local_physical_experts = num_local_physical_experts
+            moe.n_physical_experts = num_physical_experts
+            moe.n_redundant_experts = self.num_redundant_experts
+            moe.experts.update_expert_map()
+
+def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
+    # Params for weights, fp8 weight scales, fp8 activation scales
+    # (param_name, weight_name, expert_id, shard_id)
+    return FusedMoE.make_expert_params_mapping(
+        ckpt_gate_proj_name="gate_proj",
+        ckpt_down_proj_name="down_proj",
+        ckpt_up_proj_name="up_proj",
+        num_experts=self.config.n_routed_experts,
+        num_redundant_experts=self.num_redundant_experts)
+
+def load_expert_weight(self, mapping, name, loaded_weight, params_dict):
+    ignore_suffixes = (".bias", "_bias", ".k_scale", "_k_scale",
+                        ".v_scale", "_v_scale", ".weight_scale",
+                        "_weight_scale", ".input_scale", "_input_scale")
+
+    expert_matched = False
+    is_continue = False
+    success = False
+    name_mapped = ''
+    param_name, weight_name, expert_id, shard_id = mapping
+    if weight_name not in name:
+        is_continue = True
+        return expert_matched, is_continue, success, name_mapped
+
+    # Anyway, this is an expert weight and should not be
+    # attempted to load as other weights later
+    expert_matched = True
+
+    # Do not modify `name` since the loop may continue here
+    # Instead, create a new variable
+    name_mapped = name.replace(weight_name, param_name)
+
+    if is_pp_missing_parameter(name_mapped, self):
+        is_continue = True
+        return expert_matched, is_continue, success, name_mapped
+
+    # Skip loading extra parameters for GPTQ/modelopt models.
+    if name_mapped.endswith(ignore_suffixes) \
+            and name_mapped not in params_dict:
+        is_continue = True
+        return expert_matched, is_continue, success, name_mapped
+
+    param = params_dict[name_mapped]
+    # We should ask the weight loader to return success or not
+    # here since otherwise we may skip experts with other
+    # available replicas.
+    weight_loader = typing.cast(Callable[..., bool],
+                                param.weight_loader)
+    success = weight_loader(param,
+                            loaded_weight,
+                            name_mapped,
+                            shard_id=shard_id,
+                            expert_id=expert_id,
+                            return_success=True)
+    return expert_matched, is_continue, success, name_mapped
+
+def model_register(model):
+    """
+    Registers custom methods related to Expert Parallel Load Balancing (EPLB)
+    onto the vLLM model instance. It also determines the number of MoE layers
+    based on the model configuration.
+
+    Args:
+        model: The vLLM model instance to which the methods will be added.
+    """
+    model.set_eplb_state = types.MethodType(set_eplb_state, model)
+    model.load_expert_weight = types.MethodType(load_expert_weight, model)
+    model.update_physical_experts_metadata = \
+        types.MethodType(update_physical_experts_metadata, model) 
+    model.model.get_expert_mapping = \
+        types.MethodType(get_expert_mapping, model.model)
+    print("register complete")

vllm/model_executor/models/deepseek_v2.py

@@ -842,60 +842,26 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.num_expert_groups = config.n_group
 
         self.moe_layers: list[FusedMoE] = []
-        example_moe = None
+        self.example_moe = None
         for layer in self.model.layers:
             if isinstance(layer, PPMissingLayer):
                 continue
 
             assert isinstance(layer, DeepseekV2DecoderLayer)
             if isinstance(layer.mlp, DeepseekV2MoE):
                 # Pick last one layer since the first ones may be dense layers.
-                example_moe = layer.mlp
+                self.example_moe = layer.mlp
                 self.moe_layers.append(layer.mlp.experts)
 
-        if example_moe is None:
+        if self.example_moe is None:
             raise RuntimeError("No DeepseekV2MoE layer found in model.layers.")
 
-        self.num_logical_experts = example_moe.n_logical_experts
-        self.num_physical_experts = example_moe.n_physical_experts
-        self.num_local_physical_experts = example_moe.n_local_physical_experts
-        self.num_routed_experts = example_moe.n_routed_experts
-        self.num_shared_experts = example_moe.n_shared_experts
-        self.num_redundant_experts = example_moe.n_redundant_experts
-
-    def set_eplb_state(
-        self,
-        expert_load_view: torch.Tensor,
-        logical_to_physical_map: torch.Tensor,
-        logical_replica_count: torch.Tensor,
-    ) -> None:
-        for layer_idx, layer in enumerate(self.moe_layers):
-            # Register the expert weights.
-            self.expert_weights.append(layer.get_expert_weights())
-            layer.set_eplb_state(
-                moe_layer_idx=layer_idx,
-                expert_load_view=expert_load_view,
-                logical_to_physical_map=logical_to_physical_map,
-                logical_replica_count=logical_replica_count,
-            )
-
-    def update_physical_experts_metadata(
-        self,
-        num_physical_experts: int,
-        num_local_physical_experts: int,
-    ) -> None:
-        assert self.num_local_physical_experts == num_local_physical_experts
-        self.num_physical_experts = num_physical_experts
-        self.num_local_physical_experts = num_local_physical_experts
-        self.num_redundant_experts = (num_physical_experts -
-                                      self.num_logical_experts)
-        for layer in self.model.layers:
-            if isinstance(layer.mlp, DeepseekV2MoE):
-                moe = layer.mlp
-                moe.n_local_physical_experts = num_local_physical_experts
-                moe.n_physical_experts = num_physical_experts
-                moe.n_redundant_experts = self.num_redundant_experts
-                moe.experts.update_expert_map()
+        self.num_logical_experts = self.example_moe.n_logical_experts
+        self.num_physical_experts = self.example_moe.n_physical_experts
+        self.num_local_physical_experts = self.example_moe.n_local_physical_experts
+        self.num_routed_experts = self.example_moe.n_routed_experts
+        self.num_shared_experts = self.example_moe.n_shared_experts
+        self.num_redundant_experts = self.example_moe.n_redundant_experts
 
     def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.model.get_input_embeddings(input_ids)
@@ -927,16 +893,10 @@ def load_weights(self, weights: Iterable[tuple[str,
             ("fused_qkv_a_proj", "q_a_proj", 0),
             ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
         ]
-
+        from vllm.distributed.eplb.gpu_model_register import get_expert_mapping, load_expert_weight
         # Params for weights, fp8 weight scales, fp8 activation scales
         # (param_name, weight_name, expert_id, shard_id)
-        expert_params_mapping = FusedMoE.make_expert_params_mapping(
-            ckpt_gate_proj_name="gate_proj",
-            ckpt_down_proj_name="down_proj",
-            ckpt_up_proj_name="up_proj",
-            num_experts=self.config.n_routed_experts,
-            num_redundant_experts=self.num_redundant_experts)
-
+        expert_params_mapping = get_expert_mapping(self)
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
         for name, loaded_weight in weights:
@@ -982,34 +942,17 @@ def load_weights(self, weights: Iterable[tuple[str,
                 break
             else:
                 is_expert_weight = False
+                is_continue = False
                 for mapping in expert_params_mapping:
-                    param_name, weight_name, expert_id, shard_id = mapping
-                    if weight_name not in name:
-                        continue
-
-                    # Anyway, this is an expert weight and should not be
-                    # attempted to load as other weights later
-                    is_expert_weight = True
-
-                    # Do not modify `name` since the loop may continue here
-                    # Instead, create a new variable
-                    name_mapped = name.replace(weight_name, param_name)
+                    expert_matched, is_continue, success, name_mapped = \
+                        load_expert_weight(self, mapping, name,
+                                           loaded_weight, params_dict)
+                    if expert_matched:
+                        is_expert_weight = True
 
-                    if is_pp_missing_parameter(name_mapped, self):
+                    if is_continue:
                         continue
 
-                    param = params_dict[name_mapped]
-                    # We should ask the weight loader to return success or not
-                    # here since otherwise we may skip experts with other
-                    # available replicas.
-                    weight_loader = typing.cast(Callable[..., bool],
-                                                param.weight_loader)
-                    success = weight_loader(param,
-                                            loaded_weight,
-                                            name_mapped,
-                                            shard_id=shard_id,
-                                            expert_id=expert_id,
-                                            return_success=True)
                     if success:
                         name = name_mapped
                         break

vllm/model_executor/models/glm4_moe.py

@@ -495,15 +495,6 @@ def make_empty_intermediate_tensors(
                         device=device),
         })
 
-    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
-        # Params for weights, fp8 weight scales, fp8 activation scales
-        # (param_name, weight_name, expert_id, shard_id)
-        return FusedMoE.make_expert_params_mapping(
-            ckpt_gate_proj_name="gate_proj",
-            ckpt_down_proj_name="down_proj",
-            ckpt_up_proj_name="up_proj",
-            num_experts=self.config.n_routed_experts)
-
     def load_weights(self, weights: Iterable[tuple[str,
                                                    torch.Tensor]]) -> set[str]:
         stacked_params_mapping = [
@@ -514,10 +505,10 @@ def load_weights(self, weights: Iterable[tuple[str,
             ("gate_up_proj", "gate_proj", 0),
             ("gate_up_proj", "up_proj", 1),
         ]
-
+        from vllm.distributed.eplb.gpu_model_register import get_expert_mapping, load_expert_weight
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
-        expert_params_mapping = self.get_expert_mapping()
+        expert_params_mapping = get_expert_mapping(self)
         for name, loaded_weight in weights:
             spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
             if spec_layer is not None:
@@ -547,34 +538,17 @@ def load_weights(self, weights: Iterable[tuple[str,
                 break
             else:
                 is_expert_weight = False
+                is_continue = False
                 for mapping in expert_params_mapping:
-                    param_name, weight_name, expert_id, shard_id = mapping
-                    if weight_name not in name:
-                        continue
-
-                    # Anyway, this is an expert weight and should not be
-                    # attempted to load as other weights later
-                    is_expert_weight = True
+                    expert_matched, is_continue, success, name_mapped = \
+                        load_expert_weight(self, mapping, name,
+                                           loaded_weight, params_dict)
+                    if expert_matched:
+                        is_expert_weight = True
 
-                    # Do not modify `name` since the loop may continue here
-                    # Instead, create a new variable
-                    name_mapped = name.replace(weight_name, param_name)
-
-                    if is_pp_missing_parameter(name_mapped, self):
+                    if is_continue:
                         continue
 
-                    param = params_dict[name_mapped]
-                    # We should ask the weight loader to return success or not
-                    # here since otherwise we may skip experts with other
-                    # available replicas.
-                    weight_loader = typing.cast(Callable[..., bool],
-                                                param.weight_loader)
-                    success = weight_loader(param,
-                                            loaded_weight,
-                                            name_mapped,
-                                            shard_id=shard_id,
-                                            expert_id=expert_id,
-                                            return_success=True)
                     if success:
                         name = name_mapped
                         break
@@ -648,42 +622,26 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.num_expert_groups = config.n_group
 
         self.moe_layers: list[FusedMoE] = []
-        example_moe = None
+        self.example_moe = None
         for layer in self.model.layers:
             if isinstance(layer, PPMissingLayer):
                 continue
 
             assert isinstance(layer, Glm4MoeDecoderLayer)
             if isinstance(layer.mlp, Glm4MoE):
                 # Pick last one layer since the first ones may be dense layers.
-                example_moe = layer.mlp
+                self.example_moe = layer.mlp
                 self.moe_layers.append(layer.mlp.experts)
 
-        if example_moe is None:
+        if self.example_moe is None:
             raise RuntimeError("No Glm4MoE layer found in model.layers.")
 
-        self.num_logical_experts = example_moe.n_logical_experts
-        self.num_physical_experts = example_moe.n_physical_experts
-        self.num_local_physical_experts = example_moe.n_local_physical_experts
-        self.num_routed_experts = example_moe.n_routed_experts
-        self.num_shared_experts = example_moe.n_shared_experts
-        self.num_redundant_experts = example_moe.n_redundant_experts
-
-    def set_eplb_state(
-        self,
-        expert_load_view: torch.Tensor,
-        logical_to_physical_map: torch.Tensor,
-        logical_replica_count: torch.Tensor,
-    ) -> None:
-        for layer_idx, layer in enumerate(self.moe_layers):
-            # Register the expert weights.
-            self.expert_weights.append(layer.get_expert_weights())
-            layer.set_eplb_state(
-                moe_layer_idx=layer_idx,
-                expert_load_view=expert_load_view,
-                logical_to_physical_map=logical_to_physical_map,
-                logical_replica_count=logical_replica_count,
-            )
+        self.num_logical_experts = self.example_moe.n_logical_experts
+        self.num_physical_experts = self.example_moe.n_physical_experts
+        self.num_local_physical_experts = self.example_moe.n_local_physical_experts
+        self.num_routed_experts = self.example_moe.n_routed_experts
+        self.num_shared_experts = self.example_moe.n_shared_experts
+        self.num_redundant_experts = self.example_moe.n_redundant_experts
 
     def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.model.get_input_embeddings(input_ids)