|
|
@@ -1,58 +0,0 @@
|
|
|
-diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py
|
|
|
-index 15f6c5ed4..1d411ecc7 100644
|
|
|
---- a/vllm/model_executor/layers/fused_moe/routed_experts.py
|
|
|
-+++ b/vllm/model_executor/layers/fused_moe/routed_experts.py
|
|
|
-@@ -36,6 +36,28 @@ if TYPE_CHECKING:
|
|
|
- logger = init_logger(__name__)
|
|
|
-
|
|
|
-
|
|
|
-+def _index_expert_mapping(
|
|
|
-+ mapping: list[tuple[str, str, int, str]],
|
|
|
-+) -> dict[str, list[tuple[str, str, int, str]]]:
|
|
|
-+ """Index by logical checkpoint ID, preserving entry order and EPLB replicas.
|
|
|
-+
|
|
|
-+ An empty index keeps the full scan for unsupported mapping names.
|
|
|
-+ """
|
|
|
-+ mapping_by_expert: dict[str, list[tuple[str, str, int, str]]] = {}
|
|
|
-+ for entry in mapping:
|
|
|
-+ prefix, _, suffix = entry[1].partition(".")
|
|
|
-+ expert_key, sep, _ = suffix.partition(".")
|
|
|
-+ if (
|
|
|
-+ prefix != "experts"
|
|
|
-+ or not expert_key
|
|
|
-+ or (expert_key.isdecimal() and not sep)
|
|
|
-+ ):
|
|
|
-+ return {}
|
|
|
-+ if expert_key.isdecimal():
|
|
|
-+ mapping_by_expert.setdefault(expert_key, []).append(entry)
|
|
|
-+ return mapping_by_expert
|
|
|
-+
|
|
|
-+
|
|
|
- class FusedMoeWeightScaleSupported(Enum):
|
|
|
- TENSOR = "tensor"
|
|
|
- CHANNEL = "channel"
|
|
|
-@@ -889,14 +911,22 @@ class RoutedExperts(PluggableLayer):
|
|
|
- self, weights: Iterable[tuple[str, torch.Tensor]]
|
|
|
- ) -> Iterable[str]:
|
|
|
- expert_mapping = self.get_expert_mapping(include_fused=True)
|
|
|
-+ mapping_by_expert = _index_expert_mapping(expert_mapping)
|
|
|
-+
|
|
|
- for expert_name, loaded_weight in weights:
|
|
|
- qual_name = f"{self.layer_name}.{expert_name}"
|
|
|
-- # Fused expert weights can be identified by their 3D tensors
|
|
|
- is_fused = loaded_weight.dim() == 3
|
|
|
-+ # Fused tensors and ambiguous names keep the full mapping.
|
|
|
-+ candidates = expert_mapping
|
|
|
-+ if not is_fused and qual_name.count("experts.") == 1:
|
|
|
-+ expert_key = qual_name.partition("experts.")[2].partition(".")[0]
|
|
|
-+ candidates = mapping_by_expert.get(expert_key, expert_mapping)
|
|
|
-+
|
|
|
- matched = False
|
|
|
-- for param_name, weight_name, expert_id, shard_id in expert_mapping:
|
|
|
-+ for param_name, weight_name, expert_id, shard_id in candidates:
|
|
|
- if weight_name not in qual_name:
|
|
|
- if matched and is_fused:
|
|
|
-+ # Fused tensors use the first contiguous run of matches.
|
|
|
- break
|
|
|
- continue
|
|
|
- matched = True
|
