[Model] Migrate Glm, Arcee, CWM and Mellum to the Transformers modeling backend (#59679)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Harry Mellor committed
Oct 2, 2026 at 04:15 UTC
8faacd707f2f69b86ecd0bb1e7b5a705b473cb25
6 files changed
+8
-652
docs/models/supported_models.md
+4
-4
@@ -374,7 +374,6 @@ th {
374
| ------------ | ------ | ----------------- | -------------------- | ------------------------- |
375
| `AfmoeForCausalLM` | Afmoe | TBA | ✅︎ | ✅︎ |
376
| `ApertusForCausalLM` | Apertus | `swiss-ai/Apertus-8B-2509`, `swiss-ai/Apertus-70B-Instruct-2509`, etc. | ✅︎ | ✅︎ |
377
-| `ArceeForCausalLM` | Arcee (AFM) | `arcee-ai/AFM-4.5B-Base`, etc. | ✅︎ | ✅︎ |
377
| `AXK1ForCausalLM` | A.X-K1 | `skt/A.X-K1`, etc. | | ✅︎ |
378
| `BailingMoeForCausalLM` | Ling | `inclusionAI/Ling-lite-1.5`, `inclusionAI/Ling-plus`, etc. | ✅︎ | ✅︎ |
379
| `BailingMoeV2ForCausalLM` | Ling | `inclusionAI/Ling-mini-2.0`, etc. | ✅︎ | ✅︎ |
@@ -384,7 +383,6 @@ th {
383
| `ChatGLMModel`, `ChatGLMForConditionalGeneration` | ChatGLM | `zai-org/chatglm2-6b`, `zai-org/chatglm3-6b`, `thu-coai/ShieldLM-6B-chatglm3`, etc. | ✅︎ | ✅︎ |
384
| `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. | ✅︎ | ✅︎ |
385
| `Cohere2MoeForCausalLM` | North-Mini-Code | `CohereLabs/North-Mini-Code`, etc. | ✅︎ | ✅︎ |
387
-| `CwmForCausalLM` | CWM | `facebook/cwm`, etc. | ✅︎ | ✅︎ |
386
| `DbrxForCausalLM` | DBRX | `databricks/dbrx-base`, `databricks/dbrx-instruct`, etc. | | ✅︎ |
387
| `DeciLMForCausalLM` | DeciLM | `nvidia/Llama-3_3-Nemotron-Super-49B-v1`, etc. | ✅︎ | ✅︎ |
388
| `DeepseekForCausalLM` | DeepSeek | `deepseek-ai/deepseek-llm-67b-base`, `deepseek-ai/deepseek-llm-7b-chat`, etc. | ✅︎ | ✅︎ |
@@ -406,7 +404,6 @@ th {
404
| `Gemma3ForCausalLM` | Gemma 3 | `google/gemma-3-1b-it`, etc. | ✅︎ | ✅︎ |
405
| `Gemma3nForCausalLM` | Gemma 3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
406
| `Gemma4ForCausalLM` | Gemma 4 | `google/gemma-4-E2B-it`, etc. | ✅︎ | ✅︎ |
409
-| `GlmForCausalLM` | GLM-4 | `zai-org/glm-4-9b-chat-hf`, etc. | ✅︎ | ✅︎ |
407
| `Glm4ForCausalLM` | GLM-4-0414 | `zai-org/GLM-4-32B-0414`, etc. | ✅︎ | ✅︎ |
408
| `Glm4MoeForCausalLM` | GLM-4.5, GLM-4.6, GLM-4.7 | `zai-org/GLM-4.5`, etc. | ✅︎ | ✅︎ |
409
| `Glm4MoeLiteForCausalLM` | GLM-4.7-Flash | `zai-org/GLM-4.7-Flash`, etc. | ✅︎ | ✅︎ |
@@ -439,7 +436,6 @@ th {
436
| `LongcatFlashForCausalLM` | LongCat-Flash | `meituan-longcat/LongCat-Flash-Chat`, `meituan-longcat/LongCat-Flash-Chat-FP8` | ✅︎ | ✅︎ |
437
| `MambaForCausalLM` | Mamba | `state-spaces/mamba-130m-hf`, `state-spaces/mamba-790m-hf`, `state-spaces/mamba-2.8b-hf`, etc. | | ✅︎ |
438
| `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ |
442
-| `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | | ✅︎ |
439
| `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ |
440
| `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | | ✅︎ |
441
| `MiMoV2ForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |
@@ -486,10 +482,14 @@ Some models are supported only via the [Transformers modeling backend](#transfor
482
483
| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
484
| ------------ | ------ | ----------------- | -------------------- | ------------------------- |
485
+| `ArceeForCausalLM` | Arcee (AFM) | `arcee-ai/AFM-4.5B-Base`, etc. | ✅︎ | ✅︎ |
486
+| `CwmForCausalLM` | CWM | `facebook/cwm`, etc. | ✅︎ | ✅︎ |
487
| `FlexOlmoForCausalLM` | FlexOlmo | `allenai/FlexOlmo-7x7B-1T`, `allenai/FlexOlmo-7x7B-1T-RT`, etc. | | ✅︎ |
488
+| `GlmForCausalLM` | GLM-4 | `zai-org/glm-4-9b-chat-hf`, etc. | ✅︎ | ✅︎ |
489
| `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | |
490
| `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |
491
| `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |
492
+| `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | | ✅︎ |
493
| `NanbeigeForCausalLM` | Nanbeige4.2 | `Nanbeige/Nanbeige4.2-3B`, etc. | ✅︎ | ✅︎ |
494
| `OlmoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ |
495
| `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ |
vllm/model_executor/models/arcee.py
deleted
-365
@@ -1,365 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-# Copyright 2023-2025 vLLM Team
4
-# Licensed under the Apache License, Version 2.0 (the "License");
5
-# You may not use this file except in compliance with the License.
6
-# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
7
-#
8
-# Inference-only Arcee (AFM) model – adds support for ReLU^2 feed-forward
9
-# activation.
10
-
11
-from collections.abc import Iterable
12
-from itertools import islice
13
-from typing import Any
14
-
15
-import torch
16
-from torch import nn
17
-from transformers import LlamaConfig
18
-
19
-from vllm.compilation.decorators import support_torch_compile
20
-from vllm.distributed import get_pp_group
21
-from vllm.model_executor.layers.activation import ReLUSquaredActivation
22
-from vllm.model_executor.layers.layernorm import RMSNorm
23
-from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
24
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
25
-from vllm.model_executor.layers.vocab_parallel_embedding import (
26
- ParallelLMHead,
27
- VocabParallelEmbedding,
28
-)
29
-from vllm.sequence import IntermediateTensors
30
-
31
-from .interfaces import (
32
- EagleModelMixin,
33
- SupportsEagle,
34
- SupportsEagle3,
35
- SupportsLoRA,
36
- SupportsPP,
37
-)
38
-from .utils import (
39
- AutoWeightsLoader,
40
- PPMissingLayer,
41
- WeightsMapper,
42
- make_empty_intermediate_tensors_factory,
43
- make_layers,
44
- maybe_prefix,
45
-)
46
-
47
-
48
-class ArceeMLP(nn.Module):
49
- """Feed-forward layer for Arcee using ReLU^2 activation
50
- (no gating as in LLaMA)."""
51
-
52
- def __init__(
53
- self,
54
- hidden_size: int,
55
- intermediate_size: int,
56
- hidden_act: str,
57
- quant_config: Any | None = None,
58
- bias: bool = False,
59
- prefix: str = "",
60
- reduce_results: bool = True,
61
- ) -> None:
62
- super().__init__()
63
- # Single linear projection up to intermediate size
64
- # (no separate gate projection)
65
- self.up_proj = ColumnParallelLinear(
66
- input_size=hidden_size,
67
- output_size=intermediate_size,
68
- bias=bias,
69
- quant_config=quant_config,
70
- prefix=f"{prefix}.up_proj",
71
- )
72
- # Down projection back to hidden size
73
- self.down_proj = RowParallelLinear(
74
- input_size=intermediate_size,
75
- output_size=hidden_size,
76
- bias=bias,
77
- quant_config=quant_config,
78
- reduce_results=reduce_results,
79
- prefix=f"{prefix}.down_proj",
80
- )
81
- if hidden_act != "relu2":
82
- raise ValueError(
83
- f"Unsupported activation: {hidden_act}. "
84
- "Only 'relu2' is supported for AFM."
85
- )
86
- # Define ReLU^2 activation: (ReLU(x))^2 elementwise
87
- self.act_fn = ReLUSquaredActivation()
88
-
89
- def forward(self, x: torch.Tensor) -> torch.Tensor:
90
- x, _ = self.up_proj(x) # Project to intermediate size
91
- x = self.act_fn(x) # Apply ReLU^2 activation elementwise
92
- x, _ = self.down_proj(x) # Project back down to hidden size
93
- return x
94
-
95
-
96
-class ArceeDecoderLayer(nn.Module):
97
- """Transformer decoder block for Arcee, with self-attention and
98
- ReLU^2 MLP."""
99
-
100
- def __init__(
101
- self,
102
- config: LlamaConfig,
103
- cache_config: Any | None = None,
104
- quant_config: Any | None = None,
105
- prefix: str = "",
106
- ) -> None:
107
- super().__init__()
108
- self.hidden_size = config.hidden_size
109
- max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
110
- # Determine if attention bias is needed (some variants use bias terms)
111
- attention_bias = getattr(config, "attention_bias", False) or getattr(
112
- config, "bias", False
113
- )
114
- bias_o_proj = attention_bias
115
- if hasattr(config, "qkv_bias"):
116
- attention_bias = config.qkv_bias
117
-
118
- # Self-Attention (using LLaMA's attention structure)
119
- from vllm.model_executor.models.llama import (
120
- LlamaAttention, # import here to avoid circular import
121
- )
122
-
123
- self.self_attn = LlamaAttention(
124
- config=config,
125
- hidden_size=self.hidden_size,
126
- num_heads=config.num_attention_heads,
127
- num_kv_heads=getattr(
128
- config, "num_key_value_heads", config.num_attention_heads
129
- ),
130
- max_position_embeddings=max_position_embeddings,
131
- quant_config=quant_config,
132
- bias=attention_bias,
133
- bias_o_proj=bias_o_proj,
134
- cache_config=cache_config,
135
- prefix=f"{prefix}.self_attn",
136
- attn_type=getattr(
137
- config, "attn_type", "decoder"
138
- ), # assume decoder (causal) unless specified
139
- )
140
- # MLP with ReLU^2 activation
141
- self.mlp = ArceeMLP(
142
- hidden_size=self.hidden_size,
143
- intermediate_size=config.intermediate_size,
144
- hidden_act=config.hidden_act,
145
- quant_config=quant_config,
146
- bias=getattr(config, "mlp_bias", False),
147
- prefix=f"{prefix}.mlp",
148
- )
149
- # Layer normalization layers (RMSNorm as in LLaMA)
150
- self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
151
- self.post_attention_layernorm = RMSNorm(
152
- config.hidden_size, eps=config.rms_norm_eps
153
- )
154
-
155
- def forward(
156
- self,
157
- positions: torch.Tensor,
158
- hidden_states: torch.Tensor,
159
- residual: torch.Tensor | None,
160
- ) -> tuple[torch.Tensor, torch.Tensor]:
161
- # Self-Attention block
162
- if residual is None:
163
- residual = hidden_states
164
- hidden_states = self.input_layernorm(hidden_states)
165
- else:
166
- # Fused residual add + layernorm if supported
167
- hidden_states, residual = self.input_layernorm(hidden_states, residual)
168
- hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
169
- # Feed-forward block
170
- hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
171
- hidden_states = self.mlp(hidden_states)
172
- return hidden_states, residual
173
-
174
-
175
-@support_torch_compile
176
-class ArceeModel(nn.Module, EagleModelMixin):
177
- """The transformer model backbone for Arcee (embedding layer + stacked
178
- decoder blocks + final norm)."""
179
-
180
- def __init__(
181
- self,
182
- *,
183
- vllm_config,
184
- prefix: str = "",
185
- layer_type: type[nn.Module] = ArceeDecoderLayer,
186
- ) -> None:
187
- super().__init__()
188
- config: LlamaConfig = vllm_config.model_config.hf_config
189
- cache_config = vllm_config.cache_config
190
- quant_config = vllm_config.quant_config
191
- self.quant_config = quant_config
192
- self.config = config
193
- self.vocab_size = config.vocab_size
194
-
195
- # Word embeddings (parallelized if using pipeline parallel)
196
- if get_pp_group().is_first_rank or (
197
- config.tie_word_embeddings and get_pp_group().is_last_rank
198
- ):
199
- self.embed_tokens = VocabParallelEmbedding(
200
- self.vocab_size,
201
- config.hidden_size,
202
- quant_config=quant_config,
203
- )
204
- else:
205
- self.embed_tokens = PPMissingLayer() # placeholder on non-embedding ranks
206
-
207
- # Build decoder layers across pipeline ranks
208
- self.start_layer, self.end_layer, self.layers = make_layers(
209
- config.num_hidden_layers,
210
- lambda prefix: layer_type(
211
- config=config,
212
- cache_config=cache_config,
213
- quant_config=quant_config,
214
- prefix=prefix,
215
- ),
216
- prefix=f"{prefix}.layers",
217
- )
218
- # Final RMSNorm on the last pipeline stage
219
- if get_pp_group().is_last_rank:
220
- self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
221
- else:
222
- self.norm = PPMissingLayer()
223
-
224
- # Prepare factory for empty intermediate tensors
225
- # (for pipeline scheduling)
226
- self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
227
- ["hidden_states", "residual"], config.hidden_size
228
- )
229
-
230
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
231
- return self.embed_tokens(input_ids)
232
-
233
- def forward(
234
- self,
235
- input_ids: torch.Tensor | None,
236
- positions: torch.Tensor,
237
- intermediate_tensors: IntermediateTensors | None,
238
- inputs_embeds: torch.Tensor | None = None,
239
- ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
240
- # Embedding lookup (on first pipeline rank)
241
- if get_pp_group().is_first_rank:
242
- hidden_states = (
243
- inputs_embeds
244
- if inputs_embeds is not None
245
- else self.embed_input_ids(input_ids)
246
- )
247
- residual = None
248
- else:
249
- assert intermediate_tensors is not None, (
250
- "IntermediateTensors must be provided for non-first pipeline ranks"
251
- )
252
- hidden_states = intermediate_tensors["hidden_states"]
253
- residual = intermediate_tensors["residual"]
254
-
255
- aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
256
- for idx, layer in enumerate(
257
- islice(self.layers, self.start_layer, self.end_layer)
258
- ):
259
- hidden_states, residual = layer(positions, hidden_states, residual)
260
- self._maybe_add_hidden_state(
261
- aux_hidden_states, idx + 1, hidden_states, residual
262
- )
263
-
264
- if not get_pp_group().is_last_rank:
265
- # Send intermediate results to the next pipeline stage
266
- return IntermediateTensors(
267
- {"hidden_states": hidden_states, "residual": residual}
268
- )
269
- # On last rank: apply final layer norm
270
- hidden_states, _ = self.norm(hidden_states, residual)
271
- if len(aux_hidden_states) > 0:
272
- return hidden_states, aux_hidden_states
273
- return hidden_states
274
-
275
-
276
-class ArceeForCausalLM(
277
- nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3
278
-):
279
- """Arcee Model for causal language modeling, integrated with vLLM
280
- runtime."""
281
-
282
- hf_to_vllm_mapper = WeightsMapper(
283
- orig_to_new_stacked={
284
- # weight_name: (param_name, shard_id)
285
- ".q_proj": (".qkv_proj", "q"),
286
- ".k_proj": (".qkv_proj", "k"),
287
- ".v_proj": (".qkv_proj", "v"),
288
- },
289
- orig_to_new_substr={
290
- "gate_proj": None,
291
- },
292
- )
293
- # Map fused module names to their submodule components
294
- # (for quantization and LoRA)
295
- packed_modules_mapping = {
296
- "qkv_proj": ["q_proj", "k_proj", "v_proj"],
297
- }
298
-
299
- def __init__(self, *, vllm_config, prefix: str = "") -> None:
300
- super().__init__()
301
- config = vllm_config.model_config.hf_config
302
- self.config = config
303
-
304
- # Initialize the inner Transformer model (ArceeModel)
305
- self.model = ArceeModel(
306
- vllm_config=vllm_config,
307
- prefix=maybe_prefix(prefix, "model"),
308
- )
309
- # On the last pipeline stage, set up the LM head and logits processor
310
- if get_pp_group().is_last_rank:
311
- # Determine vocabulary size (including any LoRA extra tokens
312
- # for padded LM head)
313
-
314
- self.lm_head = ParallelLMHead(
315
- config.vocab_size,
316
- config.hidden_size,
317
- quant_config=vllm_config.quant_config,
318
- bias=getattr(config, "lm_head_bias", False),
319
- prefix=maybe_prefix(prefix, "lm_head"),
320
- )
321
- if config.tie_word_embeddings:
322
- # Tie output weights with input embedding matrix
323
- self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
324
- logit_scale = getattr(config, "logit_scale", 1.0)
325
- self.logits_processor = LogitsProcessor(
326
- config.vocab_size, scale=logit_scale
327
- )
328
- else:
329
- # Placeholder for lm_head on non-last ranks
330
- self.lm_head = PPMissingLayer()
331
-
332
- self.make_empty_intermediate_tensors = (
333
- self.model.make_empty_intermediate_tensors
334
- )
335
-
336
- def forward(
337
- self,
338
- input_ids: torch.Tensor | None,
339
- positions: torch.Tensor,
340
- intermediate_tensors: IntermediateTensors | None = None,
341
- inputs_embeds: torch.Tensor | None = None,
342
- ) -> torch.Tensor | IntermediateTensors:
343
- model_output = self.model(
344
- input_ids=input_ids,
345
- positions=positions,
346
- intermediate_tensors=intermediate_tensors,
347
- inputs_embeds=inputs_embeds,
348
- )
349
- return model_output
350
-
351
- def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
352
- # Compute final logits from hidden states (last pipeline rank only)
353
- logits = self.logits_processor(self.lm_head, hidden_states)
354
- return logits
355
-
356
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
357
- return self.model.embed_input_ids(input_ids)
358
-
359
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
360
- """Load weights into the model (delegates to inner model and handles
361
- tied embeddings)."""
362
- loader = AutoWeightsLoader(self)
363
- # AutoWeightLoader handles weight name remapping, including fusing
364
- # separate q_proj, k_proj, v_proj into qkv_proj
365
- return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/bert_with_rope.py
-1
@@ -572,7 +572,6 @@ class GteNewModel(BertWithRope):
572
super().__init__(vllm_config=vllm_config, prefix=prefix, **kwargs)
573
574
# GteNewModel only gate_up_proj does not have bias.
575
- # Hack method learned from vllm/model_executor/models/glm.py
575
for layer in self.encoder.layers:
576
layer.mlp.gate_up_proj.bias = None
577
layer.mlp.gate_up_proj.skip_bias_add = True
vllm/model_executor/models/glm.py
deleted
-24
@@ -1,24 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Inference-only HF format GLM-4 model compatible with THUDM weights."""
4
-
5
-from vllm.config import VllmConfig
6
-from vllm.model_executor.models.llama import LlamaForCausalLM
7
-
8
-from .utils import PPMissingLayer
9
-
10
-
11
-class GlmForCausalLM(LlamaForCausalLM):
12
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
13
- hf_config = vllm_config.model_config.hf_config
14
- hf_config.rope_parameters["partial_rotary_factor"] = 0.5
15
- super().__init__(vllm_config=vllm_config, prefix=prefix)
16
- # Hack Llama model to fit HF format GLM implementation
17
- # Attention difference between GLM and Llama:
18
- # 1. Half partial rotary_dim and no Neox style.
19
- # 2. There is no bias for o_proj in attention
20
- for layer in self.model.layers:
21
- if not isinstance(layer, PPMissingLayer):
22
- layer.self_attn.rotary_emb.is_neox_style = False
23
- layer.self_attn.o_proj.bias = None
24
- layer.self_attn.o_proj.skip_bias_add = True
vllm/model_executor/models/mellum.py
deleted
-253
@@ -1,253 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from typing import Any
5
-
6
-from torch import nn
7
-
8
-from vllm.compilation.decorators import support_torch_compile
9
-from vllm.config import VllmConfig
10
-from vllm.distributed import get_tensor_model_parallel_world_size
11
-from vllm.model_executor.layers.attention import Attention
12
-from vllm.model_executor.layers.layernorm import RMSNorm
13
-from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
14
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
15
-from vllm.model_executor.layers.rotary_embedding import get_rope
16
-from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
17
-
18
-from .qwen3_moe import (
19
- Qwen3MoeAttention,
20
- Qwen3MoeDecoderLayer,
21
- Qwen3MoeForCausalLM,
22
- Qwen3MoeMLP,
23
- Qwen3MoeModel,
24
- Qwen3MoeSparseMoeBlock,
25
-)
26
-from .utils import PPMissingLayer, extract_layer_index, maybe_prefix
27
-
28
-
29
-class MellumAttention(Qwen3MoeAttention):
30
- """Differences from `Qwen3MoeAttention`:
31
- - Supports `per_layer_sliding_window` for `Attention`.
32
- """
33
-
34
- def __init__(
35
- self,
36
- hidden_size: int,
37
- num_heads: int,
38
- num_kv_heads: int,
39
- rope_parameters: dict[str, Any],
40
- max_position_embeddings: int = 8192,
41
- head_dim: int | None = None,
42
- rms_norm_eps: float = 1e-06,
43
- qkv_bias: bool = False,
44
- cache_config: Any | None = None,
45
- quant_config: Any | None = None,
46
- prefix: str = "",
47
- dual_chunk_attention_config: dict[str, Any] | None = None,
48
- per_layer_sliding_window: int | None = None,
49
- ) -> None:
50
- nn.Module.__init__(self)
51
-
52
- self.hidden_size = hidden_size
53
- tp_size = get_tensor_model_parallel_world_size()
54
- self.total_num_heads = num_heads
55
- assert self.total_num_heads % tp_size == 0
56
- self.num_heads = self.total_num_heads // tp_size
57
- self.total_num_kv_heads = num_kv_heads
58
- if self.total_num_kv_heads >= tp_size:
59
- assert self.total_num_kv_heads % tp_size == 0
60
- else:
61
- assert tp_size % self.total_num_kv_heads == 0
62
- self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
63
- self.head_dim = head_dim or (hidden_size // self.total_num_heads)
64
- self.q_size = self.num_heads * self.head_dim
65
- self.kv_size = self.num_kv_heads * self.head_dim
66
- self.scaling = self.head_dim**-0.5
67
- self.max_position_embeddings = max_position_embeddings
68
- self.dual_chunk_attention_config = dual_chunk_attention_config
69
-
70
- self.qkv_proj = QKVParallelLinear(
71
- hidden_size,
72
- self.head_dim,
73
- self.total_num_heads,
74
- self.total_num_kv_heads,
75
- bias=qkv_bias,
76
- quant_config=quant_config,
77
- prefix=f"{prefix}.qkv_proj",
78
- )
79
-
80
- self.o_proj = RowParallelLinear(
81
- self.total_num_heads * self.head_dim,
82
- hidden_size,
83
- bias=False,
84
- quant_config=quant_config,
85
- prefix=f"{prefix}.o_proj",
86
- )
87
-
88
- self.rotary_emb = get_rope(
89
- self.head_dim,
90
- max_position=max_position_embeddings,
91
- rope_parameters=rope_parameters,
92
- dual_chunk_attention_config=dual_chunk_attention_config,
93
- )
94
- if dual_chunk_attention_config:
95
- self.attn = Attention(
96
- self.num_heads,
97
- self.head_dim,
98
- self.scaling,
99
- num_kv_heads=self.num_kv_heads,
100
- cache_config=cache_config,
101
- quant_config=quant_config,
102
- per_layer_sliding_window=per_layer_sliding_window,
103
- prefix=f"{prefix}.attn",
104
- layer_idx=extract_layer_index(prefix),
105
- dual_chunk_attention_config=dual_chunk_attention_config,
106
- )
107
- else:
108
- self.attn = Attention(
109
- self.num_heads,
110
- self.head_dim,
111
- self.scaling,
112
- num_kv_heads=self.num_kv_heads,
113
- cache_config=cache_config,
114
- quant_config=quant_config,
115
- per_layer_sliding_window=per_layer_sliding_window,
116
- prefix=f"{prefix}.attn",
117
- )
118
-
119
- self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
120
- self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
121
-
122
-
123
-class MellumDecoderLayer(Qwen3MoeDecoderLayer):
124
- """Differences from `Qwen3MoeDecoderLayer`:
125
- - Supports interleaved SWA and per-layer RoPE scaling.
126
- """
127
-
128
- def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
129
- nn.Module.__init__(self)
130
-
131
- config = vllm_config.model_config.hf_text_config
132
- cache_config = vllm_config.cache_config
133
- quant_config = vllm_config.quant_config
134
-
135
- self.hidden_size = config.hidden_size
136
- max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
137
- dual_chunk_attention_config = getattr(
138
- config, "dual_chunk_attention_config", None
139
- )
140
-
141
- layer_idx = extract_layer_index(prefix)
142
- layer_type = config.layer_types[layer_idx]
143
- if layer_type == "sliding_attention":
144
- sliding_window = getattr(config, "sliding_window", None)
145
- else:
146
- sliding_window = None
147
- rope_parameters = config.rope_parameters[layer_type]
148
-
149
- self.self_attn = MellumAttention(
150
- hidden_size=self.hidden_size,
151
- num_heads=config.num_attention_heads,
152
- num_kv_heads=config.num_key_value_heads,
153
- rope_parameters=rope_parameters,
154
- max_position_embeddings=max_position_embeddings,
155
- rms_norm_eps=config.rms_norm_eps,
156
- qkv_bias=getattr(config, "attention_bias", False),
157
- head_dim=getattr(config, "head_dim", None),
158
- cache_config=cache_config,
159
- quant_config=quant_config,
160
- prefix=f"{prefix}.self_attn",
161
- dual_chunk_attention_config=dual_chunk_attention_config,
162
- per_layer_sliding_window=sliding_window,
163
- )
164
-
165
- if config.mlp_layer_types[layer_idx] == "sparse":
166
- self.mlp = Qwen3MoeSparseMoeBlock(
167
- vllm_config=vllm_config, prefix=f"{prefix}.mlp"
168
- )
169
- else:
170
- self.mlp = Qwen3MoeMLP(
171
- hidden_size=config.hidden_size,
172
- intermediate_size=config.intermediate_size,
173
- hidden_act=config.hidden_act,
174
- quant_config=quant_config,
175
- prefix=f"{prefix}.mlp",
176
- )
177
-
178
- self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
179
- self.post_attention_layernorm = RMSNorm(
180
- config.hidden_size, eps=config.rms_norm_eps
181
- )
182
-
183
-
184
-@support_torch_compile
185
-class MellumModel(Qwen3MoeModel):
186
- """Differences from `Qwen3MoeModel`:
187
- - Uses `MellumDecoderLayer`.
188
- """
189
-
190
- def __init__(
191
- self,
192
- *,
193
- vllm_config: VllmConfig,
194
- prefix: str = "",
195
- ):
196
- super().__init__(
197
- vllm_config=vllm_config,
198
- prefix=prefix,
199
- decoder_layer_type=MellumDecoderLayer,
200
- )
201
-
202
-
203
-class MellumForCausalLM(Qwen3MoeForCausalLM):
204
- """Differences from `Qwen3MoeForCausalLM`:
205
- - Uses `MellumModel`.
206
- """
207
-
208
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
209
- nn.Module.__init__(self)
210
- config = vllm_config.model_config.hf_text_config
211
- quant_config = vllm_config.quant_config
212
- self.config = config
213
- self.quant_config = quant_config
214
- if "dense" in getattr(config, "mlp_layer_types", []):
215
- self.packed_modules_mapping["gate_up_proj"] = ["gate_proj", "up_proj"]
216
- self.model = MellumModel(
217
- vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
218
- )
219
- self.lm_head = ParallelLMHead(
220
- config.vocab_size,
221
- config.hidden_size,
222
- quant_config=quant_config,
223
- prefix=maybe_prefix(prefix, "lm_head"),
224
- )
225
- if self.config.tie_word_embeddings:
226
- self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
227
- self.logits_processor = LogitsProcessor(config.vocab_size)
228
- self.make_empty_intermediate_tensors = (
229
- self.model.make_empty_intermediate_tensors
230
- )
231
-
232
- self.moe_layers = []
233
- example_layer = None
234
- for layer in self.model.layers:
235
- if isinstance(layer, PPMissingLayer):
236
- continue
237
-
238
- assert isinstance(layer, Qwen3MoeDecoderLayer)
239
- if isinstance(layer.mlp, Qwen3MoeSparseMoeBlock):
240
- example_layer = layer.mlp
241
- self.moe_layers.append(layer.mlp.experts)
242
-
243
- if example_layer is None:
244
- raise RuntimeError("No MoE layer found in the model.layers.")
245
-
246
- self.num_moe_layers = len(self.moe_layers)
247
- self.num_expert_groups = 1
248
- self.num_shared_experts = 0
249
- self.num_logical_experts = example_layer.n_logical_experts
250
- self.num_physical_experts = example_layer.n_physical_experts
251
- self.num_local_physical_experts = example_layer.n_local_physical_experts
252
- self.num_routed_experts = example_layer.n_routed_experts
253
- self.num_redundant_experts = example_layer.n_redundant_experts
vllm/model_executor/models/registry.py
+4
-5
@@ -72,7 +72,6 @@ _TEXT_GENERATION_MODELS = {
72
# [Decoder-only]
73
"AfmoeForCausalLM": ("afmoe", "AfmoeForCausalLM"),
74
"ApertusForCausalLM": ("apertus", "ApertusForCausalLM"),
75
- "ArceeForCausalLM": ("arcee", "ArceeForCausalLM"),
75
"AXK1ForCausalLM": ("AXK1", "AXK1ForCausalLM"),
76
"BailingMoeForCausalLM": ("bailing_moe", "BailingMoeForCausalLM"),
77
"BailingMoeV2ForCausalLM": ("bailing_moe", "BailingMoeV2ForCausalLM"),
@@ -84,7 +83,6 @@ _TEXT_GENERATION_MODELS = {
83
"CohereForCausalLM": ("commandr", "CohereForCausalLM"),
84
"Cohere2ForCausalLM": ("commandr", "CohereForCausalLM"),
85
"Cohere2MoeForCausalLM": ("cohere2_moe", "Cohere2MoeForCausalLM"),
87
- "CwmForCausalLM": ("llama", "LlamaForCausalLM"),
86
"DbrxForCausalLM": ("dbrx", "DbrxForCausalLM"),
87
"DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
88
"DeepseekForCausalLM": ("deepseek_v2", "DeepseekForCausalLM"),
@@ -114,7 +112,6 @@ _TEXT_GENERATION_MODELS = {
112
"vllm.models.qwen4_exp",
113
"Qwen4ExpForCausalLM",
114
),
117
- "GlmForCausalLM": ("glm", "GlmForCausalLM"),
115
"Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),
116
"Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"),
117
"Glm4MoeLiteForCausalLM": ("glm4_moe_lite", "Glm4MoeLiteForCausalLM"),
@@ -160,7 +157,6 @@ _TEXT_GENERATION_MODELS = {
157
),
158
"MambaForCausalLM": ("mamba", "MambaForCausalLM"),
159
"Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"),
163
- "MellumForCausalLM": ("mellum", "MellumForCausalLM"),
160
"MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"),
161
"MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"),
162
"MiniMaxM2ForCausalLM": ("minimax_m2", "MiniMaxM2ForCausalLM"),
@@ -226,7 +222,6 @@ _EMBEDDING_MODELS = {
222
"DeepseekV3BidirectionalModel": ("deepseek_v2", "DeepseekV3ForCausalLM"),
223
"Gemma2Model": ("gemma2", "Gemma2ForCausalLM"),
224
"Gemma3TextModel": ("gemma3", "Gemma3Model"),
229
- "GlmForCausalLM": ("glm", "GlmForCausalLM"),
225
"GteModel": ("bert_with_rope", "SnowflakeGteNewModel"),
226
"GteNewModel": ("bert_with_rope", "GteNewModel"),
227
"JinaEmbeddingsV5Model": ("jina", "JinaEmbeddingsV5Model"),
@@ -709,10 +704,14 @@ _SPECULATIVE_DECODING_MODELS = {
704
705
_TRANSFORMERS_SUPPORTED_MODELS = {
706
# Text generation models
707
+ "ArceeForCausalLM": ("transformers", "TransformersForCausalLM"),
708
+ "CwmForCausalLM": ("transformers", "TransformersForCausalLM"),
709
"FlexOlmoForCausalLM": ("transformers", "TransformersMoEForCausalLM"),
710
+ "GlmForCausalLM": ("transformers", "TransformersForCausalLM"),
711
"GPTBigCodeForCausalLM": ("transformers", "TransformersForCausalLM"),
712
"HunYuanDenseV1ForCausalLM": ("transformers", "TransformersForCausalLM"),
713
"HunYuanMoEV1ForCausalLM": ("transformers", "TransformersMoEForCausalLM"),
714
+ "MellumForCausalLM": ("transformers", "TransformersMoEForCausalLM"),
715
"NanbeigeForCausalLM": ("transformers", "TransformersForCausalLM"),
716
"OlmoForCausalLM": ("transformers", "TransformersForCausalLM"),
717
"Olmo2ForCausalLM": ("transformers", "TransformersForCausalLM"),