llm-compressor/tests/llmcompressor/transformers/tracing/models.py at 0e074e4fa4b4fbfbc8f830eef42138bd5ef214dc · vllm-project/llm-compressor · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
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
import pytest
from transformers import AutoModelForCausalLM

from llmcompressor.transformers.tracing import (
    TraceableIdefics3ForConditionalGeneration,
    TraceableLlavaForConditionalGeneration,
    TraceableMllamaForConditionalGeneration,
    TraceableQwen2_5_VLForConditionalGeneration,
    TraceableQwen2VLForConditionalGeneration,
    TraceableWhisperForConditionalGeneration,
)
from llmcompressor.transformers.tracing.debug import trace


@pytest.mark.parametrize(
    "model_id,model_class,targets",
    [
        ("meta-llama/Meta-Llama-3-8B-Instruct", AutoModelForCausalLM, None),
    ],
)
def test_text_trace(model_id, model_class, targets):
    trace(
        model_id,
        model_class,
        targets,
        ignore=[],
        modality="text",
        trust_remote_code=True,
    )


@pytest.mark.parametrize(
    "model_id,model_class,targets,ignore",
    [
        (
            "Qwen/Qwen2-VL-2B-Instruct",
            TraceableQwen2VLForConditionalGeneration,
            None,
            ["lm_head", "re:visual.*"],
        ),
        (
            "Qwen/Qwen2.5-VL-7B-Instruct",
            TraceableQwen2_5_VLForConditionalGeneration,
            None,
            ["lm_head", "re:visual.*"],
        ),
        (
            "mgoin/pixtral-12b",
            TraceableLlavaForConditionalGeneration,
            ["MistralDecoderLayer"],
            ["re:.*lm_head", "re:vision_tower.*", "re:multi_modal_projector.*"],
        ),
        (
            "meta-llama/Llama-3.2-11B-Vision-Instruct",
            TraceableMllamaForConditionalGeneration,
            None,
            ["re:.*lm_head", "re:multi_modal_projector.*", "re:vision_model.*"],
        ),
        (
            "llava-hf/llava-1.5-7b-hf",
            TraceableLlavaForConditionalGeneration,
            ["LlamaDecoderLayer"],
            ["re:.*lm_head", "re:vision_tower.*", "re:multi_modal_projector.*"],
        ),
        (
            "HuggingFaceM4/Idefics3-8B-Llama3",
            TraceableIdefics3ForConditionalGeneration,
            ["Idefics3EncoderLayer", "LlamaDecoderLayer"],
            ["re:.*lm_head", "re:model.vision_model.*", "re:model.connector.*"],
        ),
    ],
)
def test_vision_trace(model_id, model_class, targets, ignore):
    trace(
        model_id,
        model_class,
        targets,
        ignore=ignore,
        modality="vision",
        trust_remote_code=True,
    )


@pytest.mark.parametrize(
    "model_id,model_class,targets,ignore",
    [
        ("openai/whisper-large-v3", TraceableWhisperForConditionalGeneration, None, []),
    ],
)
def test_audio_trace(model_id, model_class, targets, ignore):
    trace(
        model_id,
        model_class,
        targets,
        ignore=ignore,
        modality="audio",
        trust_remote_code=True,
    )