Skip to content

vllm.model_executor.layers.quantization.compressed_tensors.utils

Functions:

_find_first_match(value, targets, check_contains=False)

Returns first element of target that matches value either exactly or as a regex after 're:'. If check_contains is set to True, additionally checks if the target string is contained within the value.

Parameters:

  • value

    (str) –

    string to compare the list of targets against

  • targets

    (Iterable[str]) –

    list of targets to match the layer against

  • check_contains

    (bool, default: False ) –

    whether or not to do a substring match

Source code in vllm/model_executor/layers/quantization/compressed_tensors/utils.py 
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
def _find_first_match(
    value: str, targets: Iterable[str], check_contains: bool = False
) -> str | None:
    """Returns first element of target that matches value either
    exactly or as a regex after 're:'. If check_contains is set to True,
    additionally checks if the target string is contained within the value.

    Args:
        value: string to compare the list of targets against
        targets: list of targets to match the layer against
        check_contains: whether or not to do a substring match

    """
    for target in targets:
        if is_equal_or_regex_match(value, target, check_contains=check_contains):
            return target
    return None

_match_fused_layer(layer_name, target_layers, fused_mapping)

Match a fused layer name to its corresponding individual layer in target_layers. Returns first value in fused_mapping which matches targets

Implements an "all" matching strategy where a fused layer matches iff "all" of its components match

Parameters:

  • layer_name

    (str) –

    layer name

  • target_layers

    (Iterable[str]) –

    list of targets to match the layer against

  • fused_mapping

    (Mapping[str, list[str]]) –

    map from fused layer names to its components

Examples:

layer_name = "model.layers.0.self_attn.qkv_proj" target_layers = ["model.layers.0.self_attn.q_proj", "model.layers.0.self_attn.k_proj", "model.layers.0.self_attn.v_proj"]

Source code in vllm/model_executor/layers/quantization/compressed_tensors/utils.py 
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
def _match_fused_layer(
    layer_name: str,
    target_layers: Iterable[str],
    fused_mapping: Mapping[str, list[str]],
) -> str | None:
    """Match a fused layer name to its corresponding individual layer in
    target_layers. Returns first value in fused_mapping which matches targets

    Implements an "all" matching strategy where a fused layer matches iff
    "all" of its components match

    Args:
        layer_name: layer name
        target_layers: list of targets to match the layer against
        fused_mapping: map from fused layer names to its components

    Examples:
        layer_name = "model.layers.0.self_attn.qkv_proj"
        target_layers = ["model.layers.0.self_attn.q_proj",
                        "model.layers.0.self_attn.k_proj",
                        "model.layers.0.self_attn.v_proj"]

    """
    # find layer_name in mapping
    fused = next((key for key in fused_mapping if layer_name.endswith(key)), None)
    if fused is None:
        return None

    # expand path of unfused components
    unfused_paths = [
        layer_name.replace(fused, unfused) for unfused in fused_mapping[fused]
    ]

    # for each unfused component, find a match in targets
    unfused_matches: list[str | None] = []
    for unfused in unfused_paths:
        for target in target_layers:
            if is_equal_or_regex_match(unfused, target):
                unfused_matches.append(target)
                break
        else:
            unfused_matches.append(None)

    return unfused_matches[0] if all(unfused_matches) else None

find_matched_target(layer_name, module, targets, fused_mapping=MappingProxyType({}))

Helper function to look up which "target" in the compressed-tensors config that a layer corresponds to.

Recall that a compressed-tensors configs has a concept of config_groups, where each layer can be quantized with a different scheme.

targets in each config_group will be a list of either layer names (or regexes corresponding to layer names) or names of torch Modules.

First, we try to match the layer_name with a target Second, we try to match the module's name with a target Third, we try to map the layer_name to a list of fused module names. All component module names must match in order for a match to be successful. A successful match returns the first component target

Parameters:

Source code in vllm/model_executor/layers/quantization/compressed_tensors/utils.py 
 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
110
111
def find_matched_target(
    layer_name: str | None,
    module: Module,
    targets: Iterable[str],
    fused_mapping: Mapping[str, list[str]] = MappingProxyType({}),
) -> str | None:
    """Helper function to look up which "target" in the compressed-tensors
    config that a layer corresponds to.

    Recall that a compressed-tensors configs has a concept of
    config_groups, where each layer can be quantized with a different
    scheme.

    targets in each config_group will be a list of either layer names
    (or regexes corresponding to layer names) or names of torch Modules.

    First, we try to match the layer_name with a target
    Second, we try to match the module's name with a target
    Third, we try to map the layer_name to a list of fused module names.
        *All* component module names must match in order for a match to be
        successful. A successful match returns the first component target

    Args:
        layer_name: layer name
        module: torch.nn.Module
        targets: list of targets to match the layer against
        fused_mapping: map from fused layer names to its components

    """
    if layer_name is None:
        layer_name = ""

    matched_target = (
        _find_first_match(layer_name, targets)
        or _match_fused_layer(layer_name, targets, fused_mapping)
        or _find_first_match(module.__class__.__name__, targets, True)
    )

    return matched_target