From 8faacd707f2f69b86ecd0bb1e7b5a705b473cb25 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Fri, 2 Oct 2026 04:15:14 +0100 Subject: [PATCH] [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 --- docs/models/supported_models.md | 8 +- vllm/model_executor/models/arcee.py | 365 ------------------- vllm/model_executor/models/bert_with_rope.py | 1 - vllm/model_executor/models/glm.py | 24 -- vllm/model_executor/models/mellum.py | 253 ------------- vllm/model_executor/models/registry.py | 9 +- 6 files changed, 8 insertions(+), 652 deletions(-) delete mode 100644 vllm/model_executor/models/arcee.py delete mode 100644 vllm/model_executor/models/glm.py delete mode 100644 vllm/model_executor/models/mellum.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index df44a3e8b0..545890e152 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -374,7 +374,6 @@ th { | ------------ | ------ | ----------------- | -------------------- | ------------------------- | | `AfmoeForCausalLM` | Afmoe | TBA | ✅︎ | ✅︎ | | `ApertusForCausalLM` | Apertus | `swiss-ai/Apertus-8B-2509`, `swiss-ai/Apertus-70B-Instruct-2509`, etc. | ✅︎ | ✅︎ | -| `ArceeForCausalLM` | Arcee (AFM) | `arcee-ai/AFM-4.5B-Base`, etc. | ✅︎ | ✅︎ | | `AXK1ForCausalLM` | A.X-K1 | `skt/A.X-K1`, etc. | | ✅︎ | | `BailingMoeForCausalLM` | Ling | `inclusionAI/Ling-lite-1.5`, `inclusionAI/Ling-plus`, etc. | ✅︎ | ✅︎ | | `BailingMoeV2ForCausalLM` | Ling | `inclusionAI/Ling-mini-2.0`, etc. | ✅︎ | ✅︎ | @@ -384,7 +383,6 @@ th { | `ChatGLMModel`, `ChatGLMForConditionalGeneration` | ChatGLM | `zai-org/chatglm2-6b`, `zai-org/chatglm3-6b`, `thu-coai/ShieldLM-6B-chatglm3`, etc. | ✅︎ | ✅︎ | | `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. | ✅︎ | ✅︎ | | `Cohere2MoeForCausalLM` | North-Mini-Code | `CohereLabs/North-Mini-Code`, etc. | ✅︎ | ✅︎ | -| `CwmForCausalLM` | CWM | `facebook/cwm`, etc. | ✅︎ | ✅︎ | | `DbrxForCausalLM` | DBRX | `databricks/dbrx-base`, `databricks/dbrx-instruct`, etc. | | ✅︎ | | `DeciLMForCausalLM` | DeciLM | `nvidia/Llama-3_3-Nemotron-Super-49B-v1`, etc. | ✅︎ | ✅︎ | | `DeepseekForCausalLM` | DeepSeek | `deepseek-ai/deepseek-llm-67b-base`, `deepseek-ai/deepseek-llm-7b-chat`, etc. | ✅︎ | ✅︎ | @@ -406,7 +404,6 @@ th { | `Gemma3ForCausalLM` | Gemma 3 | `google/gemma-3-1b-it`, etc. | ✅︎ | ✅︎ | | `Gemma3nForCausalLM` | Gemma 3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | | | `Gemma4ForCausalLM` | Gemma 4 | `google/gemma-4-E2B-it`, etc. | ✅︎ | ✅︎ | -| `GlmForCausalLM` | GLM-4 | `zai-org/glm-4-9b-chat-hf`, etc. | ✅︎ | ✅︎ | | `Glm4ForCausalLM` | GLM-4-0414 | `zai-org/GLM-4-32B-0414`, etc. | ✅︎ | ✅︎ | | `Glm4MoeForCausalLM` | GLM-4.5, GLM-4.6, GLM-4.7 | `zai-org/GLM-4.5`, etc. | ✅︎ | ✅︎ | | `Glm4MoeLiteForCausalLM` | GLM-4.7-Flash | `zai-org/GLM-4.7-Flash`, etc. | ✅︎ | ✅︎ | @@ -439,7 +436,6 @@ th { | `LongcatFlashForCausalLM` | LongCat-Flash | `meituan-longcat/LongCat-Flash-Chat`, `meituan-longcat/LongCat-Flash-Chat-FP8` | ✅︎ | ✅︎ | | `MambaForCausalLM` | Mamba | `state-spaces/mamba-130m-hf`, `state-spaces/mamba-790m-hf`, `state-spaces/mamba-2.8b-hf`, etc. | | ✅︎ | | `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ | -| `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | | ✅︎ | | `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ | | `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | | ✅︎ | | `MiMoV2ForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ | @@ -486,10 +482,14 @@ Some models are supported only via the [Transformers modeling backend](#transfor | Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) | | ------------ | ------ | ----------------- | -------------------- | ------------------------- | +| `ArceeForCausalLM` | Arcee (AFM) | `arcee-ai/AFM-4.5B-Base`, etc. | ✅︎ | ✅︎ | +| `CwmForCausalLM` | CWM | `facebook/cwm`, etc. | ✅︎ | ✅︎ | | `FlexOlmoForCausalLM` | FlexOlmo | `allenai/FlexOlmo-7x7B-1T`, `allenai/FlexOlmo-7x7B-1T-RT`, etc. | | ✅︎ | +| `GlmForCausalLM` | GLM-4 | `zai-org/glm-4-9b-chat-hf`, etc. | ✅︎ | ✅︎ | | `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | | | `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ | | `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ | +| `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | | ✅︎ | | `NanbeigeForCausalLM` | Nanbeige4.2 | `Nanbeige/Nanbeige4.2-3B`, etc. | ✅︎ | ✅︎ | | `OlmoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ | | `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ | diff --git a/vllm/model_executor/models/arcee.py b/vllm/model_executor/models/arcee.py deleted file mode 100644 index 32ad19a8e7..0000000000 --- a/vllm/model_executor/models/arcee.py +++ /dev/null @@ -1,365 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Copyright 2023-2025 vLLM Team -# Licensed under the Apache License, Version 2.0 (the "License"); -# You may not use this file except in compliance with the License. -# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 -# -# Inference-only Arcee (AFM) model – adds support for ReLU^2 feed-forward -# activation. - -from collections.abc import Iterable -from itertools import islice -from typing import Any - -import torch -from torch import nn -from transformers import LlamaConfig - -from vllm.compilation.decorators import support_torch_compile -from vllm.distributed import get_pp_group -from vllm.model_executor.layers.activation import ReLUSquaredActivation -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.vocab_parallel_embedding import ( - ParallelLMHead, - VocabParallelEmbedding, -) -from vllm.sequence import IntermediateTensors - -from .interfaces import ( - EagleModelMixin, - SupportsEagle, - SupportsEagle3, - SupportsLoRA, - SupportsPP, -) -from .utils import ( - AutoWeightsLoader, - PPMissingLayer, - WeightsMapper, - make_empty_intermediate_tensors_factory, - make_layers, - maybe_prefix, -) - - -class ArceeMLP(nn.Module): - """Feed-forward layer for Arcee using ReLU^2 activation - (no gating as in LLaMA).""" - - def __init__( - self, - hidden_size: int, - intermediate_size: int, - hidden_act: str, - quant_config: Any | None = None, - bias: bool = False, - prefix: str = "", - reduce_results: bool = True, - ) -> None: - super().__init__() - # Single linear projection up to intermediate size - # (no separate gate projection) - self.up_proj = ColumnParallelLinear( - input_size=hidden_size, - output_size=intermediate_size, - bias=bias, - quant_config=quant_config, - prefix=f"{prefix}.up_proj", - ) - # Down projection back to hidden size - self.down_proj = RowParallelLinear( - input_size=intermediate_size, - output_size=hidden_size, - bias=bias, - quant_config=quant_config, - reduce_results=reduce_results, - prefix=f"{prefix}.down_proj", - ) - if hidden_act != "relu2": - raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only 'relu2' is supported for AFM." - ) - # Define ReLU^2 activation: (ReLU(x))^2 elementwise - self.act_fn = ReLUSquaredActivation() - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x, _ = self.up_proj(x) # Project to intermediate size - x = self.act_fn(x) # Apply ReLU^2 activation elementwise - x, _ = self.down_proj(x) # Project back down to hidden size - return x - - -class ArceeDecoderLayer(nn.Module): - """Transformer decoder block for Arcee, with self-attention and - ReLU^2 MLP.""" - - def __init__( - self, - config: LlamaConfig, - cache_config: Any | None = None, - quant_config: Any | None = None, - prefix: str = "", - ) -> None: - super().__init__() - self.hidden_size = config.hidden_size - max_position_embeddings = getattr(config, "max_position_embeddings", 8192) - # Determine if attention bias is needed (some variants use bias terms) - attention_bias = getattr(config, "attention_bias", False) or getattr( - config, "bias", False - ) - bias_o_proj = attention_bias - if hasattr(config, "qkv_bias"): - attention_bias = config.qkv_bias - - # Self-Attention (using LLaMA's attention structure) - from vllm.model_executor.models.llama import ( - LlamaAttention, # import here to avoid circular import - ) - - self.self_attn = LlamaAttention( - config=config, - hidden_size=self.hidden_size, - num_heads=config.num_attention_heads, - num_kv_heads=getattr( - config, "num_key_value_heads", config.num_attention_heads - ), - max_position_embeddings=max_position_embeddings, - quant_config=quant_config, - bias=attention_bias, - bias_o_proj=bias_o_proj, - cache_config=cache_config, - prefix=f"{prefix}.self_attn", - attn_type=getattr( - config, "attn_type", "decoder" - ), # assume decoder (causal) unless specified - ) - # MLP with ReLU^2 activation - self.mlp = ArceeMLP( - hidden_size=self.hidden_size, - intermediate_size=config.intermediate_size, - hidden_act=config.hidden_act, - quant_config=quant_config, - bias=getattr(config, "mlp_bias", False), - prefix=f"{prefix}.mlp", - ) - # Layer normalization layers (RMSNorm as in LLaMA) - self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = RMSNorm( - config.hidden_size, eps=config.rms_norm_eps - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - residual: torch.Tensor | None, - ) -> tuple[torch.Tensor, torch.Tensor]: - # Self-Attention block - if residual is None: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - else: - # Fused residual add + layernorm if supported - hidden_states, residual = self.input_layernorm(hidden_states, residual) - hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states) - # Feed-forward block - hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) - hidden_states = self.mlp(hidden_states) - return hidden_states, residual - - -@support_torch_compile -class ArceeModel(nn.Module, EagleModelMixin): - """The transformer model backbone for Arcee (embedding layer + stacked - decoder blocks + final norm).""" - - def __init__( - self, - *, - vllm_config, - prefix: str = "", - layer_type: type[nn.Module] = ArceeDecoderLayer, - ) -> None: - super().__init__() - config: LlamaConfig = vllm_config.model_config.hf_config - cache_config = vllm_config.cache_config - quant_config = vllm_config.quant_config - self.quant_config = quant_config - self.config = config - self.vocab_size = config.vocab_size - - # Word embeddings (parallelized if using pipeline parallel) - if get_pp_group().is_first_rank or ( - config.tie_word_embeddings and get_pp_group().is_last_rank - ): - self.embed_tokens = VocabParallelEmbedding( - self.vocab_size, - config.hidden_size, - quant_config=quant_config, - ) - else: - self.embed_tokens = PPMissingLayer() # placeholder on non-embedding ranks - - # Build decoder layers across pipeline ranks - self.start_layer, self.end_layer, self.layers = make_layers( - config.num_hidden_layers, - lambda prefix: layer_type( - config=config, - cache_config=cache_config, - quant_config=quant_config, - prefix=prefix, - ), - prefix=f"{prefix}.layers", - ) - # Final RMSNorm on the last pipeline stage - if get_pp_group().is_last_rank: - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - else: - self.norm = PPMissingLayer() - - # Prepare factory for empty intermediate tensors - # (for pipeline scheduling) - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( - ["hidden_states", "residual"], config.hidden_size - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: - # Embedding lookup (on first pipeline rank) - if get_pp_group().is_first_rank: - hidden_states = ( - inputs_embeds - if inputs_embeds is not None - else self.embed_input_ids(input_ids) - ) - residual = None - else: - assert intermediate_tensors is not None, ( - "IntermediateTensors must be provided for non-first pipeline ranks" - ) - hidden_states = intermediate_tensors["hidden_states"] - residual = intermediate_tensors["residual"] - - aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual) - for idx, layer in enumerate( - islice(self.layers, self.start_layer, self.end_layer) - ): - hidden_states, residual = layer(positions, hidden_states, residual) - self._maybe_add_hidden_state( - aux_hidden_states, idx + 1, hidden_states, residual - ) - - if not get_pp_group().is_last_rank: - # Send intermediate results to the next pipeline stage - return IntermediateTensors( - {"hidden_states": hidden_states, "residual": residual} - ) - # On last rank: apply final layer norm - hidden_states, _ = self.norm(hidden_states, residual) - if len(aux_hidden_states) > 0: - return hidden_states, aux_hidden_states - return hidden_states - - -class ArceeForCausalLM( - nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3 -): - """Arcee Model for causal language modeling, integrated with vLLM - runtime.""" - - hf_to_vllm_mapper = WeightsMapper( - orig_to_new_stacked={ - # weight_name: (param_name, shard_id) - ".q_proj": (".qkv_proj", "q"), - ".k_proj": (".qkv_proj", "k"), - ".v_proj": (".qkv_proj", "v"), - }, - orig_to_new_substr={ - "gate_proj": None, - }, - ) - # Map fused module names to their submodule components - # (for quantization and LoRA) - packed_modules_mapping = { - "qkv_proj": ["q_proj", "k_proj", "v_proj"], - } - - def __init__(self, *, vllm_config, prefix: str = "") -> None: - super().__init__() - config = vllm_config.model_config.hf_config - self.config = config - - # Initialize the inner Transformer model (ArceeModel) - self.model = ArceeModel( - vllm_config=vllm_config, - prefix=maybe_prefix(prefix, "model"), - ) - # On the last pipeline stage, set up the LM head and logits processor - if get_pp_group().is_last_rank: - # Determine vocabulary size (including any LoRA extra tokens - # for padded LM head) - - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=vllm_config.quant_config, - bias=getattr(config, "lm_head_bias", False), - prefix=maybe_prefix(prefix, "lm_head"), - ) - if config.tie_word_embeddings: - # Tie output weights with input embedding matrix - self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) - logit_scale = getattr(config, "logit_scale", 1.0) - self.logits_processor = LogitsProcessor( - config.vocab_size, scale=logit_scale - ) - else: - # Placeholder for lm_head on non-last ranks - self.lm_head = PPMissingLayer() - - self.make_empty_intermediate_tensors = ( - self.model.make_empty_intermediate_tensors - ) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: - model_output = self.model( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds, - ) - return model_output - - def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: - # Compute final logits from hidden states (last pipeline rank only) - logits = self.logits_processor(self.lm_head, hidden_states) - return logits - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.model.embed_input_ids(input_ids) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """Load weights into the model (delegates to inner model and handles - tied embeddings).""" - loader = AutoWeightsLoader(self) - # AutoWeightLoader handles weight name remapping, including fusing - # separate q_proj, k_proj, v_proj into qkv_proj - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/bert_with_rope.py b/vllm/model_executor/models/bert_with_rope.py index cd35792b28..ed0ee78b2b 100644 --- a/vllm/model_executor/models/bert_with_rope.py +++ b/vllm/model_executor/models/bert_with_rope.py @@ -572,7 +572,6 @@ class GteNewModel(BertWithRope): super().__init__(vllm_config=vllm_config, prefix=prefix, **kwargs) # GteNewModel only gate_up_proj does not have bias. - # Hack method learned from vllm/model_executor/models/glm.py for layer in self.encoder.layers: layer.mlp.gate_up_proj.bias = None layer.mlp.gate_up_proj.skip_bias_add = True diff --git a/vllm/model_executor/models/glm.py b/vllm/model_executor/models/glm.py deleted file mode 100644 index 26d7c29aae..0000000000 --- a/vllm/model_executor/models/glm.py +++ /dev/null @@ -1,24 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Inference-only HF format GLM-4 model compatible with THUDM weights.""" - -from vllm.config import VllmConfig -from vllm.model_executor.models.llama import LlamaForCausalLM - -from .utils import PPMissingLayer - - -class GlmForCausalLM(LlamaForCausalLM): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - hf_config = vllm_config.model_config.hf_config - hf_config.rope_parameters["partial_rotary_factor"] = 0.5 - super().__init__(vllm_config=vllm_config, prefix=prefix) - # Hack Llama model to fit HF format GLM implementation - # Attention difference between GLM and Llama: - # 1. Half partial rotary_dim and no Neox style. - # 2. There is no bias for o_proj in attention - for layer in self.model.layers: - if not isinstance(layer, PPMissingLayer): - layer.self_attn.rotary_emb.is_neox_style = False - layer.self_attn.o_proj.bias = None - layer.self_attn.o_proj.skip_bias_add = True diff --git a/vllm/model_executor/models/mellum.py b/vllm/model_executor/models/mellum.py deleted file mode 100644 index a75e3b0a8e..0000000000 --- a/vllm/model_executor/models/mellum.py +++ /dev/null @@ -1,253 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from typing import Any - -from torch import nn - -from vllm.compilation.decorators import support_torch_compile -from vllm.config import VllmConfig -from vllm.distributed import get_tensor_model_parallel_world_size -from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.rotary_embedding import get_rope -from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead - -from .qwen3_moe import ( - Qwen3MoeAttention, - Qwen3MoeDecoderLayer, - Qwen3MoeForCausalLM, - Qwen3MoeMLP, - Qwen3MoeModel, - Qwen3MoeSparseMoeBlock, -) -from .utils import PPMissingLayer, extract_layer_index, maybe_prefix - - -class MellumAttention(Qwen3MoeAttention): - """Differences from `Qwen3MoeAttention`: - - Supports `per_layer_sliding_window` for `Attention`. - """ - - def __init__( - self, - hidden_size: int, - num_heads: int, - num_kv_heads: int, - rope_parameters: dict[str, Any], - max_position_embeddings: int = 8192, - head_dim: int | None = None, - rms_norm_eps: float = 1e-06, - qkv_bias: bool = False, - cache_config: Any | None = None, - quant_config: Any | None = None, - prefix: str = "", - dual_chunk_attention_config: dict[str, Any] | None = None, - per_layer_sliding_window: int | None = None, - ) -> None: - nn.Module.__init__(self) - - self.hidden_size = hidden_size - tp_size = get_tensor_model_parallel_world_size() - self.total_num_heads = num_heads - assert self.total_num_heads % tp_size == 0 - self.num_heads = self.total_num_heads // tp_size - self.total_num_kv_heads = num_kv_heads - if self.total_num_kv_heads >= tp_size: - assert self.total_num_kv_heads % tp_size == 0 - else: - assert tp_size % self.total_num_kv_heads == 0 - self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - self.head_dim = head_dim or (hidden_size // self.total_num_heads) - self.q_size = self.num_heads * self.head_dim - self.kv_size = self.num_kv_heads * self.head_dim - self.scaling = self.head_dim**-0.5 - self.max_position_embeddings = max_position_embeddings - self.dual_chunk_attention_config = dual_chunk_attention_config - - self.qkv_proj = QKVParallelLinear( - hidden_size, - self.head_dim, - self.total_num_heads, - self.total_num_kv_heads, - bias=qkv_bias, - quant_config=quant_config, - prefix=f"{prefix}.qkv_proj", - ) - - self.o_proj = RowParallelLinear( - self.total_num_heads * self.head_dim, - hidden_size, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.o_proj", - ) - - self.rotary_emb = get_rope( - self.head_dim, - max_position=max_position_embeddings, - rope_parameters=rope_parameters, - dual_chunk_attention_config=dual_chunk_attention_config, - ) - if dual_chunk_attention_config: - self.attn = Attention( - self.num_heads, - self.head_dim, - self.scaling, - num_kv_heads=self.num_kv_heads, - cache_config=cache_config, - quant_config=quant_config, - per_layer_sliding_window=per_layer_sliding_window, - prefix=f"{prefix}.attn", - layer_idx=extract_layer_index(prefix), - dual_chunk_attention_config=dual_chunk_attention_config, - ) - else: - self.attn = Attention( - self.num_heads, - self.head_dim, - self.scaling, - num_kv_heads=self.num_kv_heads, - cache_config=cache_config, - quant_config=quant_config, - per_layer_sliding_window=per_layer_sliding_window, - prefix=f"{prefix}.attn", - ) - - self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) - self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) - - -class MellumDecoderLayer(Qwen3MoeDecoderLayer): - """Differences from `Qwen3MoeDecoderLayer`: - - Supports interleaved SWA and per-layer RoPE scaling. - """ - - def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: - nn.Module.__init__(self) - - config = vllm_config.model_config.hf_text_config - cache_config = vllm_config.cache_config - quant_config = vllm_config.quant_config - - self.hidden_size = config.hidden_size - max_position_embeddings = getattr(config, "max_position_embeddings", 8192) - dual_chunk_attention_config = getattr( - config, "dual_chunk_attention_config", None - ) - - layer_idx = extract_layer_index(prefix) - layer_type = config.layer_types[layer_idx] - if layer_type == "sliding_attention": - sliding_window = getattr(config, "sliding_window", None) - else: - sliding_window = None - rope_parameters = config.rope_parameters[layer_type] - - self.self_attn = MellumAttention( - hidden_size=self.hidden_size, - num_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - rope_parameters=rope_parameters, - max_position_embeddings=max_position_embeddings, - rms_norm_eps=config.rms_norm_eps, - qkv_bias=getattr(config, "attention_bias", False), - head_dim=getattr(config, "head_dim", None), - cache_config=cache_config, - quant_config=quant_config, - prefix=f"{prefix}.self_attn", - dual_chunk_attention_config=dual_chunk_attention_config, - per_layer_sliding_window=sliding_window, - ) - - if config.mlp_layer_types[layer_idx] == "sparse": - self.mlp = Qwen3MoeSparseMoeBlock( - vllm_config=vllm_config, prefix=f"{prefix}.mlp" - ) - else: - self.mlp = Qwen3MoeMLP( - hidden_size=config.hidden_size, - intermediate_size=config.intermediate_size, - hidden_act=config.hidden_act, - quant_config=quant_config, - prefix=f"{prefix}.mlp", - ) - - self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = RMSNorm( - config.hidden_size, eps=config.rms_norm_eps - ) - - -@support_torch_compile -class MellumModel(Qwen3MoeModel): - """Differences from `Qwen3MoeModel`: - - Uses `MellumDecoderLayer`. - """ - - def __init__( - self, - *, - vllm_config: VllmConfig, - prefix: str = "", - ): - super().__init__( - vllm_config=vllm_config, - prefix=prefix, - decoder_layer_type=MellumDecoderLayer, - ) - - -class MellumForCausalLM(Qwen3MoeForCausalLM): - """Differences from `Qwen3MoeForCausalLM`: - - Uses `MellumModel`. - """ - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - nn.Module.__init__(self) - config = vllm_config.model_config.hf_text_config - quant_config = vllm_config.quant_config - self.config = config - self.quant_config = quant_config - if "dense" in getattr(config, "mlp_layer_types", []): - self.packed_modules_mapping["gate_up_proj"] = ["gate_proj", "up_proj"] - self.model = MellumModel( - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") - ) - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) - if self.config.tie_word_embeddings: - self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) - self.logits_processor = LogitsProcessor(config.vocab_size) - self.make_empty_intermediate_tensors = ( - self.model.make_empty_intermediate_tensors - ) - - self.moe_layers = [] - example_layer = None - for layer in self.model.layers: - if isinstance(layer, PPMissingLayer): - continue - - assert isinstance(layer, Qwen3MoeDecoderLayer) - if isinstance(layer.mlp, Qwen3MoeSparseMoeBlock): - example_layer = layer.mlp - self.moe_layers.append(layer.mlp.experts) - - if example_layer is None: - raise RuntimeError("No MoE layer found in the model.layers.") - - self.num_moe_layers = len(self.moe_layers) - self.num_expert_groups = 1 - self.num_shared_experts = 0 - self.num_logical_experts = example_layer.n_logical_experts - self.num_physical_experts = example_layer.n_physical_experts - self.num_local_physical_experts = example_layer.n_local_physical_experts - self.num_routed_experts = example_layer.n_routed_experts - self.num_redundant_experts = example_layer.n_redundant_experts diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index b435d77ad7..81763fed40 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -72,7 +72,6 @@ _TEXT_GENERATION_MODELS = { # [Decoder-only] "AfmoeForCausalLM": ("afmoe", "AfmoeForCausalLM"), "ApertusForCausalLM": ("apertus", "ApertusForCausalLM"), - "ArceeForCausalLM": ("arcee", "ArceeForCausalLM"), "AXK1ForCausalLM": ("AXK1", "AXK1ForCausalLM"), "BailingMoeForCausalLM": ("bailing_moe", "BailingMoeForCausalLM"), "BailingMoeV2ForCausalLM": ("bailing_moe", "BailingMoeV2ForCausalLM"), @@ -84,7 +83,6 @@ _TEXT_GENERATION_MODELS = { "CohereForCausalLM": ("commandr", "CohereForCausalLM"), "Cohere2ForCausalLM": ("commandr", "CohereForCausalLM"), "Cohere2MoeForCausalLM": ("cohere2_moe", "Cohere2MoeForCausalLM"), - "CwmForCausalLM": ("llama", "LlamaForCausalLM"), "DbrxForCausalLM": ("dbrx", "DbrxForCausalLM"), "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"), "DeepseekForCausalLM": ("deepseek_v2", "DeepseekForCausalLM"), @@ -114,7 +112,6 @@ _TEXT_GENERATION_MODELS = { "vllm.models.qwen4_exp", "Qwen4ExpForCausalLM", ), - "GlmForCausalLM": ("glm", "GlmForCausalLM"), "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"), "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"), "Glm4MoeLiteForCausalLM": ("glm4_moe_lite", "Glm4MoeLiteForCausalLM"), @@ -160,7 +157,6 @@ _TEXT_GENERATION_MODELS = { ), "MambaForCausalLM": ("mamba", "MambaForCausalLM"), "Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"), - "MellumForCausalLM": ("mellum", "MellumForCausalLM"), "MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"), "MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"), "MiniMaxM2ForCausalLM": ("minimax_m2", "MiniMaxM2ForCausalLM"), @@ -226,7 +222,6 @@ _EMBEDDING_MODELS = { "DeepseekV3BidirectionalModel": ("deepseek_v2", "DeepseekV3ForCausalLM"), "Gemma2Model": ("gemma2", "Gemma2ForCausalLM"), "Gemma3TextModel": ("gemma3", "Gemma3Model"), - "GlmForCausalLM": ("glm", "GlmForCausalLM"), "GteModel": ("bert_with_rope", "SnowflakeGteNewModel"), "GteNewModel": ("bert_with_rope", "GteNewModel"), "JinaEmbeddingsV5Model": ("jina", "JinaEmbeddingsV5Model"), @@ -709,10 +704,14 @@ _SPECULATIVE_DECODING_MODELS = { _TRANSFORMERS_SUPPORTED_MODELS = { # Text generation models + "ArceeForCausalLM": ("transformers", "TransformersForCausalLM"), + "CwmForCausalLM": ("transformers", "TransformersForCausalLM"), "FlexOlmoForCausalLM": ("transformers", "TransformersMoEForCausalLM"), + "GlmForCausalLM": ("transformers", "TransformersForCausalLM"), "GPTBigCodeForCausalLM": ("transformers", "TransformersForCausalLM"), "HunYuanDenseV1ForCausalLM": ("transformers", "TransformersForCausalLM"), "HunYuanMoEV1ForCausalLM": ("transformers", "TransformersMoEForCausalLM"), + "MellumForCausalLM": ("transformers", "TransformersMoEForCausalLM"), "NanbeigeForCausalLM": ("transformers", "TransformersForCausalLM"), "OlmoForCausalLM": ("transformers", "TransformersForCausalLM"), "Olmo2ForCausalLM": ("transformers", "TransformersForCausalLM"),