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, 01:00 UTC
[Quantization] Enable BNB support for more MoE models (#21100)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Jul 19, 2025 at 08:52 UTC
466e878f2ad5e36cba4861db1cac7cd0d92055fb
5 files changed
+223
-181
docs/models/supported_models.md
+4
-4
@@ -316,7 +316,7 @@ Specified using `--task generate`.
316
| `AquilaForCausalLM` | Aquila, Aquila2 | `BAAI/Aquila-7B`, `BAAI/AquilaChat-7B`, etc. | ✅︎ | ✅︎ | ✅︎ |
317
| `ArcticForCausalLM` | Arctic | `Snowflake/snowflake-arctic-base`, `Snowflake/snowflake-arctic-instruct`, etc. | | ✅︎ | ✅︎ |
318
| `BaiChuanForCausalLM` | Baichuan2, Baichuan | `baichuan-inc/Baichuan2-13B-Chat`, `baichuan-inc/Baichuan-7B`, etc. | ✅︎ | ✅︎ | ✅︎ |
319
-| `BailingMoeForCausalLM` | Ling | `inclusionAI/Ling-lite-1.5`, `inclusionAI/Ling-plus`, etc. | | ✅︎ | ✅︎ |
319
+| `BailingMoeForCausalLM` | Ling | `inclusionAI/Ling-lite-1.5`, `inclusionAI/Ling-plus`, etc. | ✅︎ | ✅︎ | ✅︎ |
320
| `BambaForCausalLM` | Bamba | `ibm-ai-platform/Bamba-9B-fp8`, `ibm-ai-platform/Bamba-9B` | ✅︎ | ✅︎ | ✅︎ |
321
| `BloomForCausalLM` | BLOOM, BLOOMZ, BLOOMChat | `bigscience/bloom`, `bigscience/bloomz`, etc. | | ✅︎ | |
322
| `BartForConditionalGeneration` | BART | `facebook/bart-base`, `facebook/bart-large-cnn`, etc. | | | |
@@ -328,8 +328,8 @@ Specified using `--task generate`.
328
| `DeepseekV2ForCausalLM` | DeepSeek-V2 | `deepseek-ai/DeepSeek-V2`, `deepseek-ai/DeepSeek-V2-Chat`, etc. | | ✅︎ | ✅︎ |
329
| `DeepseekV3ForCausalLM` | DeepSeek-V3 | `deepseek-ai/DeepSeek-V3-Base`, `deepseek-ai/DeepSeek-V3`, etc. | | ✅︎ | ✅︎ |
330
| `Dots1ForCausalLM` | dots.llm1 | `rednote-hilab/dots.llm1.base`, `rednote-hilab/dots.llm1.inst`, etc. | | ✅︎ | ✅︎ |
331
-| `Ernie4_5_ForCausalLM` | Ernie4.5 | `baidu/ERNIE-4.5-0.3B-PT`, etc. | | ✅︎ | ✅︎ |
332
-| `Ernie4_5_MoeForCausalLM` | Ernie4.5MoE | `baidu/ERNIE-4.5-21B-A3B-PT`, `baidu/ERNIE-4.5-300B-A47B-PT`, etc. | | ✅︎ | ✅︎ |
331
+| `Ernie4_5_ForCausalLM` | Ernie4.5 | `baidu/ERNIE-4.5-0.3B-PT`, etc. | ✅︎ | ✅︎ | ✅︎ |
332
+| `Ernie4_5_MoeForCausalLM` | Ernie4.5MoE | `baidu/ERNIE-4.5-21B-A3B-PT`, `baidu/ERNIE-4.5-300B-A47B-PT`, etc. |✅︎| ✅︎ | ✅︎ |
333
| `ExaoneForCausalLM` | EXAONE-3 | `LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct`, etc. | ✅︎ | ✅︎ | ✅︎ |
334
| `Fairseq2LlamaForCausalLM` | Llama (fairseq2 format) | `mgleize/fairseq2-dummy-Llama-3.2-1B`, etc. | ✅︎ | ✅︎ | ✅︎ |
335
| `FalconForCausalLM` | Falcon | `tiiuae/falcon-7b`, `tiiuae/falcon-40b`, `tiiuae/falcon-rw-7b`, etc. | | ✅︎ | ✅︎ |
@@ -351,7 +351,7 @@ Specified using `--task generate`.
351
| `GraniteMoeSharedForCausalLM` | Granite MoE Shared | `ibm-research/moe-7b-1b-active-shared-experts` (test model) | ✅︎ | ✅︎ | ✅︎ |
352
| `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ | |
353
| `Grok1ModelForCausalLM` | Grok1 | `hpcai-tech/grok-1`. | ✅︎ | ✅︎ | ✅︎ |
354
-| `HunYuanMoEV1ForCausalLM` | Hunyuan-80B-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | | | ✅︎ |
354
+| `HunYuanMoEV1ForCausalLM` | Hunyuan-80B-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | | ✅︎ |
355
| `InternLMForCausalLM` | InternLM | `internlm/internlm-7b`, `internlm/internlm-chat-7b`, etc. | ✅︎ | ✅︎ | ✅︎ |
356
| `InternLM2ForCausalLM` | InternLM2 | `internlm/internlm2-7b`, `internlm/internlm2-chat-7b`, etc. | ✅︎ | ✅︎ | ✅︎ |
357
| `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ | ✅︎ |
vllm/model_executor/models/bailing_moe.py
+14
-7
@@ -53,7 +53,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
53
from vllm.model_executor.sampling_metadata import SamplingMetadata
54
from vllm.sequence import IntermediateTensors
55
56
-from .interfaces import SupportsPP
56
+from .interfaces import SupportsLoRA, SupportsPP
57
from .utils import (AutoWeightsLoader, PPMissingLayer, is_pp_missing_parameter,
58
make_empty_intermediate_tensors_factory, make_layers,
59
maybe_prefix)
@@ -374,6 +374,14 @@ class BailingMoeModel(nn.Module):
374
hidden_states, _ = self.norm(hidden_states, residual)
375
return hidden_states
376
377
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
378
+ return FusedMoE.make_expert_params_mapping(
379
+ ckpt_gate_proj_name="gate_proj",
380
+ ckpt_down_proj_name="down_proj",
381
+ ckpt_up_proj_name="up_proj",
382
+ num_experts=self.config.num_experts,
383
+ )
384
+
385
def load_weights(self, weights: Iterable[tuple[str,
386
torch.Tensor]]) -> set[str]:
387
stacked_params_mapping = [
@@ -381,14 +389,10 @@ class BailingMoeModel(nn.Module):
389
("gate_up_proj", "gate_proj", 0),
390
("gate_up_proj", "up_proj", 1),
391
]
384
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
385
- ckpt_gate_proj_name="gate_proj",
386
- ckpt_down_proj_name="down_proj",
387
- ckpt_up_proj_name="up_proj",
388
- num_experts=self.config.num_experts)
392
393
params_dict = dict(self.named_parameters(remove_duplicate=False))
394
loaded_params: set[str] = set()
395
+ expert_params_mapping = self.get_expert_mapping()
396
for name, loaded_weight in weights:
397
if self.config.norm_head and "lm_head.weight" in name:
398
loaded_weight = F.normalize(loaded_weight,
@@ -449,7 +453,7 @@ class BailingMoeModel(nn.Module):
453
return loaded_params
454
455
452
-class BailingMoeForCausalLM(nn.Module, SupportsPP):
456
+class BailingMoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
457
458
packed_modules_mapping = {
459
"query_key_value": ["query_key_value"],
@@ -518,3 +522,6 @@ class BailingMoeForCausalLM(nn.Module, SupportsPP):
522
if self.config.tie_word_embeddings else None),
523
)
524
return loader.load_weights(weights)
525
+
526
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
527
+ return self.model.get_expert_mapping()
vllm/model_executor/models/ernie45_moe.py
+84
-69
@@ -51,8 +51,8 @@ from vllm.model_executor.model_loader.weight_utils import (
51
from vllm.model_executor.sampling_metadata import SamplingMetadata
52
from vllm.sequence import IntermediateTensors
53
54
-from .interfaces import SupportsPP
55
-from .utils import (PPMissingLayer, extract_layer_index,
54
+from .interfaces import SupportsLoRA, SupportsPP
55
+from .utils import (AutoWeightsLoader, PPMissingLayer, extract_layer_index,
56
is_pp_missing_parameter,
57
make_empty_intermediate_tensors_factory, make_layers,
58
maybe_prefix)
@@ -427,66 +427,15 @@ class Ernie4_5_MoeModel(nn.Module):
427
428
return hidden_states
429
430
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
431
431
-class Ernie4_5_MoeForCausalLM(nn.Module, SupportsPP):
432
- packed_modules_mapping = {
433
- "qkv_proj": [
434
- "q_proj",
435
- "k_proj",
436
- "v_proj",
437
- ],
438
- "gate_up_proj": [
439
- "gate_proj",
440
- "up_proj",
441
- ],
442
- }
443
-
444
- fall_back_to_pt_during_load = False
445
-
446
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
447
- super().__init__()
448
- config = vllm_config.model_config.hf_config
449
- quant_config = vllm_config.quant_config
450
- self.config = config
451
- self.quant_config = quant_config
452
- self.model = Ernie4_5_MoeModel(vllm_config=vllm_config,
453
- prefix=maybe_prefix(prefix, "model"))
454
-
455
- if get_pp_group().is_last_rank:
456
- self.lm_head = ParallelLMHead(config.vocab_size,
457
- config.hidden_size,
458
- quant_config=quant_config)
459
- else:
460
- self.lm_head = PPMissingLayer()
461
-
462
- if self.config.tie_word_embeddings:
463
- self.lm_head.weight = self.model.embed_tokens.weight
464
- self.logits_processor = LogitsProcessor(config.vocab_size)
465
- self.make_empty_intermediate_tensors = (
466
- self.model.make_empty_intermediate_tensors)
467
-
468
- def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
469
- return self.model.get_input_embeddings(input_ids)
470
-
471
- def forward(
472
- self,
473
- input_ids: torch.Tensor,
474
- positions: torch.Tensor,
475
- intermediate_tensors: Optional[IntermediateTensors] = None,
476
- inputs_embeds: Optional[torch.Tensor] = None,
477
- ) -> Union[torch.Tensor, IntermediateTensors]:
478
- hidden_states = self.model(input_ids, positions, intermediate_tensors,
479
- inputs_embeds)
480
- return hidden_states
481
-
482
- def compute_logits(
483
- self,
484
- hidden_states: torch.Tensor,
485
- sampling_metadata: SamplingMetadata,
486
- ) -> Optional[torch.Tensor]:
487
- logits = self.logits_processor(self.lm_head, hidden_states,
488
- sampling_metadata)
489
- return logits
432
+ # Params for weights, fp8 weight scales, fp8 activation scales
433
+ # (param_name, weight_name, expert_id, shard_id)
434
+ return FusedMoE.make_expert_params_mapping(
435
+ ckpt_gate_proj_name="gate_proj",
436
+ ckpt_down_proj_name="down_proj",
437
+ ckpt_up_proj_name="up_proj",
438
+ num_experts=self.config.moe_num_experts)
439
440
def load_weights(self, weights: Iterable[tuple[str,
441
torch.Tensor]]) -> set[str]:
@@ -499,16 +448,9 @@ class Ernie4_5_MoeForCausalLM(nn.Module, SupportsPP):
448
("gate_up_proj", "up_proj", 1),
449
]
450
502
- # Params for weights, fp8 weight scales, fp8 activation scales
503
- # (param_name, weight_name, expert_id, shard_id)
504
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
505
- ckpt_gate_proj_name="gate_proj",
506
- ckpt_down_proj_name="down_proj",
507
- ckpt_up_proj_name="up_proj",
508
- num_experts=self.config.moe_num_experts)
509
-
451
params_dict = dict(self.named_parameters())
452
loaded_params: set[str] = set()
453
+ expert_params_mapping = self.get_expert_mapping()
454
for name, loaded_weight in weights:
455
if self.config.tie_word_embeddings and name.endswith(
456
"lm_head.weight"):
@@ -581,3 +523,76 @@ class Ernie4_5_MoeForCausalLM(nn.Module, SupportsPP):
523
weight_loader(param, loaded_weight)
524
loaded_params.add(name)
525
return loaded_params
526
+
527
+
528
+class Ernie4_5_MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
529
+ packed_modules_mapping = {
530
+ "qkv_proj": [
531
+ "q_proj",
532
+ "k_proj",
533
+ "v_proj",
534
+ ],
535
+ "gate_up_proj": [
536
+ "gate_proj",
537
+ "up_proj",
538
+ ],
539
+ }
540
+
541
+ fall_back_to_pt_during_load = False
542
+
543
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
544
+ super().__init__()
545
+ config = vllm_config.model_config.hf_config
546
+ quant_config = vllm_config.quant_config
547
+ self.config = config
548
+ self.quant_config = quant_config
549
+ self.model = Ernie4_5_MoeModel(vllm_config=vllm_config,
550
+ prefix=maybe_prefix(prefix, "model"))
551
+
552
+ if get_pp_group().is_last_rank:
553
+ self.lm_head = ParallelLMHead(config.vocab_size,
554
+ config.hidden_size,
555
+ quant_config=quant_config)
556
+ else:
557
+ self.lm_head = PPMissingLayer()
558
+
559
+ if self.config.tie_word_embeddings:
560
+ self.lm_head.weight = self.model.embed_tokens.weight
561
+ self.logits_processor = LogitsProcessor(config.vocab_size)
562
+ self.make_empty_intermediate_tensors = (
563
+ self.model.make_empty_intermediate_tensors)
564
+
565
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
566
+ return self.model.get_input_embeddings(input_ids)
567
+
568
+ def forward(
569
+ self,
570
+ input_ids: torch.Tensor,
571
+ positions: torch.Tensor,
572
+ intermediate_tensors: Optional[IntermediateTensors] = None,
573
+ inputs_embeds: Optional[torch.Tensor] = None,
574
+ ) -> Union[torch.Tensor, IntermediateTensors]:
575
+ hidden_states = self.model(input_ids, positions, intermediate_tensors,
576
+ inputs_embeds)
577
+ return hidden_states
578
+
579
+ def compute_logits(
580
+ self,
581
+ hidden_states: torch.Tensor,
582
+ sampling_metadata: SamplingMetadata,
583
+ ) -> Optional[torch.Tensor]:
584
+ logits = self.logits_processor(self.lm_head, hidden_states,
585
+ sampling_metadata)
586
+ return logits
587
+
588
+ def load_weights(self, weights: Iterable[tuple[str,
589
+ torch.Tensor]]) -> set[str]:
590
+ loader = AutoWeightsLoader(
591
+ self,
592
+ skip_prefixes=(["lm_head."]
593
+ if self.config.tie_word_embeddings else None),
594
+ )
595
+ return loader.load_weights(weights)
596
+
597
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
598
+ return self.model.get_expert_mapping()
vllm/model_executor/models/grok1.py
+14
-10
@@ -360,6 +360,16 @@ class Grok1Model(nn.Module):
360
hidden_states, _ = self.norm(hidden_states, residual)
361
return hidden_states
362
363
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
364
+ # Map Grok1's unique expert parameter names to standard names
365
+ # Grok1 uses "num_experts" in its config
366
+ num_experts = getattr(self.config, "num_experts", 8)
367
+ return FusedMoE.make_expert_params_mapping(
368
+ ckpt_gate_proj_name="linear", # Grok1 specific
369
+ ckpt_down_proj_name="linear_1", # Grok1 specific
370
+ ckpt_up_proj_name="linear_v", # Grok1 specific
371
+ num_experts=num_experts)
372
+
373
def load_weights(self, weights: Iterable[tuple[str,
374
torch.Tensor]]) -> set[str]:
375
stacked_params_mapping = [
@@ -369,18 +379,9 @@ class Grok1Model(nn.Module):
379
("qkv_proj", "v_proj", "v"),
380
]
381
372
- # Map Grok1's unique expert parameter names to standard names
373
- # Grok1 uses "num_experts" in its config
374
- num_experts = getattr(self.config, "num_experts", 8)
375
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
376
- ckpt_gate_proj_name="linear", # Grok1 specific
377
- ckpt_down_proj_name="linear_1", # Grok1 specific
378
- ckpt_up_proj_name="linear_v", # Grok1 specific
379
- num_experts=num_experts)
380
-
382
params_dict = dict(self.named_parameters())
383
loaded_params: set[str] = set()
383
-
384
+ expert_params_mapping = self.get_expert_mapping()
385
for name, loaded_weight in weights:
386
if (self.quant_config is not None and
387
(scale_name := self.quant_config.get_cache_scale(name))):
@@ -544,3 +545,6 @@ class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
545
skip_prefixes=skip_prefixes,
546
)
547
return loader.load_weights(weights)
548
+
549
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
550
+ return self.model.get_expert_mapping()
vllm/model_executor/models/hunyuan_v1_moe.py
+107
-91
@@ -56,7 +56,9 @@ from vllm.model_executor.model_loader.weight_utils import (
56
from vllm.model_executor.sampling_metadata import SamplingMetadata
57
from vllm.sequence import IntermediateTensors
58
59
-from .utils import PPMissingLayer, is_pp_missing_parameter, make_layers
59
+from .interfaces import SupportsLoRA
60
+from .utils import (AutoWeightsLoader, PPMissingLayer, is_pp_missing_parameter,
61
+ make_layers)
62
63
64
def _get_cla_factor(config: PretrainedConfig) -> int:
@@ -617,86 +619,6 @@ class HunYuanModel(nn.Module):
619
hidden_states, _ = self.norm(hidden_states, residual)
620
return hidden_states
621
620
-
621
-class HunYuanMoEV1ForCausalLM(nn.Module):
622
- packed_modules_mapping = {
623
- "qkv_proj": [
624
- "q_proj",
625
- "k_proj",
626
- "v_proj",
627
- ],
628
- "gate_up_proj": [
629
- "gate_proj",
630
- "up_proj",
631
- ],
632
- }
633
-
634
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
635
- super().__init__()
636
-
637
- config = vllm_config.model_config.hf_config
638
- quant_config = vllm_config.quant_config
639
- lora_config = vllm_config.lora_config
640
- self.config = config
641
- self.quant_config = quant_config
642
- self.lora_config = lora_config
643
-
644
- self.model = HunYuanModel(vllm_config=vllm_config, prefix="model")
645
- if get_pp_group().is_last_rank:
646
- self.unpadded_vocab_size = config.vocab_size
647
- if lora_config:
648
- self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
649
- self.lm_head = ParallelLMHead(
650
- self.unpadded_vocab_size,
651
- config.hidden_size,
652
- org_num_embeddings=config.vocab_size,
653
- padding_size=DEFAULT_VOCAB_PADDING_SIZE,
654
- quant_config=quant_config,
655
- )
656
- if config.tie_word_embeddings:
657
- self.lm_head.weight = self.model.embed_tokens.weight
658
-
659
- logit_scale = getattr(config, "logit_scale", 1.0)
660
- self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
661
- config.vocab_size,
662
- logit_scale)
663
- else:
664
- self.lm_head = PPMissingLayer()
665
-
666
- def forward(
667
- self,
668
- input_ids: torch.Tensor,
669
- positions: torch.Tensor,
670
- intermediate_tensors: Optional[IntermediateTensors] = None,
671
- inputs_embeds: Optional[torch.Tensor] = None,
672
- ) -> Union[torch.Tensor, IntermediateTensors]:
673
- model_output = self.model(input_ids, positions, intermediate_tensors,
674
- inputs_embeds)
675
- return model_output
676
-
677
- def compute_logits(
678
- self,
679
- hidden_states: torch.Tensor,
680
- sampling_metadata: SamplingMetadata,
681
- ) -> Optional[torch.Tensor]:
682
- logits = self.logits_processor(self.lm_head, hidden_states,
683
- sampling_metadata)
684
- return logits
685
-
686
- def make_empty_intermediate_tensors(
687
- self, batch_size: int, dtype: torch.dtype,
688
- device: torch.device) -> IntermediateTensors:
689
- return IntermediateTensors({
690
- "hidden_states":
691
- torch.zeros((batch_size, self.config.hidden_size),
692
- dtype=dtype,
693
- device=device),
694
- "residual":
695
- torch.zeros((batch_size, self.config.hidden_size),
696
- dtype=dtype,
697
- device=device),
698
- })
699
-
622
def _split_qkv_weight(self, qkv: torch.Tensor):
623
num_attention_heads = self.config.num_attention_heads
624
num_kv_heads = getattr(self.config, "num_key_value_heads",
@@ -719,6 +641,17 @@ class HunYuanMoEV1ForCausalLM(nn.Module):
641
v = v.reshape(-1, hidden_size)
642
return torch.concat((q, k, v))
643
644
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
645
+
646
+ # Params for weights, fp8 weight scales, fp8 activation scales
647
+ # (param_name, weight_name, expert_id, shard_id)
648
+ return FusedMoE.make_expert_params_mapping(
649
+ ckpt_gate_proj_name="gate_proj",
650
+ ckpt_down_proj_name="down_proj",
651
+ ckpt_up_proj_name="up_proj",
652
+ num_experts=self.config.num_experts,
653
+ )
654
+
655
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
656
cla_factor = _get_cla_factor(self.config)
657
stacked_params_mapping = [
@@ -745,16 +678,9 @@ class HunYuanMoEV1ForCausalLM(nn.Module):
678
),
679
]
680
748
- # Params for weights, fp8 weight scales, fp8 activation scales
749
- # (param_name, weight_name, expert_id, shard_id)
750
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
751
- ckpt_gate_proj_name="gate_proj",
752
- ckpt_down_proj_name="down_proj",
753
- ckpt_up_proj_name="up_proj",
754
- num_experts=self.config.num_experts,
755
- )
756
-
681
params_dict = dict(self.named_parameters())
682
+ loaded_params: set[str] = set()
683
+ expert_params_mapping = self.get_expert_mapping()
684
for name, loaded_weight in weights:
685
if "rotary_emb.inv_freq" in name:
686
continue
@@ -806,7 +732,7 @@ class HunYuanMoEV1ForCausalLM(nn.Module):
732
param = params_dict[name]
733
weight_loader = param.weight_loader
734
weight_loader(param, loaded_weight, shard_id)
809
-
735
+ loaded_params.add(name)
736
is_found = True
737
break
738
if is_found:
@@ -885,3 +811,93 @@ class HunYuanMoEV1ForCausalLM(nn.Module):
811
weight_loader = getattr(param, "weight_loader",
812
default_weight_loader)
813
weight_loader(param, loaded_weight)
814
+ loaded_params.add(name)
815
+ return loaded_params
816
+
817
+
818
+class HunYuanMoEV1ForCausalLM(nn.Module, SupportsLoRA):
819
+ packed_modules_mapping = {
820
+ "qkv_proj": [
821
+ "q_proj",
822
+ "k_proj",
823
+ "v_proj",
824
+ ],
825
+ "gate_up_proj": [
826
+ "gate_proj",
827
+ "up_proj",
828
+ ],
829
+ }
830
+
831
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
832
+ super().__init__()
833
+
834
+ config = vllm_config.model_config.hf_config
835
+ quant_config = vllm_config.quant_config
836
+ self.config = config
837
+ self.quant_config = quant_config
838
+
839
+ self.model = HunYuanModel(vllm_config=vllm_config, prefix="model")
840
+ if get_pp_group().is_last_rank:
841
+ self.unpadded_vocab_size = config.vocab_size
842
+ self.lm_head = ParallelLMHead(
843
+ self.unpadded_vocab_size,
844
+ config.hidden_size,
845
+ org_num_embeddings=config.vocab_size,
846
+ padding_size=DEFAULT_VOCAB_PADDING_SIZE,
847
+ quant_config=quant_config,
848
+ )
849
+ if config.tie_word_embeddings:
850
+ self.lm_head.weight = self.model.embed_tokens.weight
851
+
852
+ logit_scale = getattr(config, "logit_scale", 1.0)
853
+ self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
854
+ config.vocab_size,
855
+ logit_scale)
856
+ else:
857
+ self.lm_head = PPMissingLayer()
858
+
859
+ def forward(
860
+ self,
861
+ input_ids: torch.Tensor,
862
+ positions: torch.Tensor,
863
+ intermediate_tensors: Optional[IntermediateTensors] = None,
864
+ inputs_embeds: Optional[torch.Tensor] = None,
865
+ ) -> Union[torch.Tensor, IntermediateTensors]:
866
+ model_output = self.model(input_ids, positions, intermediate_tensors,
867
+ inputs_embeds)
868
+ return model_output
869
+
870
+ def compute_logits(
871
+ self,
872
+ hidden_states: torch.Tensor,
873
+ sampling_metadata: SamplingMetadata,
874
+ ) -> Optional[torch.Tensor]:
875
+ logits = self.logits_processor(self.lm_head, hidden_states,
876
+ sampling_metadata)
877
+ return logits
878
+
879
+ def make_empty_intermediate_tensors(
880
+ self, batch_size: int, dtype: torch.dtype,
881
+ device: torch.device) -> IntermediateTensors:
882
+ return IntermediateTensors({
883
+ "hidden_states":
884
+ torch.zeros((batch_size, self.config.hidden_size),
885
+ dtype=dtype,
886
+ device=device),
887
+ "residual":
888
+ torch.zeros((batch_size, self.config.hidden_size),
889
+ dtype=dtype,
890
+ device=device),
891
+ })
892
+
893
+ def load_weights(self, weights: Iterable[tuple[str,
894
+ torch.Tensor]]) -> set[str]:
895
+ loader = AutoWeightsLoader(
896
+ self,
897
+ skip_prefixes=(["lm_head."]
898
+ if self.config.tie_word_embeddings else None),
899
+ )
900
+ return loader.load_weights(weights)
901
+
902
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
903
+ return self.model.get_expert_mapping()