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 14, 05:30 UTC
[Model] Remove quantized mixtral (#24437)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Sep 9, 2025 at 02:02 UTC
8d7f39b48c9df750678171e37586dd84def5737e
4 files changed
-472
tests/models/registry.py
-1
@@ -285,7 +285,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
285
"MistralForCausalLM": _HfExamplesInfo("mistralai/Mistral-7B-Instruct-v0.1"),
286
"MixtralForCausalLM": _HfExamplesInfo("mistralai/Mixtral-8x7B-Instruct-v0.1", # noqa: E501
287
{"tiny": "TitanML/tiny-mixtral"}), # noqa: E501
288
- "QuantMixtralForCausalLM": _HfExamplesInfo("mistral-community/Mixtral-8x22B-v0.1-AWQ"), # noqa: E501
288
"MptForCausalLM": _HfExamplesInfo("mpt", is_available_online=False),
289
"MPTForCausalLM": _HfExamplesInfo("mosaicml/mpt-7b"),
290
"NemotronForCausalLM": _HfExamplesInfo("nvidia/Minitron-8B-Base"),
vllm/model_executor/model_loader/utils.py
-16
@@ -169,22 +169,6 @@ def get_model_architecture(
169
model_config: ModelConfig) -> tuple[type[nn.Module], str]:
170
architectures = getattr(model_config.hf_config, "architectures", [])
171
172
- # Special handling for quantized Mixtral.
173
- # FIXME(woosuk): This is a temporary hack.
174
- mixtral_supported = [
175
- "fp8",
176
- "compressed-tensors",
177
- "gptq_marlin",
178
- "awq_marlin",
179
- "quark",
180
- "bitsandbytes",
181
- ]
182
-
183
- if (model_config.quantization is not None
184
- and model_config.quantization not in mixtral_supported
185
- and "MixtralForCausalLM" in architectures):
186
- architectures = ["QuantMixtralForCausalLM"]
187
-
172
model_cls, arch = model_config.registry.resolve_model_cls(
173
architectures,
174
model_config=model_config,
vllm/model_executor/models/mixtral_quant.py
deleted
-454
@@ -1,454 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-# Adapted from
5
-# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
6
-# Copyright 2023 The vLLM team.
7
-# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
8
-#
9
-# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
10
-# and OPT implementations in this library. It has been modified from its
11
-# original forms to accommodate minor architectural differences compared
12
-# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
13
-#
14
-# Licensed under the Apache License, Version 2.0 (the "License");
15
-# you may not use this file except in compliance with the License.
16
-# You may obtain a copy of the License at
17
-#
18
-# http://www.apache.org/licenses/LICENSE-2.0
19
-#
20
-# Unless required by applicable law or agreed to in writing, software
21
-# distributed under the License is distributed on an "AS IS" BASIS,
22
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
23
-# See the License for the specific language governing permissions and
24
-# limitations under the License.
25
-"""Inference-only Mixtral model."""
26
-from collections.abc import Iterable
27
-from itertools import islice
28
-from typing import Optional, Union
29
-
30
-import numpy as np
31
-import torch
32
-import torch.nn.functional as F
33
-from torch import nn
34
-from transformers import MixtralConfig
35
-
36
-from vllm.attention import Attention
37
-from vllm.config import CacheConfig, VllmConfig
38
-from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
39
- get_tensor_model_parallel_world_size,
40
- tensor_model_parallel_all_reduce)
41
-from vllm.model_executor.layers.layernorm import RMSNorm
42
-from vllm.model_executor.layers.linear import (QKVParallelLinear,
43
- ReplicatedLinear,
44
- RowParallelLinear)
45
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
46
-from vllm.model_executor.layers.quantization import QuantizationConfig
47
-from vllm.model_executor.layers.rotary_embedding import get_rope
48
-from vllm.model_executor.layers.vocab_parallel_embedding import (
49
- ParallelLMHead, VocabParallelEmbedding)
50
-from vllm.model_executor.model_loader.weight_utils import (
51
- default_weight_loader, maybe_remap_kv_scale_name)
52
-from vllm.model_executor.sampling_metadata import SamplingMetadata
53
-from vllm.sequence import IntermediateTensors
54
-
55
-from .interfaces import SupportsPP
56
-from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
57
- make_empty_intermediate_tensors_factory, make_layers,
58
- maybe_prefix)
59
-
60
-
61
-class MixtralMLP(nn.Module):
62
-
63
- def __init__(
64
- self,
65
- num_experts: int,
66
- hidden_size: int,
67
- intermediate_size: int,
68
- quant_config: Optional[QuantizationConfig] = None,
69
- ) -> None:
70
- super().__init__()
71
- self.num_experts = num_experts
72
- self.ffn_dim = intermediate_size
73
- self.hidden_dim = hidden_size
74
-
75
- self.w1 = ReplicatedLinear(self.hidden_dim,
76
- self.ffn_dim,
77
- bias=False,
78
- quant_config=quant_config)
79
- self.w2 = ReplicatedLinear(self.ffn_dim,
80
- self.hidden_dim,
81
- bias=False,
82
- quant_config=quant_config)
83
- self.w3 = ReplicatedLinear(self.hidden_dim,
84
- self.ffn_dim,
85
- bias=False,
86
- quant_config=quant_config)
87
-
88
- # TODO: Use vllm's SiluAndMul
89
- self.act_fn = nn.SiLU()
90
-
91
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
92
- w1_out, _ = self.w1(hidden_states)
93
- w1_out = self.act_fn(w1_out)
94
- w3_out, _ = self.w3(hidden_states)
95
- current_hidden_states = w1_out * w3_out
96
- current_hidden_states, _ = self.w2(current_hidden_states)
97
- return current_hidden_states
98
-
99
-
100
-class MixtralMoE(nn.Module):
101
-
102
- def __init__(
103
- self,
104
- config: MixtralConfig,
105
- quant_config: Optional[QuantizationConfig] = None,
106
- ):
107
- super().__init__()
108
- self.config = config
109
- self.rank = get_tensor_model_parallel_rank()
110
- self.tp_size = get_tensor_model_parallel_world_size()
111
- self.num_total_experts = config.num_local_experts
112
- self.top_k = config.num_experts_per_tok
113
- if self.tp_size > self.num_total_experts:
114
- raise ValueError(
115
- f"Tensor parallel size {self.tp_size} is greater than "
116
- f"the number of experts {self.num_total_experts}.")
117
- # Split experts equally between ranks
118
- self.expert_indices = np.array_split(range(self.num_total_experts),
119
- self.tp_size)[self.rank].tolist()
120
- if not self.expert_indices:
121
- raise ValueError(
122
- f"Rank {self.rank} has no experts assigned to it.")
123
-
124
- self.experts = nn.ModuleList([
125
- MixtralMLP(self.num_total_experts,
126
- config.hidden_size,
127
- config.intermediate_size,
128
- quant_config=quant_config)
129
- if idx in self.expert_indices else None
130
- for idx in range(self.num_total_experts)
131
- ])
132
- self.gate = ReplicatedLinear(config.hidden_size,
133
- self.num_total_experts,
134
- bias=False,
135
- quant_config=None)
136
-
137
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
138
- num_tokens, hidden_dim = hidden_states.shape
139
- hidden_states = hidden_states.view(-1, hidden_dim)
140
- # router_logits: (num_tokens, n_experts)
141
- router_logits, _ = self.gate(hidden_states)
142
-
143
- routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
144
- routing_weights, selected_experts = torch.topk(routing_weights,
145
- self.top_k,
146
- dim=-1)
147
- routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
148
-
149
- final_hidden_states = None
150
- for expert_idx in self.expert_indices:
151
- expert_layer = self.experts[expert_idx]
152
- expert_mask = (selected_experts == expert_idx)
153
- expert_weights = (routing_weights * expert_mask).sum(dim=-1,
154
- keepdim=True)
155
-
156
- current_hidden_states = expert_layer(hidden_states).mul_(
157
- expert_weights)
158
- if final_hidden_states is None:
159
- final_hidden_states = current_hidden_states
160
- else:
161
- final_hidden_states.add_(current_hidden_states)
162
-
163
- return tensor_model_parallel_all_reduce(final_hidden_states).view(
164
- num_tokens, hidden_dim)
165
-
166
-
167
-class MixtralAttention(nn.Module):
168
-
169
- def __init__(
170
- self,
171
- config: MixtralConfig,
172
- hidden_size: int,
173
- num_heads: int,
174
- num_kv_heads: int,
175
- max_position: int = 4096 * 32,
176
- rope_theta: float = 10000,
177
- quant_config: Optional[QuantizationConfig] = None,
178
- cache_config: Optional[CacheConfig] = None,
179
- prefix: str = "",
180
- ) -> None:
181
- super().__init__()
182
- self.hidden_size = hidden_size
183
- tp_size = get_tensor_model_parallel_world_size()
184
- self.total_num_heads = num_heads
185
- assert self.total_num_heads % tp_size == 0
186
- self.num_heads = self.total_num_heads // tp_size
187
- self.total_num_kv_heads = num_kv_heads
188
- if self.total_num_kv_heads >= tp_size:
189
- # Number of KV heads is greater than TP size, so we partition
190
- # the KV heads across multiple tensor parallel GPUs.
191
- assert self.total_num_kv_heads % tp_size == 0
192
- else:
193
- # Number of KV heads is less than TP size, so we replicate
194
- # the KV heads across multiple tensor parallel GPUs.
195
- assert tp_size % self.total_num_kv_heads == 0
196
- self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
197
- # MixtralConfig has an optional head_dim argument
198
- self.head_dim = getattr(config, "head_dim", None)
199
- if self.head_dim is None:
200
- self.head_dim = self.hidden_size // self.total_num_heads
201
- self.q_size = self.num_heads * self.head_dim
202
- self.kv_size = self.num_kv_heads * self.head_dim
203
- self.scaling = self.head_dim**-0.5
204
- self.rope_theta = rope_theta
205
-
206
- self.qkv_proj = QKVParallelLinear(
207
- hidden_size,
208
- self.head_dim,
209
- self.total_num_heads,
210
- self.total_num_kv_heads,
211
- bias=False,
212
- quant_config=quant_config,
213
- )
214
- self.o_proj = RowParallelLinear(
215
- self.total_num_heads * self.head_dim,
216
- hidden_size,
217
- bias=False,
218
- quant_config=quant_config,
219
- )
220
- self.rotary_emb = get_rope(
221
- self.head_dim,
222
- rotary_dim=self.head_dim,
223
- max_position=max_position,
224
- base=int(self.rope_theta),
225
- is_neox_style=True,
226
- )
227
- self.attn = Attention(self.num_heads,
228
- self.head_dim,
229
- self.scaling,
230
- num_kv_heads=self.num_kv_heads,
231
- cache_config=cache_config,
232
- quant_config=quant_config,
233
- prefix=f"{prefix}.attn")
234
-
235
- def forward(
236
- self,
237
- positions: torch.Tensor,
238
- hidden_states: torch.Tensor,
239
- ) -> torch.Tensor:
240
- qkv, _ = self.qkv_proj(hidden_states)
241
- q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
242
- q, k = self.rotary_emb(positions, q, k)
243
- attn_output = self.attn(q, k, v)
244
- output, _ = self.o_proj(attn_output)
245
- return output
246
-
247
-
248
-class MixtralDecoderLayer(nn.Module):
249
-
250
- def __init__(
251
- self,
252
- config: MixtralConfig,
253
- cache_config: Optional[CacheConfig] = None,
254
- quant_config: Optional[QuantizationConfig] = None,
255
- prefix: str = "",
256
- ) -> None:
257
- super().__init__()
258
- self.hidden_size = config.hidden_size
259
- # Requires transformers > 4.32.0
260
- rope_theta = getattr(config, "rope_theta", 10000)
261
- self.self_attn = MixtralAttention(
262
- config=config,
263
- hidden_size=self.hidden_size,
264
- num_heads=config.num_attention_heads,
265
- max_position=config.max_position_embeddings,
266
- num_kv_heads=config.num_key_value_heads,
267
- rope_theta=rope_theta,
268
- cache_config=cache_config,
269
- quant_config=quant_config,
270
- prefix=f"{prefix}.self_attn",
271
- )
272
- self.block_sparse_moe = MixtralMoE(config=config,
273
- quant_config=quant_config)
274
- self.input_layernorm = RMSNorm(config.hidden_size,
275
- eps=config.rms_norm_eps)
276
- self.post_attention_layernorm = RMSNorm(config.hidden_size,
277
- eps=config.rms_norm_eps)
278
-
279
- def forward(
280
- self,
281
- positions: torch.Tensor,
282
- hidden_states: torch.Tensor,
283
- residual: Optional[torch.Tensor],
284
- ) -> torch.Tensor:
285
- # Self Attention
286
- if residual is None:
287
- residual = hidden_states
288
- hidden_states = self.input_layernorm(hidden_states)
289
- else:
290
- hidden_states, residual = self.input_layernorm(
291
- hidden_states, residual)
292
- hidden_states = self.self_attn(
293
- positions=positions,
294
- hidden_states=hidden_states,
295
- )
296
-
297
- # Fully Connected
298
- hidden_states, residual = self.post_attention_layernorm(
299
- hidden_states, residual)
300
- hidden_states = self.block_sparse_moe(hidden_states)
301
- return hidden_states, residual
302
-
303
-
304
-class MixtralModel(nn.Module):
305
-
306
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
307
- super().__init__()
308
-
309
- config = vllm_config.model_config.hf_config
310
- cache_config = vllm_config.cache_config
311
- quant_config = vllm_config.quant_config
312
-
313
- self.vocab_size = config.vocab_size
314
-
315
- self.embed_tokens = VocabParallelEmbedding(
316
- config.vocab_size,
317
- config.hidden_size,
318
- )
319
- self.start_layer, self.end_layer, self.layers = make_layers(
320
- config.num_hidden_layers,
321
- lambda prefix: MixtralDecoderLayer(
322
- config, cache_config, quant_config=quant_config, prefix=prefix
323
- ),
324
- prefix=f"{prefix}.layers")
325
- self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
326
- self.make_empty_intermediate_tensors = (
327
- make_empty_intermediate_tensors_factory(
328
- ["hidden_states", "residual"], config.hidden_size))
329
-
330
- def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
331
- return self.embed_tokens(input_ids)
332
-
333
- def forward(
334
- self,
335
- input_ids: torch.Tensor,
336
- positions: torch.Tensor,
337
- intermediate_tensors: Optional[IntermediateTensors],
338
- inputs_embeds: Optional[torch.Tensor] = None,
339
- ) -> Union[torch.Tensor, IntermediateTensors]:
340
- if get_pp_group().is_first_rank:
341
- if inputs_embeds is not None:
342
- hidden_states = inputs_embeds
343
- else:
344
- hidden_states = self.get_input_embeddings(input_ids)
345
- residual = None
346
- else:
347
- assert intermediate_tensors is not None
348
- hidden_states = intermediate_tensors["hidden_states"]
349
- residual = intermediate_tensors["residual"]
350
- for layer in islice(self.layers, self.start_layer, self.end_layer):
351
- hidden_states, residual = layer(positions, hidden_states, residual)
352
- if not get_pp_group().is_last_rank:
353
- return IntermediateTensors({
354
- "hidden_states": hidden_states,
355
- "residual": residual
356
- })
357
- hidden_states, _ = self.norm(hidden_states, residual)
358
- return hidden_states
359
-
360
- def load_weights(self, weights: Iterable[tuple[str,
361
- torch.Tensor]]) -> set[str]:
362
- stacked_params_mapping = [
363
- # (param_name, shard_name, shard_id)
364
- ("qkv_proj", "q_proj", "q"),
365
- ("qkv_proj", "k_proj", "k"),
366
- ("qkv_proj", "v_proj", "v"),
367
- ]
368
-
369
- params_dict = dict(self.named_parameters())
370
- loaded_params: set[str] = set()
371
- for name, loaded_weight in weights:
372
- if name.endswith("scale"):
373
- # Remapping the name of FP8 kv-scale.
374
- name = maybe_remap_kv_scale_name(name, params_dict)
375
- if name is None:
376
- continue
377
- for (param_name, weight_name, shard_id) in stacked_params_mapping:
378
- if weight_name not in name:
379
- continue
380
- name = name.replace(weight_name, param_name)
381
- # Skip loading extra bias for GPTQ models.
382
- if name.endswith(".bias") and name not in params_dict:
383
- continue
384
- if is_pp_missing_parameter(name, self):
385
- continue
386
- param = params_dict[name]
387
- weight_loader = param.weight_loader
388
- weight_loader(param, loaded_weight, shard_id)
389
- break
390
- else:
391
- # Skip loading extra bias for GPTQ models.
392
- if name.endswith(".bias") and name not in params_dict:
393
- continue
394
- # Skip experts that are not assigned to this worker.
395
- if ("block_sparse_moe.experts." in name
396
- and name not in params_dict):
397
- continue
398
- if is_pp_missing_parameter(name, self):
399
- continue
400
- param = params_dict[name]
401
- weight_loader = getattr(param, "weight_loader",
402
- default_weight_loader)
403
- weight_loader(param, loaded_weight)
404
- loaded_params.add(name)
405
- return loaded_params
406
-
407
-
408
-class MixtralForCausalLM(nn.Module, SupportsPP):
409
- fall_back_to_pt_during_load = False
410
-
411
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
412
- super().__init__()
413
- config = vllm_config.model_config.hf_config
414
- quant_config = vllm_config.quant_config
415
- self.config = config
416
- self.quant_config = quant_config
417
- self.model = MixtralModel(vllm_config=vllm_config,
418
- prefix=maybe_prefix(prefix, "model"))
419
- self.lm_head = ParallelLMHead(config.vocab_size,
420
- config.hidden_size,
421
- quant_config=quant_config)
422
- if self.config.tie_word_embeddings:
423
- self.lm_head.weight = self.model.embed_tokens.weight
424
- self.logits_processor = LogitsProcessor(config.vocab_size)
425
- self.make_empty_intermediate_tensors = (
426
- self.model.make_empty_intermediate_tensors)
427
-
428
- def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
429
- return self.model.get_input_embeddings(input_ids)
430
-
431
- def forward(
432
- self,
433
- input_ids: torch.Tensor,
434
- positions: torch.Tensor,
435
- intermediate_tensors: Optional[IntermediateTensors] = None,
436
- inputs_embeds: Optional[torch.Tensor] = None,
437
- ) -> Union[torch.Tensor, IntermediateTensors]:
438
- hidden_states = self.model(input_ids, positions, intermediate_tensors,
439
- inputs_embeds)
440
- return hidden_states
441
-
442
- def compute_logits(
443
- self,
444
- hidden_states: torch.Tensor,
445
- sampling_metadata: SamplingMetadata,
446
- ) -> Optional[torch.Tensor]:
447
- logits = self.logits_processor(self.lm_head, hidden_states,
448
- sampling_metadata)
449
- return logits
450
-
451
- def load_weights(self, weights: Iterable[tuple[str,
452
- torch.Tensor]]) -> set[str]:
453
- loader = AutoWeightsLoader(self)
454
- return loader.load_weights(weights)
vllm/model_executor/models/registry.py
-1
@@ -110,7 +110,6 @@ _TEXT_GENERATION_MODELS = {
110
"MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"),
111
"MistralForCausalLM": ("llama", "LlamaForCausalLM"),
112
"MixtralForCausalLM": ("mixtral", "MixtralForCausalLM"),
113
- "QuantMixtralForCausalLM": ("mixtral_quant", "MixtralForCausalLM"),
113
# transformers's mpt class has lower case
114
"MptForCausalLM": ("mpt", "MPTForCausalLM"),
115
"MPTForCausalLM": ("mpt", "MPTForCausalLM"),