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 phimoe,qwen2_moe,qwen3_moe (#16203)
Signed-off-by: rongfu.leng <rongfu.leng@daocloud.io>
rongfu.leng committed
Apr 8, 2025 at 19:05 UTC
5a1e1c8353b9c3b5356d3ac071e01e310f4f8c69
3 files changed
+220
-198
vllm/model_executor/models/phimoe.py
+90
-83
@@ -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
@@ -448,6 +448,8 @@ class PhiMoEModel(nn.Module):
448
(lora_config.max_loras or 1)) if lora_config else 0)
449
self.vocab_size = config.vocab_size + lora_vocab
450
self.org_vocab_size = config.vocab_size
451
+ self.config = config
452
+ self.quant_config = quant_config
453
454
self.embed_tokens = VocabParallelEmbedding(
455
self.vocab_size,
@@ -504,85 +506,6 @@ class PhiMoEModel(nn.Module):
506
hidden_states = self.norm(hidden_states)
507
return hidden_states
508
507
-
508
-class PhiMoEForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
509
- fall_back_to_pt_during_load = False
510
-
511
- packed_modules_mapping = {
512
- "qkv_proj": [
513
- "q_proj",
514
- "k_proj",
515
- "v_proj",
516
- ],
517
- }
518
-
519
- # LoRA specific attributes
520
- embedding_modules = {
521
- "embed_tokens": "input_embeddings",
522
- "lm_head": "output_embeddings",
523
- }
524
- embedding_padding_modules = ["lm_head"]
525
-
526
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
527
- super().__init__()
528
- config = vllm_config.model_config.hf_config
529
- lora_config = vllm_config.lora_config
530
- self.config = config
531
- self.lora_config = lora_config
532
- self.quant_config = vllm_config.quant_config
533
-
534
- self.model = PhiMoEModel(vllm_config=vllm_config,
535
- prefix=maybe_prefix(prefix, "model"))
536
- self.unpadded_vocab_size = config.vocab_size
537
- if lora_config:
538
- self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
539
- self.lm_head = ParallelLMHead(
540
- self.unpadded_vocab_size,
541
- config.hidden_size,
542
- org_num_embeddings=config.vocab_size,
543
- padding_size=(
544
- DEFAULT_VOCAB_PADDING_SIZE
545
- # We need bigger padding if using lora for kernel
546
- # compatibility
547
- if not lora_config else lora_config.lora_vocab_padding_size),
548
- quant_config=None,
549
- bias=True,
550
- )
551
- self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
552
- config.vocab_size)
553
- self.sampler = get_sampler()
554
-
555
- self.make_empty_intermediate_tensors = (
556
- self.model.make_empty_intermediate_tensors)
557
-
558
- def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
559
- return self.model.get_input_embeddings(input_ids)
560
-
561
- def forward(
562
- self,
563
- input_ids: torch.Tensor,
564
- positions: torch.Tensor,
565
- intermediate_tensors: Optional[IntermediateTensors] = None,
566
- inputs_embeds: Optional[torch.Tensor] = None,
567
- ) -> Union[torch.Tensor, IntermediateTensors]:
568
- hidden_states = self.model(input_ids, positions, intermediate_tensors,
569
- inputs_embeds)
570
- return hidden_states
571
-
572
- def compute_logits(self, hidden_states: torch.Tensor,
573
- sampling_metadata: SamplingMetadata) -> torch.Tensor:
574
- logits = self.logits_processor(self.lm_head, hidden_states,
575
- sampling_metadata)
576
- return logits
577
-
578
- def sample(
579
- self,
580
- logits: Optional[torch.Tensor],
581
- sampling_metadata: SamplingMetadata,
582
- ) -> Optional[SamplerOutput]:
583
- next_tokens = self.sampler(logits, sampling_metadata)
584
- return next_tokens
585
-
509
def load_weights(self, weights: Iterable[Tuple[str,
510
torch.Tensor]]) -> Set[str]:
511
stacked_params_mapping = [
@@ -601,9 +524,6 @@ class PhiMoEForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
524
params_dict = dict(self.named_parameters())
525
loaded_params: Set[str] = set()
526
for name, loaded_weight in weights:
604
- if "rotary_emb.inv_freq" in name:
605
- continue
606
-
527
if (self.quant_config is not None and
528
(scale_name := self.quant_config.get_cache_scale(name))):
529
# Loading kv cache quantization scales
@@ -667,3 +587,90 @@ class PhiMoEForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
587
weight_loader(param, loaded_weight)
588
loaded_params.add(name)
589
return loaded_params
590
+
591
+
592
+class PhiMoEForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
593
+ fall_back_to_pt_during_load = False
594
+
595
+ packed_modules_mapping = {
596
+ "qkv_proj": [
597
+ "q_proj",
598
+ "k_proj",
599
+ "v_proj",
600
+ ],
601
+ }
602
+
603
+ # LoRA specific attributes
604
+ embedding_modules = {
605
+ "embed_tokens": "input_embeddings",
606
+ "lm_head": "output_embeddings",
607
+ }
608
+ embedding_padding_modules = ["lm_head"]
609
+
610
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
611
+ super().__init__()
612
+ config = vllm_config.model_config.hf_config
613
+ lora_config = vllm_config.lora_config
614
+ self.config = config
615
+ self.lora_config = lora_config
616
+ self.quant_config = vllm_config.quant_config
617
+
618
+ self.model = PhiMoEModel(vllm_config=vllm_config,
619
+ prefix=maybe_prefix(prefix, "model"))
620
+ self.unpadded_vocab_size = config.vocab_size
621
+ if lora_config:
622
+ self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
623
+ self.lm_head = ParallelLMHead(
624
+ self.unpadded_vocab_size,
625
+ config.hidden_size,
626
+ org_num_embeddings=config.vocab_size,
627
+ padding_size=(
628
+ DEFAULT_VOCAB_PADDING_SIZE
629
+ # We need bigger padding if using lora for kernel
630
+ # compatibility
631
+ if not lora_config else lora_config.lora_vocab_padding_size),
632
+ quant_config=None,
633
+ bias=True,
634
+ )
635
+ self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
636
+ config.vocab_size)
637
+ self.sampler = get_sampler()
638
+
639
+ self.make_empty_intermediate_tensors = (
640
+ self.model.make_empty_intermediate_tensors)
641
+
642
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
643
+ return self.model.get_input_embeddings(input_ids)
644
+
645
+ def forward(
646
+ self,
647
+ input_ids: torch.Tensor,
648
+ positions: torch.Tensor,
649
+ intermediate_tensors: Optional[IntermediateTensors] = None,
650
+ inputs_embeds: Optional[torch.Tensor] = None,
651
+ ) -> Union[torch.Tensor, IntermediateTensors]:
652
+ hidden_states = self.model(input_ids, positions, intermediate_tensors,
653
+ inputs_embeds)
654
+ return hidden_states
655
+
656
+ def compute_logits(self, hidden_states: torch.Tensor,
657
+ sampling_metadata: SamplingMetadata) -> torch.Tensor:
658
+ logits = self.logits_processor(self.lm_head, hidden_states,
659
+ sampling_metadata)
660
+ return logits
661
+
662
+ def sample(
663
+ self,
664
+ logits: Optional[torch.Tensor],
665
+ sampling_metadata: SamplingMetadata,
666
+ ) -> Optional[SamplerOutput]:
667
+ next_tokens = self.sampler(logits, sampling_metadata)
668
+ return next_tokens
669
+
670
+ def load_weights(self, weights: Iterable[Tuple[str,
671
+ torch.Tensor]]) -> Set[str]:
672
+ loader = AutoWeightsLoader(
673
+ self,
674
+ skip_prefixes=(["rotary_emb.inv_freq"]),
675
+ )
676
+ return loader.load_weights(weights)
vllm/model_executor/models/qwen2_moe.py
+65
-57
@@ -55,7 +55,8 @@ from vllm.model_executor.sampling_metadata import SamplingMetadata
55
from vllm.sequence import IntermediateTensors
56
57
from .interfaces import SupportsPP
58
-from .utils import (extract_layer_index, is_pp_missing_parameter,
58
+from .utils import (AutoWeightsLoader, extract_layer_index,
59
+ is_pp_missing_parameter,
60
make_empty_intermediate_tensors_factory, make_layers,
61
maybe_prefix)
62
@@ -329,6 +330,7 @@ class Qwen2MoeModel(nn.Module):
330
quant_config = vllm_config.quant_config
331
332
self.vocab_size = config.vocab_size
333
+ self.config = config
334
335
self.embed_tokens = VocabParallelEmbedding(
336
config.vocab_size,
@@ -377,60 +379,6 @@ class Qwen2MoeModel(nn.Module):
379
hidden_states, _ = self.norm(hidden_states, residual)
380
return hidden_states
381
380
-
381
-class Qwen2MoeForCausalLM(nn.Module, SupportsPP):
382
-
383
- fall_back_to_pt_during_load = False
384
-
385
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
386
- super().__init__()
387
- config = vllm_config.model_config.hf_config
388
- quant_config = vllm_config.quant_config
389
- self.config = config
390
- self.quant_config = quant_config
391
- self.model = Qwen2MoeModel(vllm_config=vllm_config,
392
- prefix=maybe_prefix(prefix, "model"))
393
- self.lm_head = ParallelLMHead(config.vocab_size,
394
- config.hidden_size,
395
- quant_config=quant_config)
396
- if self.config.tie_word_embeddings:
397
- self.lm_head.weight = self.model.embed_tokens.weight
398
- self.logits_processor = LogitsProcessor(config.vocab_size)
399
- self.sampler = get_sampler()
400
- self.make_empty_intermediate_tensors = (
401
- self.model.make_empty_intermediate_tensors)
402
-
403
- def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
404
- return self.model.get_input_embeddings(input_ids)
405
-
406
- def forward(
407
- self,
408
- input_ids: torch.Tensor,
409
- positions: torch.Tensor,
410
- intermediate_tensors: Optional[IntermediateTensors] = None,
411
- inputs_embeds: Optional[torch.Tensor] = None,
412
- ) -> Union[torch.Tensor, IntermediateTensors]:
413
- hidden_states = self.model(input_ids, positions, intermediate_tensors,
414
- inputs_embeds)
415
- return hidden_states
416
-
417
- def compute_logits(
418
- self,
419
- hidden_states: torch.Tensor,
420
- sampling_metadata: SamplingMetadata,
421
- ) -> Optional[torch.Tensor]:
422
- logits = self.logits_processor(self.lm_head, hidden_states,
423
- sampling_metadata)
424
- return logits
425
-
426
- def sample(
427
- self,
428
- logits: Optional[torch.Tensor],
429
- sampling_metadata: SamplingMetadata,
430
- ) -> Optional[SamplerOutput]:
431
- next_tokens = self.sampler(logits, sampling_metadata)
432
- return next_tokens
433
-
382
def load_weights(self, weights: Iterable[Tuple[str,
383
torch.Tensor]]) -> Set[str]:
384
stacked_params_mapping = [
@@ -453,8 +401,6 @@ class Qwen2MoeForCausalLM(nn.Module, SupportsPP):
401
params_dict = dict(self.named_parameters())
402
loaded_params: Set[str] = set()
403
for name, loaded_weight in weights:
456
- if "rotary_emb.inv_freq" in name:
457
- continue
404
for (param_name, weight_name, shard_id) in stacked_params_mapping:
405
# Skip non-stacked layers and experts (experts handled below).
406
if weight_name not in name:
@@ -531,3 +477,65 @@ class Qwen2MoeForCausalLM(nn.Module, SupportsPP):
477
weight_loader(param, loaded_weight)
478
loaded_params.add(name)
479
return loaded_params
480
+
481
+
482
+class Qwen2MoeForCausalLM(nn.Module, SupportsPP):
483
+
484
+ fall_back_to_pt_during_load = False
485
+
486
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
487
+ super().__init__()
488
+ config = vllm_config.model_config.hf_config
489
+ quant_config = vllm_config.quant_config
490
+ self.config = config
491
+ self.quant_config = quant_config
492
+ self.model = Qwen2MoeModel(vllm_config=vllm_config,
493
+ prefix=maybe_prefix(prefix, "model"))
494
+ self.lm_head = ParallelLMHead(config.vocab_size,
495
+ config.hidden_size,
496
+ quant_config=quant_config)
497
+ if self.config.tie_word_embeddings:
498
+ self.lm_head.weight = self.model.embed_tokens.weight
499
+ self.logits_processor = LogitsProcessor(config.vocab_size)
500
+ self.sampler = get_sampler()
501
+ self.make_empty_intermediate_tensors = (
502
+ self.model.make_empty_intermediate_tensors)
503
+
504
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
505
+ return self.model.get_input_embeddings(input_ids)
506
+
507
+ def forward(
508
+ self,
509
+ input_ids: torch.Tensor,
510
+ positions: torch.Tensor,
511
+ intermediate_tensors: Optional[IntermediateTensors] = None,
512
+ inputs_embeds: Optional[torch.Tensor] = None,
513
+ ) -> Union[torch.Tensor, IntermediateTensors]:
514
+ hidden_states = self.model(input_ids, positions, intermediate_tensors,
515
+ inputs_embeds)
516
+ return hidden_states
517
+
518
+ def compute_logits(
519
+ self,
520
+ hidden_states: torch.Tensor,
521
+ sampling_metadata: SamplingMetadata,
522
+ ) -> Optional[torch.Tensor]:
523
+ logits = self.logits_processor(self.lm_head, hidden_states,
524
+ sampling_metadata)
525
+ return logits
526
+
527
+ def sample(
528
+ self,
529
+ logits: Optional[torch.Tensor],
530
+ sampling_metadata: SamplingMetadata,
531
+ ) -> Optional[SamplerOutput]:
532
+ next_tokens = self.sampler(logits, sampling_metadata)
533
+ return next_tokens
534
+
535
+ def load_weights(self, weights: Iterable[Tuple[str,
536
+ torch.Tensor]]) -> Set[str]:
537
+ loader = AutoWeightsLoader(
538
+ self,
539
+ skip_prefixes=(["rotary_emb.inv_freq"]),
540
+ )
541
+ return loader.load_weights(weights)
vllm/model_executor/models/qwen3_moe.py
+65
-58
@@ -52,7 +52,8 @@ from vllm.model_executor.sampling_metadata import SamplingMetadata
52
from vllm.sequence import IntermediateTensors
53
54
from .interfaces import SupportsPP
55
-from .utils import (extract_layer_index, is_pp_missing_parameter,
55
+from .utils import (AutoWeightsLoader, extract_layer_index,
56
+ is_pp_missing_parameter,
57
make_empty_intermediate_tensors_factory, make_layers,
58
maybe_prefix)
59
@@ -326,7 +327,7 @@ class Qwen3MoeModel(nn.Module):
327
328
self.padding_idx = config.pad_token_id
329
self.vocab_size = config.vocab_size
329
-
330
+ self.config = config
331
self.embed_tokens = VocabParallelEmbedding(
332
config.vocab_size,
333
config.hidden_size,
@@ -375,60 +376,6 @@ class Qwen3MoeModel(nn.Module):
376
hidden_states, _ = self.norm(hidden_states, residual)
377
return hidden_states
378
378
-
379
-class Qwen3MoeForCausalLM(nn.Module, SupportsPP):
380
-
381
- fall_back_to_pt_during_load = False
382
-
383
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
384
- super().__init__()
385
- config = vllm_config.model_config.hf_config
386
- quant_config = vllm_config.quant_config
387
- self.config = config
388
- self.quant_config = quant_config
389
- self.model = Qwen3MoeModel(vllm_config=vllm_config,
390
- prefix=maybe_prefix(prefix, "model"))
391
- self.lm_head = ParallelLMHead(config.vocab_size,
392
- config.hidden_size,
393
- quant_config=quant_config)
394
- if self.config.tie_word_embeddings:
395
- self.lm_head.weight = self.model.embed_tokens.weight
396
- self.logits_processor = LogitsProcessor(config.vocab_size)
397
- self.sampler = get_sampler()
398
- self.make_empty_intermediate_tensors = (
399
- self.model.make_empty_intermediate_tensors)
400
-
401
- def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
402
- return self.model.get_input_embeddings(input_ids)
403
-
404
- def forward(
405
- self,
406
- input_ids: torch.Tensor,
407
- positions: torch.Tensor,
408
- intermediate_tensors: Optional[IntermediateTensors] = None,
409
- inputs_embeds: Optional[torch.Tensor] = None,
410
- ) -> Union[torch.Tensor, IntermediateTensors]:
411
- hidden_states = self.model(input_ids, positions, intermediate_tensors,
412
- inputs_embeds)
413
- return hidden_states
414
-
415
- def compute_logits(
416
- self,
417
- hidden_states: torch.Tensor,
418
- sampling_metadata: SamplingMetadata,
419
- ) -> Optional[torch.Tensor]:
420
- logits = self.logits_processor(self.lm_head, hidden_states,
421
- sampling_metadata)
422
- return logits
423
-
424
- def sample(
425
- self,
426
- logits: Optional[torch.Tensor],
427
- sampling_metadata: SamplingMetadata,
428
- ) -> Optional[SamplerOutput]:
429
- next_tokens = self.sampler(logits, sampling_metadata)
430
- return next_tokens
431
-
379
def load_weights(self, weights: Iterable[Tuple[str,
380
torch.Tensor]]) -> Set[str]:
381
stacked_params_mapping = [
@@ -451,8 +398,6 @@ class Qwen3MoeForCausalLM(nn.Module, SupportsPP):
398
params_dict = dict(self.named_parameters())
399
loaded_params: Set[str] = set()
400
for name, loaded_weight in weights:
454
- if "rotary_emb.inv_freq" in name:
455
- continue
401
for (param_name, weight_name, shard_id) in stacked_params_mapping:
402
# Skip non-stacked layers and experts (experts handled below).
403
if weight_name not in name:
@@ -529,3 +474,65 @@ class Qwen3MoeForCausalLM(nn.Module, SupportsPP):
474
weight_loader(param, loaded_weight)
475
loaded_params.add(name)
476
return loaded_params
477
+
478
+
479
+class Qwen3MoeForCausalLM(nn.Module, SupportsPP):
480
+
481
+ fall_back_to_pt_during_load = False
482
+
483
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
484
+ super().__init__()
485
+ config = vllm_config.model_config.hf_config
486
+ quant_config = vllm_config.quant_config
487
+ self.config = config
488
+ self.quant_config = quant_config
489
+ self.model = Qwen3MoeModel(vllm_config=vllm_config,
490
+ prefix=maybe_prefix(prefix, "model"))
491
+ self.lm_head = ParallelLMHead(config.vocab_size,
492
+ config.hidden_size,
493
+ quant_config=quant_config)
494
+ if self.config.tie_word_embeddings:
495
+ self.lm_head.weight = self.model.embed_tokens.weight
496
+ self.logits_processor = LogitsProcessor(config.vocab_size)
497
+ self.sampler = get_sampler()
498
+ self.make_empty_intermediate_tensors = (
499
+ self.model.make_empty_intermediate_tensors)
500
+
501
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
502
+ return self.model.get_input_embeddings(input_ids)
503
+
504
+ def forward(
505
+ self,
506
+ input_ids: torch.Tensor,
507
+ positions: torch.Tensor,
508
+ intermediate_tensors: Optional[IntermediateTensors] = None,
509
+ inputs_embeds: Optional[torch.Tensor] = None,
510
+ ) -> Union[torch.Tensor, IntermediateTensors]:
511
+ hidden_states = self.model(input_ids, positions, intermediate_tensors,
512
+ inputs_embeds)
513
+ return hidden_states
514
+
515
+ def compute_logits(
516
+ self,
517
+ hidden_states: torch.Tensor,
518
+ sampling_metadata: SamplingMetadata,
519
+ ) -> Optional[torch.Tensor]:
520
+ logits = self.logits_processor(self.lm_head, hidden_states,
521
+ sampling_metadata)
522
+ return logits
523
+
524
+ def sample(
525
+ self,
526
+ logits: Optional[torch.Tensor],
527
+ sampling_metadata: SamplingMetadata,
528
+ ) -> Optional[SamplerOutput]:
529
+ next_tokens = self.sampler(logits, sampling_metadata)
530
+ return next_tokens
531
+
532
+ def load_weights(self, weights: Iterable[Tuple[str,
533
+ torch.Tensor]]) -> Set[str]:
534
+ loader = AutoWeightsLoader(
535
+ self,
536
+ skip_prefixes=(["rotary_emb.inv_freq"]),
537
+ )
538
+ return loader.load_weights(weights)