[bug] `apply_quantization_config` target resolution incompatible with vllm

[brian-dellabetta]

Summary

Currently on main, apply_quantization_config calls match_targets to determine which scheme a given target attaches to. Rather than going with the first scheme with a matching target, match_targets tries to prioritize first before deciding (match on exact string BEFORE matching on regex). This means that

recipe = QuantizationModifier(
    config_groups={
        "config_group_0": QuantizationScheme(
            targets=[
                r"re:.*self_attn\.(q|k|v|o)_proj$",
            ],
            **FP8_BLOCK,
        ),
        "config_group_1": QuantizationScheme(
            targets=[
                "model.layers.4.self_attn.o_proj",
            ],
            **W4A16,
        ),
    },
    ignore=["lm_head", "re:.*mlp.gate$"],
)

will yield a safetensors file with an entry model.layers.4.self_attn.o_proj.weight_packed, meaning the W4A16 scheme is selected even though it matches the preceding config_group_0 regex r"re:.*self_attn\.(q|k|v|o)_proj$".

This is NOT how vllm handles scheme resolution -- find_matched_targets resolves the first scheme to match, no prioritization of exact strings over regex. Any model created in this way fails to load in vllm:

(EngineCore pid=288636) ValueError: There is no module or parameter named 'layers.4.self_attn.o_proj.weight_packed' in LlamaModel. The available parameters belonging to layers.4.self_attn.o_proj (RowParallelLinear) are: {'layers.4.self_attn.o_proj.weight', 'layers.4.self_attn.o_proj.weight_scale'}

Full MRE

from transformers import AutoModelForCausalLM, AutoTokenizer
from compressed_tensors.quantization.quant_scheme import (
    FP8_BLOCK,
    W4A16,
    QuantizationScheme,
)

from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier

MODEL_ID = "meta-llama/Meta-Llama-3.1-8B-Instruct"

model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

recipe = QuantizationModifier(
    config_groups={
        "config_group_0": QuantizationScheme(
            targets=[
                r"re:.*self_attn\.(q|k|v|o)_proj$",
            ],
            **FP8_BLOCK,
        ),
        "config_group_1": QuantizationScheme(
            targets=[
                "model.layers.4.self_attn.o_proj",
            ],
            **W4A16,
        ),
    },
    ignore=["lm_head", "re:.*mlp.gate$"],
)

# Apply quantization.
oneshot(model=model, recipe=recipe)

# Save to disk in compressed-tensors format.
SAVE_DIR = MODEL_ID.split("/")[1] + "-FP8-BLOCK-W4A16"
model.save_pretrained(SAVE_DIR)
tokenizer.save_pretrained(SAVE_DIR)

Then run vllm serve Meta-Llama-3.1-8B-Instruct-FP8-BLOCK-W4A16

Resolution

Update match_targets to not do any prioritization. This is an edge case bug that was never triggered by our happy path quant flows, but is leading to user confusion (see #847)

[brian-dellabetta]

Discussed with @kylesayrs that config groups are actually sorted by key during serialization

• Model free ptq
• CT Model Compressor
• transformers to json

so relying on order is equally tenuous. This is currently only an issue for edge cases, when one config group has overlapping targets with another config group, and none of our happy path flows would trigger this. Moreover, a long term solution would require breaking changes to the quantization config in compressed-tensors, (e.g. turning config_groups into a list instead of a dict). For now will leave this open in case others hit this, but no active work underway to resolve

[yashb98]

Thanks @brian-dellabetta. I see why declaration order is not a reliable fix either. On 099fa98, a config-only repro with z_first (8-bit) followed by a_second (4-bit), both targeting the same linear module, selects 4-bit before ModelCompressor.update_config and 8-bit after reloading the saved config. Would an opt-in diagnostic reporting the module and conflicting groups be useful here? I'd keep current matching behavior and avoid reporting groups with equivalent quantization settings. @kylesayrs, I can prepare that if it fits the intended scope.