Skip to content

vllm.model_executor.models.expert_substitution

Functions:

_parse_approx_value_name(name)

Return (layer index, expert ID) of an approx_value tensor name.

Source code in vllm/model_executor/models/expert_substitution.py 
38
39
40
41
42
43
44
45
46
def _parse_approx_value_name(name: str) -> tuple[int, int] | None:
    """Return (layer index, expert ID) of an ``approx_value`` tensor name."""
    if not name.endswith(APPROX_VALUE_SUFFIX):
        return None
    experts_path, _, expert_id = name.removesuffix(APPROX_VALUE_SUFFIX).rpartition(".")
    layer_idx = decoder_layer_index(experts_path)
    if layer_idx is None or not expert_id.isdigit():
        return None
    return layer_idx, int(expert_id)

as_expert_substitution_model(model_cls, model_config)

Load constant expert values at the load_weights boundary.

Initial loading and direct weight updates both go through load_weights. Explicitly named values are always consumed and loaded if their layer is built by this model. approx_value tensors of layers built elsewhere (other pipeline stages, MTP layers) are passed through, so the model's own loader skips them like any other weight of a layer it does not own.

Source code in vllm/model_executor/models/expert_substitution.py 
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def as_expert_substitution_model(model_cls: type[_T], model_config: Any) -> type[_T]:
    """Load constant expert values at the ``load_weights`` boundary.

    Initial loading and direct weight updates both go through ``load_weights``.
    Explicitly named values are always consumed and loaded if their layer is
    built by this model. ``approx_value`` tensors of layers built elsewhere
    (other pipeline stages, MTP layers) are passed through, so the model's own
    loader skips them like any other weight of a layer it does not own.
    """
    spec = get_expert_substitution_spec(model_config)
    if spec is None:
        return model_cls
    value_names = spec.value_names

    class ExpertSubstitutionModel(model_cls):  # type: ignore[valid-type, misc]
        @wraps(model_cls.__init__)  # type: ignore[misc]
        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            self._expert_substitutions = _collect_expert_substitutions(self)

        def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
            if any(s.values.is_meta for _, s in self._expert_substitutions.values()):
                # Put the previous weights back so the model stays usable.
                finalize_layerwise_reload(self, model_config)
                raise NotImplementedError(
                    "expert substitution does not support layerwise weight reload"
                )
            loaded_substitutions: set[str] = set()

            def remaining_weights():
                for name, loaded_weight in weights:
                    targets = value_names.get(name)
                    if targets is None:
                        key = _parse_approx_value_name(name)
                        if key is None or key[0] not in self._expert_substitutions:
                            yield name, loaded_weight
                            continue
                        targets = (key,)
                    for layer_idx, expert_id in targets:
                        if layer_idx not in self._expert_substitutions:
                            continue
                        param_name, substitution = self._expert_substitutions[layer_idx]
                        param = substitution.values
                        param.weight_loader(param, loaded_weight, expert_id)
                        loaded_substitutions.add(param_name)

            loaded = super().load_weights(remaining_weights())
            if loaded is not None:
                loaded.update(loaded_substitutions)
            return loaded

        def process_weights_after_loading(self) -> None:
            for param_name, substitution in self._expert_substitutions.values():
                substitution.validate_loaded_values(param_name)
            finalize = getattr(super(), "process_weights_after_loading", None)
            if finalize is not None:
                finalize()

    ExpertSubstitutionModel.__name__ = model_cls.__name__
    ExpertSubstitutionModel.__qualname__ = model_cls.__qualname__
    return ExpertSubstitutionModel