vllm-58720-routed-experts.patch 2.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758
  1. diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py
  2. index 15f6c5ed4..1d411ecc7 100644
  3. --- a/vllm/model_executor/layers/fused_moe/routed_experts.py
  4. +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py
  5. @@ -36,6 +36,28 @@ if TYPE_CHECKING:
  6. logger = init_logger(__name__)
  7. +def _index_expert_mapping(
  8. + mapping: list[tuple[str, str, int, str]],
  9. +) -> dict[str, list[tuple[str, str, int, str]]]:
  10. + """Index by logical checkpoint ID, preserving entry order and EPLB replicas.
  11. +
  12. + An empty index keeps the full scan for unsupported mapping names.
  13. + """
  14. + mapping_by_expert: dict[str, list[tuple[str, str, int, str]]] = {}
  15. + for entry in mapping:
  16. + prefix, _, suffix = entry[1].partition(".")
  17. + expert_key, sep, _ = suffix.partition(".")
  18. + if (
  19. + prefix != "experts"
  20. + or not expert_key
  21. + or (expert_key.isdecimal() and not sep)
  22. + ):
  23. + return {}
  24. + if expert_key.isdecimal():
  25. + mapping_by_expert.setdefault(expert_key, []).append(entry)
  26. + return mapping_by_expert
  27. +
  28. +
  29. class FusedMoeWeightScaleSupported(Enum):
  30. TENSOR = "tensor"
  31. CHANNEL = "channel"
  32. @@ -889,14 +911,22 @@ class RoutedExperts(PluggableLayer):
  33. self, weights: Iterable[tuple[str, torch.Tensor]]
  34. ) -> Iterable[str]:
  35. expert_mapping = self.get_expert_mapping(include_fused=True)
  36. + mapping_by_expert = _index_expert_mapping(expert_mapping)
  37. +
  38. for expert_name, loaded_weight in weights:
  39. qual_name = f"{self.layer_name}.{expert_name}"
  40. - # Fused expert weights can be identified by their 3D tensors
  41. is_fused = loaded_weight.dim() == 3
  42. + # Fused tensors and ambiguous names keep the full mapping.
  43. + candidates = expert_mapping
  44. + if not is_fused and qual_name.count("experts.") == 1:
  45. + expert_key = qual_name.partition("experts.")[2].partition(".")[0]
  46. + candidates = mapping_by_expert.get(expert_key, expert_mapping)
  47. +
  48. matched = False
  49. - for param_name, weight_name, expert_id, shard_id in expert_mapping:
  50. + for param_name, weight_name, expert_id, shard_id in candidates:
  51. if weight_name not in qual_name:
  52. if matched and is_fused:
  53. + # Fused tensors use the first contiguous run of matches.
  54. break
  55. continue
  56. matched = True