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, 12:30 UTC
[Models] Cohere MoE (#40817)
Signed-off-by: Terrencezzj <terrence@cohere.ai>
Terrence Zhao committed
Apr 29, 2026 at 11:54 UTC
91a2d3901416fcff11e192f32683ca963726989b
7 files changed
+554
-3
docs/models/supported_models.md
+1
@@ -378,6 +378,7 @@ th {
378
| `BloomForCausalLM` | BLOOM, BLOOMZ, BLOOMChat | `bigscience/bloom`, `bigscience/bloomz`, etc. | | ✅︎ |
379
| `ChatGLMModel`, `ChatGLMForConditionalGeneration` | ChatGLM | `zai-org/chatglm2-6b`, `zai-org/chatglm3-6b`, `thu-coai/ShieldLM-6B-chatglm3`, etc. | ✅︎ | ✅︎ |
380
| `CohereForCausalLM`, `Cohere2ForCausalLM` | Command-R, Command-A | `CohereLabs/c4ai-command-r-v01`, `CohereLabs/c4ai-command-r7b-12-2024`, `CohereLabs/c4ai-command-a-03-2025`, `CohereLabs/command-a-reasoning-08-2025`, etc. | ✅︎ | ✅︎ |
381
+| `CohereMoeForCausalLM` | Command (MoE) | (model checkpoints loaded with `trust_remote_code=True`) | ✅︎ | ✅︎ |
382
| `CwmForCausalLM` | CWM | `facebook/cwm`, etc. | ✅︎ | ✅︎ |
383
| `DbrxForCausalLM` | DBRX | `databricks/dbrx-base`, `databricks/dbrx-instruct`, etc. | | ✅︎ |
384
| `DeciLMForCausalLM` | DeciLM | `nvidia/Llama-3_3-Nemotron-Super-49B-v1`, etc. | ✅︎ | ✅︎ |
tests/models/registry.py
+5
@@ -238,6 +238,11 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
238
"CohereLabs/c4ai-command-r7b-12-2024",
239
trust_remote_code=True,
240
),
241
+ "CohereMoeForCausalLM": _HfExamplesInfo(
242
+ "/host/engines/cohere-moe",
243
+ trust_remote_code=True,
244
+ is_available_online=False,
245
+ ),
246
"CwmForCausalLM": _HfExamplesInfo("facebook/cwm", min_transformers_version="4.58"),
247
# FIXME: databricks/dbrx-instruct has been deleted
248
"DbrxForCausalLM": _HfExamplesInfo(
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+1
@@ -263,6 +263,7 @@ class TrtLlmNvFp4ExpertsMonolithic(
263
RoutingMethodType.SigmoidRenorm,
264
RoutingMethodType.MiniMax2,
265
RoutingMethodType.Simulated,
266
+ RoutingMethodType.SigmoidRenorm,
267
]
268
269
@staticmethod
vllm/model_executor/layers/fused_moe/router/custom_routing_router.py
+4
@@ -34,11 +34,15 @@ class CustomRoutingRouter(BaseRouter):
34
35
@property
36
def routing_method_type(self) -> RoutingMethodType:
37
+ from vllm.model_executor.models.cohere_moe import token_choice_with_bias
38
from vllm.model_executor.models.llama4 import Llama4MoE
39
40
# NOTE: FLASHINFER_TRTLLM support the Llama4 router.
41
if self.custom_routing_function == Llama4MoE.custom_routing_function:
42
return RoutingMethodType.Llama4
43
+ # Cohere MoE uses a sigmoid -> top-k -> renormalize routing function.
44
+ if self.custom_routing_function == token_choice_with_bias:
45
+ return RoutingMethodType.SigmoidRenorm
46
return RoutingMethodType.Custom
47
48
def _compute_routing(
vllm/model_executor/models/cohere2_vision.py
+14
-3
@@ -44,7 +44,12 @@ from vllm.multimodal.processing import (
44
from vllm.sequence import IntermediateTensors
45
from vllm.utils.tensor_schema import TensorSchema, TensorShape
46
47
-from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
47
+from .interfaces import (
48
+ MultiModalEmbeddings,
49
+ SupportsMultiModal,
50
+ SupportsPP,
51
+ SupportsQuant,
52
+)
53
from .siglip import SiglipVisionModel
54
from .utils import (
55
AutoWeightsLoader,
@@ -309,16 +314,22 @@ class Cohere2VisionMultiModalProcessor(
314
info=Cohere2VisionProcessingInfo,
315
dummy_inputs=Cohere2VisionDummyInputsBuilder,
316
)
312
-class Cohere2VisionForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
317
+class Cohere2VisionForConditionalGeneration(
318
+ nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant
319
+):
320
hf_to_vllm_mapper = WeightsMapper(
321
orig_to_new_prefix={
322
"model.vision_tower.": "vision_tower.",
323
"model.multi_modal_projector.": "multi_modal_projector.",
324
"model.language_model.": "language_model.model.",
318
- "lm_head.": "language_model.lm_head.",
325
}
326
)
327
328
+ packed_modules_mapping = {
329
+ "qkv_proj": ["q_proj", "k_proj", "v_proj"],
330
+ "gate_up_proj": ["gate_proj", "up_proj"],
331
+ }
332
+
333
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
334
super().__init__()
335
config: Cohere2VisionConfig = vllm_config.model_config.hf_config
vllm/model_executor/models/cohere_moe.py
new
+528
@@ -0,0 +1,528 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from collections.abc import Iterable
5
+from itertools import islice
6
+
7
+import torch
8
+from torch import nn
9
+from transformers import CohereConfig
10
+
11
+from vllm.compilation.decorators import support_torch_compile
12
+from vllm.config import CacheConfig, VllmConfig
13
+from vllm.distributed import (
14
+ get_pp_group,
15
+ get_tensor_model_parallel_world_size,
16
+)
17
+from vllm.model_executor.layers.activation import SiluAndMul
18
+from vllm.model_executor.layers.attention import Attention
19
+from vllm.model_executor.layers.fused_moe import FusedMoE
20
+from vllm.model_executor.layers.linear import (
21
+ MergedColumnParallelLinear,
22
+ QKVParallelLinear,
23
+ ReplicatedLinear,
24
+ RowParallelLinear,
25
+)
26
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
27
+from vllm.model_executor.layers.quantization import QuantizationConfig
28
+from vllm.model_executor.layers.rotary_embedding import get_rope
29
+from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
30
+from vllm.model_executor.model_loader.weight_utils import (
31
+ default_weight_loader,
32
+ maybe_remap_kv_scale_name,
33
+)
34
+from vllm.platforms import current_platform
35
+from vllm.sequence import IntermediateTensors
36
+
37
+from .commandr import LayerNorm
38
+from .interfaces import SupportsPP, SupportsQuant
39
+from .utils import (
40
+ extract_layer_index,
41
+ is_pp_missing_parameter,
42
+ make_empty_intermediate_tensors_factory,
43
+ make_layers,
44
+ maybe_prefix,
45
+)
46
+
47
+
48
+@torch.compile(backend=current_platform.simple_compile_backend)
49
+def token_choice_with_bias(
50
+ hidden_states: torch.Tensor,
51
+ gating_output: torch.Tensor,
52
+ topk: int,
53
+ renormalize: bool,
54
+):
55
+ """Sigmoid -> top-k (-> renormalize) custom routing for CohereMoe."""
56
+ assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
57
+
58
+ scores = gating_output.float().sigmoid()
59
+ topk_weights, topk_ids = torch.topk(scores, k=topk, dim=-1, sorted=False)
60
+
61
+ if renormalize:
62
+ topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
63
+
64
+ return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
65
+
66
+
67
+class CohereMoeMLP(nn.Module):
68
+ """Cohere MLP used as shared experts in the MoE block."""
69
+
70
+ def __init__(
71
+ self,
72
+ config: CohereConfig,
73
+ intermediate_size: int | None = None,
74
+ quant_config: QuantizationConfig | None = None,
75
+ prefix: str = "",
76
+ ):
77
+ super().__init__()
78
+ self.config = config
79
+ self.hidden_size = config.hidden_size
80
+ self.intermediate_size = (
81
+ intermediate_size
82
+ if intermediate_size is not None
83
+ else config.intermediate_size
84
+ )
85
+ self.gate_up_proj = MergedColumnParallelLinear(
86
+ self.hidden_size,
87
+ [self.intermediate_size] * 2,
88
+ bias=False,
89
+ quant_config=quant_config,
90
+ prefix=f"{prefix}.gate_up_proj",
91
+ )
92
+ self.down_proj = RowParallelLinear(
93
+ self.intermediate_size,
94
+ self.hidden_size,
95
+ bias=False,
96
+ quant_config=quant_config,
97
+ reduce_results=False,
98
+ prefix=f"{prefix}.down_proj",
99
+ )
100
+ self.act_fn = SiluAndMul()
101
+
102
+ def forward(self, x):
103
+ gate_up, _ = self.gate_up_proj(x)
104
+ x = self.act_fn(gate_up)
105
+ x, _ = self.down_proj(x)
106
+ return x
107
+
108
+
109
+class CohereMoeAttention(nn.Module):
110
+ """Cohere MoE attention with sliding-window interleave."""
111
+
112
+ def __init__(
113
+ self,
114
+ config: CohereConfig,
115
+ cache_config: CacheConfig | None = None,
116
+ quant_config: QuantizationConfig | None = None,
117
+ prefix: str = "",
118
+ ):
119
+ super().__init__()
120
+ tp_size = get_tensor_model_parallel_world_size()
121
+ self.config = config
122
+ self.layer_idx = extract_layer_index(prefix)
123
+ self.hidden_size = config.hidden_size
124
+ self.total_num_heads = config.num_attention_heads
125
+ self.num_heads = self.total_num_heads // tp_size
126
+ self.head_dim = getattr(
127
+ config, "head_dim", self.hidden_size // self.total_num_heads
128
+ )
129
+ self.total_num_kv_heads = config.num_key_value_heads
130
+ if self.total_num_kv_heads >= tp_size:
131
+ assert self.total_num_kv_heads % tp_size == 0
132
+ else:
133
+ assert tp_size % self.total_num_kv_heads == 0
134
+ self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
135
+ self.q_size = self.num_heads * self.head_dim
136
+ self.kv_size = self.num_kv_heads * self.head_dim
137
+ self.scaling = self.head_dim**-0.5
138
+ self.max_position_embeddings = getattr(
139
+ config, "model_max_length", None
140
+ ) or getattr(config, "max_position_embeddings", 8192)
141
+ self.qkv_proj = QKVParallelLinear(
142
+ self.hidden_size,
143
+ self.head_dim,
144
+ self.total_num_heads,
145
+ self.total_num_kv_heads,
146
+ bias=False,
147
+ quant_config=quant_config,
148
+ prefix=f"{prefix}.qkv_proj",
149
+ )
150
+ self.o_proj = RowParallelLinear(
151
+ self.total_num_heads * self.head_dim,
152
+ self.hidden_size,
153
+ bias=False,
154
+ quant_config=quant_config,
155
+ prefix=f"{prefix}.o_proj",
156
+ )
157
+ self.rotary_emb = get_rope(
158
+ self.head_dim,
159
+ max_position=self.max_position_embeddings,
160
+ rope_parameters=config.rope_parameters,
161
+ is_neox_style=False,
162
+ )
163
+
164
+ self.sliding_window = None
165
+ layer_types = getattr(config, "layer_types", None)
166
+ if (
167
+ layer_types is not None
168
+ and layer_types[self.layer_idx] == "sliding_attention"
169
+ ):
170
+ self.sliding_window = config.sliding_window
171
+
172
+ self.attn = Attention(
173
+ self.num_heads,
174
+ self.head_dim,
175
+ self.scaling,
176
+ num_kv_heads=self.num_kv_heads,
177
+ cache_config=cache_config,
178
+ quant_config=quant_config,
179
+ per_layer_sliding_window=self.sliding_window,
180
+ prefix=f"{prefix}.attn",
181
+ )
182
+
183
+ def forward(
184
+ self,
185
+ positions: torch.Tensor,
186
+ hidden_states: torch.Tensor,
187
+ ) -> torch.Tensor:
188
+ qkv, _ = self.qkv_proj(hidden_states)
189
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
190
+ if self.sliding_window:
191
+ q, k = self.rotary_emb(positions, q, k)
192
+ attn_output = self.attn(q, k, v)
193
+ output, _ = self.o_proj(attn_output)
194
+ return output
195
+
196
+
197
+class CohereMoe(nn.Module):
198
+ """Tensor-parallel MoE block for CohereMoe with shared experts."""
199
+
200
+ def __init__(
201
+ self,
202
+ config: CohereConfig,
203
+ params_dtype: torch.dtype | None = None,
204
+ quant_config: QuantizationConfig | None = None,
205
+ tp_size: int | None = None,
206
+ prefix: str = "",
207
+ ):
208
+ super().__init__()
209
+ self.hidden_size = config.hidden_size
210
+ self.tp_size = get_tensor_model_parallel_world_size()
211
+
212
+ if self.tp_size > config.num_experts:
213
+ raise ValueError(
214
+ f"Tensor parallel size {self.tp_size} is greater than "
215
+ f"the number of experts {config.num_experts}."
216
+ )
217
+
218
+ if (
219
+ hasattr(config, "expert_selection_fn")
220
+ and config.expert_selection_fn == "sigmoid"
221
+ ):
222
+ self.custom_routing_function = token_choice_with_bias
223
+ else:
224
+ self.custom_routing_function = None
225
+
226
+ self.gate = ReplicatedLinear(
227
+ config.hidden_size,
228
+ config.num_experts,
229
+ bias=False,
230
+ params_dtype=params_dtype,
231
+ quant_config=None,
232
+ prefix=f"{prefix}.gate",
233
+ )
234
+
235
+ if hasattr(config, "num_shared_experts") and config.num_shared_experts > 0:
236
+ self.shared_experts = CohereMoeMLP(
237
+ config=config,
238
+ intermediate_size=config.intermediate_size * config.num_shared_experts,
239
+ quant_config=quant_config,
240
+ prefix=f"{prefix}.shared_experts",
241
+ )
242
+ self.shared_expert_combination_strategy = getattr(
243
+ config, "shared_expert_combination_strategy", "sum"
244
+ )
245
+ assert self.shared_expert_combination_strategy in ("average", "sum"), (
246
+ "shared_expert_combination_strategy must be one of ['average', 'sum']"
247
+ )
248
+ else:
249
+ self.shared_experts = None
250
+ self.shared_expert_combination_strategy = None
251
+
252
+ self.experts = FusedMoE(
253
+ num_experts=config.num_experts,
254
+ top_k=config.num_experts_per_tok,
255
+ hidden_size=config.hidden_size,
256
+ intermediate_size=config.intermediate_size,
257
+ params_dtype=params_dtype,
258
+ renormalize=getattr(config, "norm_topk_prob", True),
259
+ quant_config=quant_config,
260
+ tp_size=tp_size,
261
+ prefix=f"{prefix}.experts",
262
+ custom_routing_function=self.custom_routing_function,
263
+ shared_experts=self.shared_experts,
264
+ )
265
+
266
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
267
+ orig_shape = hidden_states.shape
268
+ hidden_states = hidden_states.view(-1, self.hidden_size)
269
+ router_logits, _ = self.gate(hidden_states)
270
+ # FusedMoE handles shared expert overlap internally and returns
271
+ # shared_output + routed_output when shared_experts is set.
272
+ final_hidden_states = self.experts(hidden_states, router_logits)
273
+ if self.shared_expert_combination_strategy == "average":
274
+ final_hidden_states = final_hidden_states / 2
275
+ return final_hidden_states.view(orig_shape)
276
+
277
+
278
+class CohereMoeDecoderLayer(nn.Module):
279
+ def __init__(
280
+ self,
281
+ config: CohereConfig,
282
+ cache_config: CacheConfig | None = None,
283
+ quant_config: QuantizationConfig | None = None,
284
+ prefix: str = "",
285
+ ):
286
+ super().__init__()
287
+ self.config = config
288
+ self.hidden_size = config.hidden_size
289
+
290
+ self.self_attn = CohereMoeAttention(
291
+ config,
292
+ cache_config,
293
+ quant_config=quant_config,
294
+ prefix=f"{prefix}.self_attn",
295
+ )
296
+ self.mlp = CohereMoe(
297
+ config=config, quant_config=quant_config, prefix=f"{prefix}.mlp"
298
+ )
299
+ self.input_layernorm = LayerNorm(
300
+ param_shape=(config.hidden_size,), eps=config.layer_norm_eps
301
+ )
302
+
303
+ def forward(
304
+ self,
305
+ positions: torch.Tensor,
306
+ hidden_states: torch.Tensor,
307
+ residual: torch.Tensor | None,
308
+ ) -> tuple[torch.Tensor, torch.Tensor]:
309
+ residual = hidden_states
310
+ hidden_states, residual = self.input_layernorm(hidden_states, residual)
311
+ hidden_states_attention = self.self_attn(
312
+ positions=positions,
313
+ hidden_states=hidden_states,
314
+ )
315
+ hidden_states_mlp = self.mlp(hidden_states)
316
+
317
+ hidden_states = residual + hidden_states_attention + hidden_states_mlp
318
+ return hidden_states, residual
319
+
320
+
321
+@support_torch_compile
322
+class CohereMoeModel(nn.Module):
323
+ """Transformer decoder for CohereMoe."""
324
+
325
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
326
+ super().__init__()
327
+
328
+ config = vllm_config.model_config.hf_config
329
+ cache_config = vllm_config.cache_config
330
+ quant_config = vllm_config.quant_config
331
+
332
+ self.config = config
333
+ self.vocab_size = config.vocab_size
334
+ self.org_vocab_size = config.vocab_size
335
+ self.embed_tokens = VocabParallelEmbedding(
336
+ config.vocab_size, config.hidden_size
337
+ )
338
+ self.start_layer, self.end_layer, self.layers = make_layers(
339
+ config.num_hidden_layers,
340
+ lambda prefix: CohereMoeDecoderLayer(
341
+ config, cache_config, quant_config, prefix=prefix
342
+ ),
343
+ prefix=f"{prefix}.layers",
344
+ )
345
+ self.norm = LayerNorm(
346
+ param_shape=(config.hidden_size,), eps=config.layer_norm_eps
347
+ )
348
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
349
+ ["hidden_states", "residual"], config.hidden_size
350
+ )
351
+
352
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
353
+ return self.embed_tokens(input_ids)
354
+
355
+ def forward(
356
+ self,
357
+ input_ids: torch.Tensor,
358
+ positions: torch.Tensor,
359
+ intermediate_tensors: IntermediateTensors | None = None,
360
+ inputs_embeds: torch.Tensor | None = None,
361
+ ) -> torch.Tensor | IntermediateTensors:
362
+ if get_pp_group().is_first_rank:
363
+ if inputs_embeds is not None:
364
+ hidden_states = inputs_embeds
365
+ else:
366
+ hidden_states = self.get_input_embeddings(input_ids)
367
+ residual = None
368
+ else:
369
+ assert intermediate_tensors is not None
370
+ hidden_states = intermediate_tensors["hidden_states"]
371
+ residual = intermediate_tensors["residual"]
372
+ for layer in islice(self.layers, self.start_layer, self.end_layer):
373
+ hidden_states, residual = layer(positions, hidden_states, residual)
374
+ if not get_pp_group().is_last_rank:
375
+ return IntermediateTensors(
376
+ {"hidden_states": hidden_states, "residual": residual}
377
+ )
378
+ hidden_states, _ = self.norm(hidden_states, residual)
379
+ return hidden_states
380
+
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
+
438
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
439
+ stacked_params_mapping = [
440
+ ("qkv_proj", "q_proj", "q"),
441
+ ("qkv_proj", "k_proj", "k"),
442
+ ("qkv_proj", "v_proj", "v"),
443
+ ("gate_up_proj", "gate_proj", 0),
444
+ ("gate_up_proj", "up_proj", 1),
445
+ ]
446
+
447
+ expert_params_mapping = FusedMoE.make_expert_params_mapping(
448
+ self,
449
+ ckpt_gate_proj_name="gate_proj",
450
+ ckpt_down_proj_name="down_proj",
451
+ ckpt_up_proj_name="up_proj",
452
+ num_experts=self.config.num_experts,
453
+ )
454
+
455
+ params_dict = dict(self.named_parameters())
456
+ loaded_params: set[str] = set()
457
+ for name, loaded_weight in weights:
458
+ if "rotary_emb.inv_freq" in name:
459
+ continue
460
+
461
+ if self.quant_config is not None and (
462
+ scale_name := self.quant_config.get_cache_scale(name)
463
+ ):
464
+ param = params_dict[scale_name]
465
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
466
+ loaded_weight = (
467
+ loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
468
+ )
469
+ weight_loader(param, loaded_weight)
470
+ loaded_params.add(scale_name)
471
+ continue
472
+
473
+ for param_name, shard_name, shard_id in stacked_params_mapping:
474
+ if shard_name not in name:
475
+ continue
476
+ if "mlp.experts" in name:
477
+ continue
478
+ name = name.replace(shard_name, param_name)
479
+ if name.endswith(".bias") and name not in params_dict:
480
+ continue
481
+ if is_pp_missing_parameter(name, self):
482
+ continue
483
+ param = params_dict[name]
484
+ weight_loader = param.weight_loader
485
+ weight_loader(param, loaded_weight, shard_id)
486
+ break
487
+ else:
488
+ for mapping in expert_params_mapping:
489
+ param_name, weight_name, expert_id, shard_id = mapping
490
+ if weight_name not in name:
491
+ continue
492
+ name = name.replace(weight_name, param_name)
493
+ if is_pp_missing_parameter(name, self):
494
+ continue
495
+ if (
496
+ name.endswith(".bias") or name.endswith("_bias")
497
+ ) and name not in params_dict:
498
+ continue
499
+ param = params_dict[name]
500
+ weight_loader = param.weight_loader
501
+ weight_loader(
502
+ param,
503
+ loaded_weight,
504
+ name,
505
+ shard_id=shard_id,
506
+ expert_id=expert_id,
507
+ )
508
+ break
509
+ else:
510
+ if "lm_head.weight" in name:
511
+ continue
512
+ if (
513
+ name.endswith(".bias") or name.endswith("_bias")
514
+ ) and name not in params_dict:
515
+ continue
516
+ if is_pp_missing_parameter(name, self):
517
+ continue
518
+ name = maybe_remap_kv_scale_name(name, params_dict)
519
+ if name is None:
520
+ continue
521
+ param = params_dict[name]
522
+ weight_loader = getattr(
523
+ param, "weight_loader", default_weight_loader
524
+ )
525
+ weight_loader(param, loaded_weight)
526
+ loaded_params.add(name)
527
+
528
+ return loaded_params
vllm/model_executor/models/registry.py
+1
@@ -89,6 +89,7 @@ _TEXT_GENERATION_MODELS = {
89
"ChatGLMForConditionalGeneration": ("chatglm", "ChatGLMForCausalLM"),
90
"CohereForCausalLM": ("commandr", "CohereForCausalLM"),
91
"Cohere2ForCausalLM": ("commandr", "CohereForCausalLM"),
92
+ "CohereMoeForCausalLM": ("cohere_moe", "CohereMoeForCausalLM"),
93
"CwmForCausalLM": ("llama", "LlamaForCausalLM"),
94
"DbrxForCausalLM": ("dbrx", "DbrxForCausalLM"),
95
"DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),