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, 14:45 UTC
[Model] use AutoWeightsLoader for granite, granitemoe, granitemoeshared, grok1, mixtral (#16325)
Signed-off-by: Aaron Ang <aaron.angyd@gmail.com>
Aaron Ang committed
Apr 9, 2025 at 23:07 UTC
a564797151a0cdf2698802997f9dba675fb87dd5
5 files changed
+353
-323
vllm/model_executor/models/granite.py
+76
-69
@@ -50,8 +50,8 @@ from vllm.model_executor.sampling_metadata import SamplingMetadata
50
from vllm.sequence import IntermediateTensors
51
52
from .interfaces import SupportsLoRA, SupportsPP
53
-from .utils import (PPMissingLayer, is_pp_missing_parameter, make_layers,
54
- maybe_prefix)
53
+from .utils import (AutoWeightsLoader, PPMissingLayer, is_pp_missing_parameter,
54
+ make_layers, maybe_prefix)
55
56
57
class GraniteMLP(nn.Module):
@@ -260,6 +260,7 @@ class GraniteModel(nn.Module):
260
lora_config = vllm_config.lora_config
261
262
self.config = config
263
+ self.quant_config = quant_config
264
lora_vocab = (lora_config.lora_extra_vocab_size *
265
(lora_config.max_loras or 1)) if lora_config else 0
266
self.vocab_size = config.vocab_size + lora_vocab
@@ -321,6 +322,65 @@ class GraniteModel(nn.Module):
322
hidden_states = self.norm(hidden_states)
323
return hidden_states
324
325
+ def load_weights(self, weights: Iterable[Tuple[str,
326
+ torch.Tensor]]) -> Set[str]:
327
+ stacked_params_mapping = [
328
+ # (param_name, shard_name, shard_id)
329
+ (".qkv_proj", ".q_proj", "q"),
330
+ (".qkv_proj", ".k_proj", "k"),
331
+ (".qkv_proj", ".v_proj", "v"),
332
+ (".gate_up_proj", ".gate_proj", 0),
333
+ (".gate_up_proj", ".up_proj", 1),
334
+ ]
335
+ params_dict = dict(self.named_parameters())
336
+ loaded_params: Set[str] = set()
337
+ for name, loaded_weight in weights:
338
+ if (self.quant_config is not None and
339
+ (scale_name := self.quant_config.get_cache_scale(name))):
340
+ # Loading kv cache quantization scales
341
+ param = params_dict[scale_name]
342
+ weight_loader = getattr(param, "weight_loader",
343
+ default_weight_loader)
344
+ loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
345
+ loaded_weight[0])
346
+ weight_loader(param, loaded_weight)
347
+ loaded_params.add(scale_name)
348
+ continue
349
+ for (param_name, weight_name, shard_id) in stacked_params_mapping:
350
+ if weight_name not in name:
351
+ continue
352
+ name = name.replace(weight_name, param_name)
353
+ # Skip loading extra bias for GPTQ models.
354
+ if name.endswith(".bias") and name not in params_dict:
355
+ continue
356
+
357
+ if is_pp_missing_parameter(name, self):
358
+ continue
359
+
360
+ param = params_dict[name]
361
+ weight_loader = param.weight_loader
362
+ weight_loader(param, loaded_weight, shard_id)
363
+
364
+ break
365
+ else:
366
+ # Skip loading extra bias for GPTQ models.
367
+ if name.endswith(".bias") and name not in params_dict:
368
+ continue
369
+ # Remapping the name of FP8 kv-scale.
370
+ name = maybe_remap_kv_scale_name(name, params_dict)
371
+ if name is None:
372
+ continue
373
+
374
+ if is_pp_missing_parameter(name, self):
375
+ continue
376
+
377
+ param = params_dict[name]
378
+ weight_loader = getattr(param, "weight_loader",
379
+ default_weight_loader)
380
+ weight_loader(param, loaded_weight)
381
+ loaded_params.add(name)
382
+ return loaded_params
383
+
384
385
class GraniteForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
386
packed_modules_mapping = {
@@ -428,71 +488,18 @@ class GraniteForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
488
489
def load_weights(self, weights: Iterable[Tuple[str,
490
torch.Tensor]]) -> Set[str]:
431
- stacked_params_mapping = [
432
- # (param_name, shard_name, shard_id)
433
- (".qkv_proj", ".q_proj", "q"),
434
- (".qkv_proj", ".k_proj", "k"),
435
- (".qkv_proj", ".v_proj", "v"),
436
- (".gate_up_proj", ".gate_proj", 0),
437
- (".gate_up_proj", ".up_proj", 1),
491
+ skip_prefixes = [
492
+ "rotary_emb.inv_freq",
493
+ # Models trained using ColossalAI may include these tensors in
494
+ # the checkpoint. Skip them.
495
+ "rotary_emb.cos_cached",
496
+ "rotary_emb.sin_cached",
497
]
439
- params_dict = dict(self.named_parameters())
440
- loaded_params: Set[str] = set()
441
- for name, loaded_weight in weights:
442
- if "rotary_emb.inv_freq" in name:
443
- continue
444
- if ("rotary_emb.cos_cached" in name
445
- or "rotary_emb.sin_cached" in name):
446
- # Models trained using ColossalAI may include these tensors in
447
- # the checkpoint. Skip them.
448
- continue
449
- # With tie_word_embeddings, we can skip lm_head.weight
450
- # The weight might appear unnecessarily in the files if the model is
451
- # processed with quantization, LoRA, fine-tuning, etc.
452
- if self.config.tie_word_embeddings and "lm_head.weight" in name:
453
- continue
454
- if (self.quant_config is not None and
455
- (scale_name := self.quant_config.get_cache_scale(name))):
456
- # Loading kv cache quantization scales
457
- param = params_dict[scale_name]
458
- weight_loader = getattr(param, "weight_loader",
459
- default_weight_loader)
460
- loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
461
- loaded_weight[0])
462
- weight_loader(param, loaded_weight)
463
- loaded_params.add(scale_name)
464
- continue
465
- for (param_name, weight_name, shard_id) in stacked_params_mapping:
466
- if weight_name not in name:
467
- continue
468
- name = name.replace(weight_name, param_name)
469
- # Skip loading extra bias for GPTQ models.
470
- if name.endswith(".bias") and name not in params_dict:
471
- continue
472
-
473
- if is_pp_missing_parameter(name, self):
474
- continue
475
-
476
- param = params_dict[name]
477
- weight_loader = param.weight_loader
478
- weight_loader(param, loaded_weight, shard_id)
479
-
480
- break
481
- else:
482
- # Skip loading extra bias for GPTQ models.
483
- if name.endswith(".bias") and name not in params_dict:
484
- continue
485
- # Remapping the name of FP8 kv-scale.
486
- name = maybe_remap_kv_scale_name(name, params_dict)
487
- if name is None:
488
- continue
489
-
490
- if is_pp_missing_parameter(name, self):
491
- continue
492
-
493
- param = params_dict[name]
494
- weight_loader = getattr(param, "weight_loader",
495
- default_weight_loader)
496
- weight_loader(param, loaded_weight)
497
- loaded_params.add(name)
498
- return loaded_params
498
+ # With tie_word_embeddings, we can skip lm_head.weight
499
+ # The weight might appear unnecessarily in the files if the model is
500
+ # processed with quantization, LoRA, fine-tuning, etc.
501
+ if self.config.tie_word_embeddings:
502
+ skip_prefixes.append("lm_head.weight")
503
+
504
+ loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
505
+ return loader.load_weights(weights)
vllm/model_executor/models/granitemoe.py
+43
-36
@@ -49,7 +49,7 @@ from vllm.sequence import IntermediateTensors
49
50
from . import mixtral
51
from .interfaces import SupportsLoRA, SupportsPP
52
-from .utils import make_layers, maybe_prefix
52
+from .utils import AutoWeightsLoader, make_layers, maybe_prefix
53
54
55
class GraniteMoeMoE(nn.Module):
@@ -252,6 +252,8 @@ class GraniteMoeModel(nn.Module):
252
quant_config = vllm_config.quant_config
253
lora_config = vllm_config.lora_config
254
255
+ self.config = config
256
+ self.quant_config = quant_config # Required by MixtralModel
257
lora_vocab = (lora_config.lora_extra_vocab_size *
258
(lora_config.max_loras or 1)) if lora_config else 0
259
self.vocab_size = config.vocab_size + lora_vocab
@@ -304,6 +306,40 @@ class GraniteMoeModel(nn.Module):
306
hidden_states = self.norm(hidden_states)
307
return hidden_states
308
309
+ def load_weights(self, weights: Iterable[Tuple[str,
310
+ torch.Tensor]]) -> Set[str]:
311
+ new_weights = {}
312
+ for n, p in weights:
313
+ if n.endswith('.block_sparse_moe.input_linear.weight'):
314
+ for e in range(p.size(0)):
315
+ w1_name = n.replace(
316
+ '.block_sparse_moe.input_linear.weight',
317
+ f".block_sparse_moe.experts.{e}.w1.weight")
318
+ w3_name = n.replace(
319
+ '.block_sparse_moe.input_linear.weight',
320
+ f".block_sparse_moe.experts.{e}.w3.weight")
321
+ w1_param, w3_param = p[e].chunk(2, dim=0)
322
+ assert w1_name not in new_weights
323
+ assert w3_name not in new_weights
324
+ new_weights[w1_name] = w1_param
325
+ new_weights[w3_name] = w3_param
326
+ elif n.endswith('.block_sparse_moe.output_linear.weight'):
327
+ for e in range(p.size(0)):
328
+ w2_name = n.replace(
329
+ '.block_sparse_moe.output_linear.weight',
330
+ f".block_sparse_moe.experts.{e}.w2.weight")
331
+ w2_param = p[e]
332
+ assert w2_name not in new_weights
333
+ new_weights[w2_name] = w2_param
334
+ elif n.endswith('.block_sparse_moe.router.layer.weight'):
335
+ gate_name = n.replace('.block_sparse_moe.router.layer.weight',
336
+ ".block_sparse_moe.gate.weight")
337
+ assert gate_name not in new_weights
338
+ new_weights[gate_name] = p
339
+ else:
340
+ new_weights[n] = p
341
+ return mixtral.MixtralModel.load_weights(self, new_weights.items())
342
+
343
344
class GraniteMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
345
fall_back_to_pt_during_load = False
@@ -331,7 +367,6 @@ class GraniteMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
367
368
self.config = config
369
self.lora_config = lora_config
334
- self.quant_config = quant_config # Required by MixtralForCausalLM
370
371
self.model = GraniteMoeModel(vllm_config=vllm_config,
372
prefix=maybe_prefix(prefix, "model"))
@@ -403,37 +438,9 @@ class GraniteMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
438
439
def load_weights(self, weights: Iterable[Tuple[str,
440
torch.Tensor]]) -> Set[str]:
406
- new_weights = {}
407
- for n, p in weights:
408
- if n.endswith('.block_sparse_moe.input_linear.weight'):
409
- for e in range(p.size(0)):
410
- w1_name = n.replace(
411
- '.block_sparse_moe.input_linear.weight',
412
- f".block_sparse_moe.experts.{e}.w1.weight")
413
- w3_name = n.replace(
414
- '.block_sparse_moe.input_linear.weight',
415
- f".block_sparse_moe.experts.{e}.w3.weight")
416
- w1_param, w3_param = p[e].chunk(2, dim=0)
417
- assert w1_name not in new_weights
418
- assert w3_name not in new_weights
419
- new_weights[w1_name] = w1_param
420
- new_weights[w3_name] = w3_param
421
- elif n.endswith('.block_sparse_moe.output_linear.weight'):
422
- for e in range(p.size(0)):
423
- w2_name = n.replace(
424
- '.block_sparse_moe.output_linear.weight',
425
- f".block_sparse_moe.experts.{e}.w2.weight")
426
- w2_param = p[e]
427
- assert w2_name not in new_weights
428
- new_weights[w2_name] = w2_param
429
- elif n.endswith('.block_sparse_moe.router.layer.weight'):
430
- gate_name = n.replace('.block_sparse_moe.router.layer.weight',
431
- ".block_sparse_moe.gate.weight")
432
- assert gate_name not in new_weights
433
- new_weights[gate_name] = p
434
- elif n == 'lm_head.weight' and self.config.tie_word_embeddings:
435
- pass
436
- else:
437
- new_weights[n] = p
438
- return mixtral.MixtralForCausalLM.load_weights(self,
439
- new_weights.items())
441
+ loader = AutoWeightsLoader(
442
+ self,
443
+ skip_prefixes=(["lm_head."]
444
+ if self.config.tie_word_embeddings else None),
445
+ )
446
+ return loader.load_weights(weights)
vllm/model_executor/models/granitemoeshared.py
+43
-36
@@ -29,7 +29,7 @@ from vllm.sequence import IntermediateTensors
29
from . import mixtral
30
from .granitemoe import GraniteMoeAttention, GraniteMoeMoE
31
from .interfaces import SupportsLoRA, SupportsPP
32
-from .utils import make_layers, maybe_prefix
32
+from .utils import AutoWeightsLoader, make_layers, maybe_prefix
33
34
35
class GraniteMoeSharedMLP(nn.Module):
@@ -152,6 +152,8 @@ class GraniteMoeSharedModel(nn.Module):
152
quant_config = vllm_config.quant_config
153
lora_config = vllm_config.lora_config
154
155
+ self.config = config
156
+ self.quant_config = quant_config # Required by MixtralModel
157
self.padding_idx = config.pad_token_id
158
lora_vocab = (lora_config.lora_extra_vocab_size *
159
(lora_config.max_loras or 1)) if lora_config else 0
@@ -207,6 +209,40 @@ class GraniteMoeSharedModel(nn.Module):
209
hidden_states = self.norm(hidden_states)
210
return hidden_states
211
212
+ def load_weights(self, weights: Iterable[Tuple[str,
213
+ torch.Tensor]]) -> Set[str]:
214
+ new_weights = {}
215
+ for n, p in weights:
216
+ if n.endswith('.block_sparse_moe.input_linear.weight'):
217
+ for e in range(p.size(0)):
218
+ w1_name = n.replace(
219
+ '.block_sparse_moe.input_linear.weight',
220
+ f".block_sparse_moe.experts.{e}.w1.weight")
221
+ w3_name = n.replace(
222
+ '.block_sparse_moe.input_linear.weight',
223
+ f".block_sparse_moe.experts.{e}.w3.weight")
224
+ w1_param, w3_param = p[e].chunk(2, dim=0)
225
+ assert w1_name not in new_weights
226
+ assert w3_name not in new_weights
227
+ new_weights[w1_name] = w1_param
228
+ new_weights[w3_name] = w3_param
229
+ elif n.endswith('.block_sparse_moe.output_linear.weight'):
230
+ for e in range(p.size(0)):
231
+ w2_name = n.replace(
232
+ '.block_sparse_moe.output_linear.weight',
233
+ f".block_sparse_moe.experts.{e}.w2.weight")
234
+ w2_param = p[e]
235
+ assert w2_name not in new_weights
236
+ new_weights[w2_name] = w2_param
237
+ elif n.endswith('.block_sparse_moe.router.layer.weight'):
238
+ gate_name = n.replace('.block_sparse_moe.router.layer.weight',
239
+ ".block_sparse_moe.gate.weight")
240
+ assert gate_name not in new_weights
241
+ new_weights[gate_name] = p
242
+ else:
243
+ new_weights[n] = p
244
+ return mixtral.MixtralModel.load_weights(self, new_weights.items())
245
+
246
247
class GraniteMoeSharedForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
248
fall_back_to_pt_during_load = False
@@ -234,7 +270,6 @@ class GraniteMoeSharedForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
270
271
self.config = config
272
self.lora_config = lora_config
237
- self.quant_config = quant_config
273
274
self.model = GraniteMoeSharedModel(vllm_config=vllm_config,
275
prefix=maybe_prefix(
@@ -307,37 +342,9 @@ class GraniteMoeSharedForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
342
343
def load_weights(self, weights: Iterable[Tuple[str,
344
torch.Tensor]]) -> Set[str]:
310
- new_weights = {}
311
- for n, p in weights:
312
- if n.endswith('.block_sparse_moe.input_linear.weight'):
313
- for e in range(p.size(0)):
314
- w1_name = n.replace(
315
- '.block_sparse_moe.input_linear.weight',
316
- f".block_sparse_moe.experts.{e}.w1.weight")
317
- w3_name = n.replace(
318
- '.block_sparse_moe.input_linear.weight',
319
- f".block_sparse_moe.experts.{e}.w3.weight")
320
- w1_param, w3_param = p[e].chunk(2, dim=0)
321
- assert w1_name not in new_weights
322
- assert w3_name not in new_weights
323
- new_weights[w1_name] = w1_param
324
- new_weights[w3_name] = w3_param
325
- elif n.endswith('.block_sparse_moe.output_linear.weight'):
326
- for e in range(p.size(0)):
327
- w2_name = n.replace(
328
- '.block_sparse_moe.output_linear.weight',
329
- f".block_sparse_moe.experts.{e}.w2.weight")
330
- w2_param = p[e]
331
- assert w2_name not in new_weights
332
- new_weights[w2_name] = w2_param
333
- elif n.endswith('.block_sparse_moe.router.layer.weight'):
334
- gate_name = n.replace('.block_sparse_moe.router.layer.weight',
335
- ".block_sparse_moe.gate.weight")
336
- assert gate_name not in new_weights
337
- new_weights[gate_name] = p
338
- elif n == 'lm_head.weight' and self.config.tie_word_embeddings:
339
- pass
340
- else:
341
- new_weights[n] = p
342
- return mixtral.MixtralForCausalLM.load_weights(self,
343
- new_weights.items())
345
+ loader = AutoWeightsLoader(
346
+ self,
347
+ skip_prefixes=(["lm_head."]
348
+ if self.config.tie_word_embeddings else None),
349
+ )
350
+ return loader.load_weights(weights)
vllm/model_executor/models/grok1.py
+101
-96
@@ -48,7 +48,7 @@ from vllm.model_executor.sampling_metadata import SamplingMetadata
48
from vllm.sequence import IntermediateTensors
49
50
from .interfaces import SupportsLoRA, SupportsPP
51
-from .utils import (is_pp_missing_parameter,
51
+from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
52
make_empty_intermediate_tensors_factory, make_layers,
53
maybe_prefix)
54
@@ -302,6 +302,8 @@ class Grok1Model(nn.Module):
302
quant_config = vllm_config.quant_config
303
lora_config = vllm_config.lora_config
304
305
+ self.config = config
306
+ self.quant_config = quant_config
307
self.padding_idx = config.pad_token_id
308
lora_vocab = (lora_config.lora_extra_vocab_size *
309
(lora_config.max_loras or 1)) if lora_config else 0
@@ -370,94 +372,6 @@ class Grok1Model(nn.Module):
372
hidden_states, _ = self.norm(hidden_states, residual)
373
return hidden_states
374
373
-
374
-class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
375
- fall_back_to_pt_during_load = False
376
-
377
- packed_modules_mapping = {
378
- "qkv_proj": [
379
- "q_proj",
380
- "k_proj",
381
- "v_proj",
382
- ],
383
- }
384
-
385
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
386
- super().__init__()
387
-
388
- config = vllm_config.model_config.hf_config
389
- quant_config = vllm_config.quant_config
390
- lora_config = vllm_config.lora_config
391
-
392
- self.config = config
393
- self.lora_config = lora_config
394
- self.quant_config = quant_config
395
-
396
- self.model = Grok1Model(vllm_config=vllm_config,
397
- prefix=maybe_prefix(prefix, "model"))
398
-
399
- self.unpadded_vocab_size = config.vocab_size
400
- if lora_config:
401
- self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
402
-
403
- self.lm_head = ParallelLMHead(
404
- self.unpadded_vocab_size,
405
- config.hidden_size,
406
- org_num_embeddings=config.vocab_size,
407
- padding_size=DEFAULT_VOCAB_PADDING_SIZE
408
- # We need bigger padding if using lora for kernel compatibility
409
- if not lora_config else lora_config.lora_vocab_padding_size,
410
- quant_config=quant_config,
411
- prefix=maybe_prefix(prefix, "lm_head"),
412
- )
413
-
414
- if self.config.tie_word_embeddings:
415
- self.lm_head.weight = self.model.embed_tokens.weight
416
-
417
- self.output_multiplier_scale = getattr(
418
- config, "output_multiplier_scale", DEFAULT_OUTPUT_MULTIPLIER_SCALE)
419
- self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
420
- config.vocab_size,
421
- self.output_multiplier_scale)
422
-
423
- self.sampler = get_sampler()
424
- self.make_empty_intermediate_tensors = (
425
- self.model.make_empty_intermediate_tensors)
426
-
427
- def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
428
- return self.model.get_input_embeddings(input_ids)
429
-
430
- def forward(
431
- self,
432
- input_ids: torch.Tensor,
433
- positions: torch.Tensor,
434
- kv_caches: List[torch.Tensor],
435
- attn_metadata: AttentionMetadata,
436
- intermediate_tensors: Optional[IntermediateTensors] = None,
437
- inputs_embeds: Optional[torch.Tensor] = None,
438
- ) -> Union[torch.Tensor, IntermediateTensors]:
439
- hidden_states = self.model(input_ids, positions, kv_caches,
440
- attn_metadata, intermediate_tensors,
441
- inputs_embeds)
442
- return hidden_states
443
-
444
- def compute_logits(
445
- self,
446
- hidden_states: torch.Tensor,
447
- sampling_metadata: SamplingMetadata,
448
- ) -> Optional[torch.Tensor]:
449
- logits = self.logits_processor(self.lm_head, hidden_states,
450
- sampling_metadata)
451
- return logits
452
-
453
- def sample(
454
- self,
455
- logits: Optional[torch.Tensor],
456
- sampling_metadata: SamplingMetadata,
457
- ) -> Optional[SamplerOutput]:
458
- next_tokens = self.sampler(logits, sampling_metadata)
459
- return next_tokens
460
-
375
def load_weights(self, weights: Iterable[Tuple[str,
376
torch.Tensor]]) -> Set[str]:
377
stacked_params_mapping = [
@@ -480,9 +394,6 @@ class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
394
loaded_params: Set[str] = set()
395
396
for name, loaded_weight in weights:
483
- if "rotary_emb.inv_freq" in name:
484
- continue
485
-
397
if (self.quant_config is not None and
398
(scale_name := self.quant_config.get_cache_scale(name))):
399
# Loading kv cache quantization scales
@@ -553,13 +464,107 @@ class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
464
if "norm.scale" in name:
465
name = name.replace("scale", "weight")
466
556
- # Skip lm_head when tie_word_embeddings is True
557
- if "lm_head" in name and self.config.tie_word_embeddings:
558
- continue
559
-
467
param = params_dict[name]
468
weight_loader = getattr(param, "weight_loader",
469
default_weight_loader)
470
weight_loader(param, loaded_weight)
471
loaded_params.add(name)
472
return loaded_params
473
+
474
+
475
+class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
476
+ fall_back_to_pt_during_load = False
477
+
478
+ packed_modules_mapping = {
479
+ "qkv_proj": [
480
+ "q_proj",
481
+ "k_proj",
482
+ "v_proj",
483
+ ],
484
+ }
485
+
486
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
487
+ super().__init__()
488
+
489
+ config = vllm_config.model_config.hf_config
490
+ quant_config = vllm_config.quant_config
491
+ lora_config = vllm_config.lora_config
492
+
493
+ self.config = config
494
+ self.lora_config = lora_config
495
+ self.quant_config = quant_config
496
+
497
+ self.model = Grok1Model(vllm_config=vllm_config,
498
+ prefix=maybe_prefix(prefix, "model"))
499
+
500
+ self.unpadded_vocab_size = config.vocab_size
501
+ if lora_config:
502
+ self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
503
+
504
+ self.lm_head = ParallelLMHead(
505
+ self.unpadded_vocab_size,
506
+ config.hidden_size,
507
+ org_num_embeddings=config.vocab_size,
508
+ padding_size=DEFAULT_VOCAB_PADDING_SIZE
509
+ # We need bigger padding if using lora for kernel compatibility
510
+ if not lora_config else lora_config.lora_vocab_padding_size,
511
+ quant_config=quant_config,
512
+ prefix=maybe_prefix(prefix, "lm_head"),
513
+ )
514
+
515
+ if self.config.tie_word_embeddings:
516
+ self.lm_head.weight = self.model.embed_tokens.weight
517
+
518
+ self.output_multiplier_scale = getattr(
519
+ config, "output_multiplier_scale", DEFAULT_OUTPUT_MULTIPLIER_SCALE)
520
+ self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
521
+ config.vocab_size,
522
+ self.output_multiplier_scale)
523
+
524
+ self.sampler = get_sampler()
525
+ self.make_empty_intermediate_tensors = (
526
+ self.model.make_empty_intermediate_tensors)
527
+
528
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
529
+ return self.model.get_input_embeddings(input_ids)
530
+
531
+ def forward(
532
+ self,
533
+ input_ids: torch.Tensor,
534
+ positions: torch.Tensor,
535
+ kv_caches: List[torch.Tensor],
536
+ attn_metadata: AttentionMetadata,
537
+ intermediate_tensors: Optional[IntermediateTensors] = None,
538
+ inputs_embeds: Optional[torch.Tensor] = None,
539
+ ) -> Union[torch.Tensor, IntermediateTensors]:
540
+ hidden_states = self.model(input_ids, positions, kv_caches,
541
+ attn_metadata, intermediate_tensors,
542
+ inputs_embeds)
543
+ return hidden_states
544
+
545
+ def compute_logits(
546
+ self,
547
+ hidden_states: torch.Tensor,
548
+ sampling_metadata: SamplingMetadata,
549
+ ) -> Optional[torch.Tensor]:
550
+ logits = self.logits_processor(self.lm_head, hidden_states,
551
+ sampling_metadata)
552
+ return logits
553
+
554
+ def sample(
555
+ self,
556
+ logits: Optional[torch.Tensor],
557
+ sampling_metadata: SamplingMetadata,
558
+ ) -> Optional[SamplerOutput]:
559
+ next_tokens = self.sampler(logits, sampling_metadata)
560
+ return next_tokens
561
+
562
+ def load_weights(self, weights: Iterable[Tuple[str,
563
+ torch.Tensor]]) -> Set[str]:
564
+ skip_prefixes = ["rotary_emb.inv_freq"]
565
+ # Skip lm_head when tie_word_embeddings is True
566
+ if self.config.tie_word_embeddings:
567
+ skip_prefixes.append("lm_head")
568
+
569
+ loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
570
+ return loader.load_weights(weights)
vllm/model_executor/models/mixtral.py
+90
-86
@@ -49,7 +49,7 @@ from vllm.model_executor.sampling_metadata import SamplingMetadata
49
from vllm.sequence import IntermediateTensors
50
51
from .interfaces import SupportsLoRA, SupportsPP
52
-from .utils import (is_pp_missing_parameter,
52
+from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
53
make_empty_intermediate_tensors_factory, make_layers,
54
maybe_prefix)
55
@@ -260,6 +260,8 @@ class MixtralModel(nn.Module):
260
quant_config = vllm_config.quant_config
261
lora_config = vllm_config.lora_config
262
263
+ self.config = config
264
+ self.quant_config = quant_config
265
lora_vocab = (lora_config.lora_extra_vocab_size *
266
(lora_config.max_loras or 1)) if lora_config else 0
267
self.vocab_size = config.vocab_size + lora_vocab
@@ -313,88 +315,6 @@ class MixtralModel(nn.Module):
315
hidden_states, _ = self.norm(hidden_states, residual)
316
return hidden_states
317
316
-
317
-class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
318
- fall_back_to_pt_during_load = False
319
-
320
- packed_modules_mapping = {
321
- "qkv_proj": [
322
- "q_proj",
323
- "k_proj",
324
- "v_proj",
325
- ],
326
- }
327
-
328
- # LoRA specific attributes
329
- embedding_modules = {
330
- "embed_tokens": "input_embeddings",
331
- "lm_head": "output_embeddings",
332
- }
333
- embedding_padding_modules = ["lm_head"]
334
-
335
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
336
- super().__init__()
337
- config = vllm_config.model_config.hf_config
338
- quant_config = vllm_config.quant_config
339
- lora_config = vllm_config.lora_config
340
- self.config = config
341
- self.lora_config = lora_config
342
- self.quant_config = quant_config
343
-
344
- self.model = MixtralModel(vllm_config=vllm_config,
345
- prefix=maybe_prefix(prefix, "model"))
346
- self.unpadded_vocab_size = config.vocab_size
347
- if lora_config:
348
- self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
349
- self.lm_head = ParallelLMHead(
350
- self.unpadded_vocab_size,
351
- config.hidden_size,
352
- org_num_embeddings=config.vocab_size,
353
- padding_size=DEFAULT_VOCAB_PADDING_SIZE
354
- # We need bigger padding if using lora for kernel
355
- # compatibility
356
- if not lora_config else lora_config.lora_vocab_padding_size,
357
- quant_config=quant_config,
358
- )
359
- if self.config.tie_word_embeddings:
360
- self.lm_head.weight = self.model.embed_tokens.weight
361
- self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
362
- config.vocab_size)
363
- self.sampler = get_sampler()
364
- self.make_empty_intermediate_tensors = (
365
- self.model.make_empty_intermediate_tensors)
366
-
367
- def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
368
- return self.model.get_input_embeddings(input_ids)
369
-
370
- def forward(
371
- self,
372
- input_ids: torch.Tensor,
373
- positions: torch.Tensor,
374
- intermediate_tensors: Optional[IntermediateTensors] = None,
375
- inputs_embeds: Optional[torch.Tensor] = None,
376
- ) -> Union[torch.Tensor, IntermediateTensors]:
377
- hidden_states = self.model(input_ids, positions, intermediate_tensors,
378
- inputs_embeds)
379
- return hidden_states
380
-
381
- def compute_logits(
382
- self,
383
- hidden_states: torch.Tensor,
384
- sampling_metadata: SamplingMetadata,
385
- ) -> Optional[torch.Tensor]:
386
- logits = self.logits_processor(self.lm_head, hidden_states,
387
- sampling_metadata)
388
- return logits
389
-
390
- def sample(
391
- self,
392
- logits: Optional[torch.Tensor],
393
- sampling_metadata: SamplingMetadata,
394
- ) -> Optional[SamplerOutput]:
395
- next_tokens = self.sampler(logits, sampling_metadata)
396
- return next_tokens
397
-
318
def load_weights(self, weights: Iterable[Tuple[str,
319
torch.Tensor]]) -> Set[str]:
320
stacked_params_mapping = [
@@ -415,9 +335,6 @@ class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
335
params_dict = dict(self.named_parameters())
336
loaded_params: Set[str] = set()
337
for name, loaded_weight in weights:
418
- if "rotary_emb.inv_freq" in name:
419
- continue
420
-
338
if (self.quant_config is not None and
339
(scale_name := self.quant_config.get_cache_scale(name))):
340
# Loading kv cache quantization scales
@@ -489,3 +406,90 @@ class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
406
weight_loader(param, loaded_weight)
407
loaded_params.add(name)
408
return loaded_params
409
+
410
+
411
+class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
412
+ fall_back_to_pt_during_load = False
413
+
414
+ packed_modules_mapping = {
415
+ "qkv_proj": [
416
+ "q_proj",
417
+ "k_proj",
418
+ "v_proj",
419
+ ],
420
+ }
421
+
422
+ # LoRA specific attributes
423
+ embedding_modules = {
424
+ "embed_tokens": "input_embeddings",
425
+ "lm_head": "output_embeddings",
426
+ }
427
+ embedding_padding_modules = ["lm_head"]
428
+
429
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
430
+ super().__init__()
431
+ config = vllm_config.model_config.hf_config
432
+ quant_config = vllm_config.quant_config
433
+ lora_config = vllm_config.lora_config
434
+ self.config = config
435
+ self.lora_config = lora_config
436
+ self.quant_config = quant_config
437
+
438
+ self.model = MixtralModel(vllm_config=vllm_config,
439
+ prefix=maybe_prefix(prefix, "model"))
440
+ self.unpadded_vocab_size = config.vocab_size
441
+ if lora_config:
442
+ self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
443
+ self.lm_head = ParallelLMHead(
444
+ self.unpadded_vocab_size,
445
+ config.hidden_size,
446
+ org_num_embeddings=config.vocab_size,
447
+ padding_size=DEFAULT_VOCAB_PADDING_SIZE
448
+ # We need bigger padding if using lora for kernel
449
+ # compatibility
450
+ if not lora_config else lora_config.lora_vocab_padding_size,
451
+ quant_config=quant_config,
452
+ )
453
+ if self.config.tie_word_embeddings:
454
+ self.lm_head.weight = self.model.embed_tokens.weight
455
+ self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
456
+ config.vocab_size)
457
+ self.sampler = get_sampler()
458
+ self.make_empty_intermediate_tensors = (
459
+ self.model.make_empty_intermediate_tensors)
460
+
461
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
462
+ return self.model.get_input_embeddings(input_ids)
463
+
464
+ def forward(
465
+ self,
466
+ input_ids: torch.Tensor,
467
+ positions: torch.Tensor,
468
+ intermediate_tensors: Optional[IntermediateTensors] = None,
469
+ inputs_embeds: Optional[torch.Tensor] = None,
470
+ ) -> Union[torch.Tensor, IntermediateTensors]:
471
+ hidden_states = self.model(input_ids, positions, intermediate_tensors,
472
+ inputs_embeds)
473
+ return hidden_states
474
+
475
+ def compute_logits(
476
+ self,
477
+ hidden_states: torch.Tensor,
478
+ sampling_metadata: SamplingMetadata,
479
+ ) -> Optional[torch.Tensor]:
480
+ logits = self.logits_processor(self.lm_head, hidden_states,
481
+ sampling_metadata)
482
+ return logits
483
+
484
+ def sample(
485
+ self,
486
+ logits: Optional[torch.Tensor],
487
+ sampling_metadata: SamplingMetadata,
488
+ ) -> Optional[SamplerOutput]:
489
+ next_tokens = self.sampler(logits, sampling_metadata)
490
+ return next_tokens
491
+
492
+ def load_weights(self, weights: Iterable[Tuple[str,
493
+ torch.Tensor]]) -> Set[str]:
494
+ loader = AutoWeightsLoader(self, skip_prefixes=["rotary_emb.inv_freq"])
495
+ return loader.load_weights(weights)