A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 23:30 UTC
[Model] Enable quantization support for `transformers` backend (#12960)
Isotr0py committed
Feb 18, 2025 at 11:52 UTC
67ef8f666aa5d681b9ab2c126caf2b0d8596074a
3 files changed
+66
-23
docs/source/models/supported_models.md
+7
-3
@@ -42,7 +42,7 @@ Alternatively, you can [open an issue on GitHub](https://github.com/vllm-project
42
43
### Transformers fallback
44
45
-After the merge of <gh-pr:11330>, `vllm` can fallback to models that are available in `transformers`. This does not work for all models for now, but most decoder language models are supported, and vision language model support is planned!
45
+`vllm` can fallback to models that are available in `transformers`. This does not work for all models for now, but most decoder language models are supported, and vision language model support is planned!
46
47
To check if the backend is `transformers`, you can simply do this:
48
@@ -56,9 +56,13 @@ If it is `TransformersModel` then it means it's based on `transformers`!
56
57
#### Supported features
58
59
-##### LORA and quantization
59
+##### Quantization
60
61
-Both are not supported yet! Make sure to open an issue and we'll work on this together with the `transformers` team!
61
+Transformers fallback has supported most of available quantization in vLLM (except GGUF). See [Quantization page](#quantization-index) for more information about supported quantization in vllm.
62
+
63
+##### LoRA
64
+
65
+LoRA hasn't supported on transformers fallback yet! Make sure to open an issue and we'll work on this together with the `transformers` team!
66
67
Usually `transformers` model load weights via the `load_adapters` API, that depends on PEFT. We need to work a bit to either use this api (for now this would result in some weights not being marked as loaded) or replace modules accordingly.
68
tests/models/test_transformers.py
+49
-5
@@ -45,10 +45,14 @@ def check_implementation(
45
("meta-llama/Llama-3.2-1B-Instruct", "transformers"),
46
("openai-community/gpt2", "transformers"),
47
("ArthurZ/Ilama-3.2-1B", "auto"), # CUSTOM CODE
48
- ("meta-llama/Llama-3.2-1B-Instruct", "auto"),
48
]) # trust_remote_code=True by default
50
-def test_models(hf_runner, vllm_runner, example_prompts, model,
51
- model_impl) -> None:
49
+def test_models(
50
+ hf_runner: Type[HfRunner],
51
+ vllm_runner: Type[VllmRunner],
52
+ example_prompts: list[str],
53
+ model: str,
54
+ model_impl: str,
55
+) -> None:
56
57
maybe_raises = nullcontext()
58
if model == "openai-community/gpt2" and model_impl == "transformers":
@@ -67,10 +71,50 @@ def test_models(hf_runner, vllm_runner, example_prompts, model,
71
72
@multi_gpu_test(num_gpus=2)
73
def test_distributed(
70
- hf_runner,
71
- vllm_runner,
74
+ hf_runner: Type[HfRunner],
75
+ vllm_runner: Type[VllmRunner],
76
example_prompts,
77
):
78
kwargs = {"model_impl": "transformers", "tensor_parallel_size": 2}
79
check_implementation(hf_runner, vllm_runner, example_prompts,
80
"meta-llama/Llama-3.2-1B-Instruct", **kwargs)
81
+
82
+
83
+@pytest.mark.parametrize("model, quantization_kwargs", [
84
+ (
85
+ "meta-llama/Llama-3.2-1B-Instruct",
86
+ {
87
+ "quantization": "bitsandbytes",
88
+ "load_format": "bitsandbytes",
89
+ },
90
+ ),
91
+])
92
+@pytest.mark.parametrize("max_tokens", [32])
93
+@pytest.mark.parametrize("num_logprobs", [5])
94
+def test_quantization(
95
+ vllm_runner: Type[VllmRunner],
96
+ example_prompts: list[str],
97
+ model: str,
98
+ quantization_kwargs: dict[str, str],
99
+ max_tokens: int,
100
+ num_logprobs: int,
101
+) -> None:
102
+ with vllm_runner(
103
+ model, model_impl="auto", enforce_eager=True,
104
+ **quantization_kwargs) as vllm_model: # type: ignore[arg-type]
105
+ vllm_outputs = vllm_model.generate_greedy_logprobs(
106
+ example_prompts, max_tokens=max_tokens, num_logprobs=num_logprobs)
107
+
108
+ with vllm_runner(
109
+ model,
110
+ model_impl="transformers",
111
+ enforce_eager=True,
112
+ **quantization_kwargs) as vllm_model: # type: ignore[arg-type]
113
+ transformers_outputs = vllm_model.generate_greedy_logprobs(
114
+ example_prompts, max_tokens=max_tokens, num_logprobs=num_logprobs)
115
+ check_logprobs_close(
116
+ outputs_0_lst=transformers_outputs,
117
+ outputs_1_lst=vllm_outputs,
118
+ name_0="transformers",
119
+ name_1="vllm",
120
+ )
vllm/model_executor/models/transformers.py
+10
-15
@@ -28,6 +28,7 @@ from vllm.distributed import get_tensor_model_parallel_world_size
28
from vllm.distributed.utils import divide
29
from vllm.logger import init_logger
30
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
31
+ ReplicatedLinear,
32
RowParallelLinear)
33
from vllm.model_executor.layers.logits_processor import LogitsProcessor
34
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
@@ -37,6 +38,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
38
from vllm.model_executor.sampling_metadata import SamplingMetadata
39
from vllm.sequence import IntermediateTensors
40
41
+from .interfaces import SupportsQuant
42
from .utils import maybe_prefix
43
44
logger = init_logger(__name__)
@@ -50,10 +52,10 @@ def vllm_flash_attention_forward(
52
value: torch.Tensor,
53
attention_mask: torch.Tensor,
54
# Transformers kwargs
53
- scaling: float = None,
55
+ scaling: Optional[float] = None,
56
# vLLM kwargs
55
- attn_metadata: AttentionMetadata = None,
56
- attention_instances: list[Attention] = None,
57
+ attn_metadata: Optional[AttentionMetadata] = None,
58
+ attention_instances: Optional[list[Attention]] = None,
59
**kwargs):
60
self_attn = attention_instances[module.layer_idx]
61
if scaling is not None:
@@ -99,13 +101,7 @@ def replace_linear_class(
101
vllm_linear_cls = {
102
"colwise": ColumnParallelLinear,
103
"rowwise": RowParallelLinear,
102
- }.get(style)
103
-
104
- if vllm_linear_cls is None:
105
- logger.warning(
106
- "Unsupported parallel style value: %s. "
107
- "This layer will not be tensor parallelized.", style)
108
- return linear
104
+ }.get(style, ReplicatedLinear)
105
106
class HFCompatibleLinear(vllm_linear_cls):
107
"""
@@ -119,10 +115,11 @@ def replace_linear_class(
115
input_size=linear.in_features,
116
output_size=linear.out_features,
117
bias=linear.bias is not None,
118
+ quant_config=quant_config,
119
)
120
121
125
-class TransformersModel(nn.Module):
122
+class TransformersModel(nn.Module, SupportsQuant):
123
embedding_padding_modules = ["lm_head"]
124
embedding_modules = ["embed_tokens"
125
] # TODO transformers will have a util to get it
@@ -133,10 +130,8 @@ class TransformersModel(nn.Module):
130
131
config = vllm_config.model_config.hf_config
132
cache_config = vllm_config.cache_config
136
- quant_config = vllm_config.quant_config
133
134
self.config = config
139
- self.quant_config = quant_config
135
self.vocab_size = config.vocab_size
136
self.unpadded_vocab_size = config.vocab_size
137
@@ -162,7 +157,7 @@ class TransformersModel(nn.Module):
157
scale=config.head_dim**-0.5,
158
num_kv_heads=divide(config.num_key_value_heads, tp_size),
159
cache_config=cache_config,
165
- quant_config=None,
160
+ quant_config=self.quant_config,
161
prefix=f"{i}.attn") for i in range(config.num_hidden_layers)
162
]
163
@@ -172,7 +167,7 @@ class TransformersModel(nn.Module):
167
# ForCausalLM modifications
168
self.lm_head = ParallelLMHead(config.vocab_size,
169
config.hidden_size,
175
- quant_config=None,
170
+ quant_config=self.quant_config,
171
prefix=maybe_prefix(prefix, "lm_head"))
172
if config.tie_word_embeddings:
173
self.lm_head.weight = self.model.get_input_embeddings().weight