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 4, 07:30 UTC
[Model] Use AutoWeightsLoader for CohereMoe (#41690)
Signed-off-by: bittoby <218712309+bittoby@users.noreply.github.com>
BitToby committed
May 4, 2026 at 18:44 UTC
0c620d2e083a49ba40c2a5df318fa246d7e7a59b
1 file changed
+63
-59
vllm/model_executor/models/cohere_moe.py
+63
-59
@@ -37,6 +37,7 @@ from vllm.sequence import IntermediateTensors
37
from .commandr import LayerNorm
38
from .interfaces import SupportsPP, SupportsQuant
39
from .utils import (
40
+ AutoWeightsLoader,
41
extract_layer_index,
42
is_pp_missing_parameter,
43
make_empty_intermediate_tensors_factory,
@@ -330,6 +331,7 @@ class CohereMoeModel(nn.Module):
331
quant_config = vllm_config.quant_config
332
333
self.config = config
334
+ self.quant_config = quant_config
335
self.vocab_size = config.vocab_size
336
self.org_vocab_size = config.vocab_size
337
self.embed_tokens = VocabParallelEmbedding(
@@ -378,63 +380,6 @@ class CohereMoeModel(nn.Module):
380
hidden_states, _ = self.norm(hidden_states, residual)
381
return hidden_states
382
381
-
382
-class CohereMoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
383
- is_text_generation_model = True
384
-
385
- packed_modules_mapping = {
386
- "qkv_proj": [
387
- "q_proj",
388
- "k_proj",
389
- "v_proj",
390
- ],
391
- "gate_up_proj": [
392
- "gate_proj",
393
- "up_proj",
394
- ],
395
- }
396
-
397
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
398
- super().__init__()
399
- config = vllm_config.model_config.hf_config
400
- quant_config = vllm_config.quant_config
401
- self.config = config
402
- assert getattr(config, "tie_word_embeddings", True)
403
- self.unpadded_vocab_size = config.vocab_size
404
- self.quant_config = quant_config
405
- self.logits_scale = config.logit_scale
406
- self.logits_processor = LogitsProcessor(
407
- self.unpadded_vocab_size, config.vocab_size, scale=self.logits_scale
408
- )
409
- self.model = CohereMoeModel(
410
- vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
411
- )
412
- self.make_empty_intermediate_tensors = (
413
- self.model.make_empty_intermediate_tensors
414
- )
415
-
416
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
417
- return self.model.get_input_embeddings(input_ids)
418
-
419
- def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
420
- return self.model.get_input_embeddings(input_ids)
421
-
422
- @torch.no_grad()
423
- def forward(
424
- self,
425
- input_ids: torch.Tensor,
426
- positions: torch.Tensor,
427
- intermediate_tensors: IntermediateTensors | None = None,
428
- inputs_embeds: torch.Tensor | None = None,
429
- ) -> torch.Tensor | IntermediateTensors:
430
- return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
431
-
432
- def compute_logits(
433
- self,
434
- hidden_states: torch.Tensor,
435
- ) -> torch.Tensor | None:
436
- return self.logits_processor(self.model.embed_tokens, hidden_states)
437
-
383
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
384
stacked_params_mapping = [
385
("qkv_proj", "q_proj", "q"),
@@ -507,8 +452,6 @@ class CohereMoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
452
)
453
break
454
else:
510
- if "lm_head.weight" in name:
511
- continue
455
if (
456
name.endswith(".bias") or name.endswith("_bias")
457
) and name not in params_dict:
@@ -526,3 +469,64 @@ class CohereMoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
469
loaded_params.add(name)
470
471
return loaded_params
472
+
473
+
474
+class CohereMoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
475
+ is_text_generation_model = True
476
+
477
+ packed_modules_mapping = {
478
+ "qkv_proj": [
479
+ "q_proj",
480
+ "k_proj",
481
+ "v_proj",
482
+ ],
483
+ "gate_up_proj": [
484
+ "gate_proj",
485
+ "up_proj",
486
+ ],
487
+ }
488
+
489
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
490
+ super().__init__()
491
+ config = vllm_config.model_config.hf_config
492
+ quant_config = vllm_config.quant_config
493
+ self.config = config
494
+ assert getattr(config, "tie_word_embeddings", True)
495
+ self.unpadded_vocab_size = config.vocab_size
496
+ self.quant_config = quant_config
497
+ self.logits_scale = config.logit_scale
498
+ self.logits_processor = LogitsProcessor(
499
+ self.unpadded_vocab_size, config.vocab_size, scale=self.logits_scale
500
+ )
501
+ self.model = CohereMoeModel(
502
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
503
+ )
504
+ self.make_empty_intermediate_tensors = (
505
+ self.model.make_empty_intermediate_tensors
506
+ )
507
+
508
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
509
+ return self.model.get_input_embeddings(input_ids)
510
+
511
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
512
+ return self.model.get_input_embeddings(input_ids)
513
+
514
+ @torch.no_grad()
515
+ def forward(
516
+ self,
517
+ input_ids: torch.Tensor,
518
+ positions: torch.Tensor,
519
+ intermediate_tensors: IntermediateTensors | None = None,
520
+ inputs_embeds: torch.Tensor | None = None,
521
+ ) -> torch.Tensor | IntermediateTensors:
522
+ return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
523
+
524
+ def compute_logits(
525
+ self,
526
+ hidden_states: torch.Tensor,
527
+ ) -> torch.Tensor | None:
528
+ return self.logits_processor(self.model.embed_tokens, hidden_states)
529
+
530
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
531
+ loader = AutoWeightsLoader(self, skip_prefixes=["lm_head."])
532
+ return loader.load_weights(weights)