diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index b78428f68c..6678cc8fb0 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -378,6 +378,7 @@ th { | `BloomForCausalLM` | BLOOM, BLOOMZ, BLOOMChat | `bigscience/bloom`, `bigscience/bloomz`, etc. | | ✅︎ | | `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. | ✅︎ | ✅︎ | +| `CohereMoeForCausalLM` | Command (MoE) | (model checkpoints loaded with `trust_remote_code=True`) | ✅︎ | ✅︎ | | `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. | ✅︎ | ✅︎ | diff --git a/tests/models/registry.py b/tests/models/registry.py index 3a82442381..36c69160e0 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -238,6 +238,11 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "CohereLabs/c4ai-command-r7b-12-2024", trust_remote_code=True, ), + "CohereMoeForCausalLM": _HfExamplesInfo( + "/host/engines/cohere-moe", + trust_remote_code=True, + is_available_online=False, + ), "CwmForCausalLM": _HfExamplesInfo("facebook/cwm", min_transformers_version="4.58"), # FIXME: databricks/dbrx-instruct has been deleted "DbrxForCausalLM": _HfExamplesInfo( diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index baa7d3fd3e..16d3f8d01d 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -263,6 +263,7 @@ class TrtLlmNvFp4ExpertsMonolithic( RoutingMethodType.SigmoidRenorm, RoutingMethodType.MiniMax2, RoutingMethodType.Simulated, + RoutingMethodType.SigmoidRenorm, ] @staticmethod diff --git a/vllm/model_executor/layers/fused_moe/router/custom_routing_router.py b/vllm/model_executor/layers/fused_moe/router/custom_routing_router.py index c1bd7a6993..8e35169d90 100644 --- a/vllm/model_executor/layers/fused_moe/router/custom_routing_router.py +++ b/vllm/model_executor/layers/fused_moe/router/custom_routing_router.py @@ -34,11 +34,15 @@ class CustomRoutingRouter(BaseRouter): @property def routing_method_type(self) -> RoutingMethodType: + from vllm.model_executor.models.cohere_moe import token_choice_with_bias from vllm.model_executor.models.llama4 import Llama4MoE # NOTE: FLASHINFER_TRTLLM support the Llama4 router. if self.custom_routing_function == Llama4MoE.custom_routing_function: return RoutingMethodType.Llama4 + # Cohere MoE uses a sigmoid -> top-k -> renormalize routing function. + if self.custom_routing_function == token_choice_with_bias: + return RoutingMethodType.SigmoidRenorm return RoutingMethodType.Custom def _compute_routing( diff --git a/vllm/model_executor/models/cohere2_vision.py b/vllm/model_executor/models/cohere2_vision.py index c3118ee777..cdb3bcbedb 100644 --- a/vllm/model_executor/models/cohere2_vision.py +++ b/vllm/model_executor/models/cohere2_vision.py @@ -44,7 +44,12 @@ from vllm.multimodal.processing import ( from vllm.sequence import IntermediateTensors from vllm.utils.tensor_schema import TensorSchema, TensorShape -from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP +from .interfaces import ( + MultiModalEmbeddings, + SupportsMultiModal, + SupportsPP, + SupportsQuant, +) from .siglip import SiglipVisionModel from .utils import ( AutoWeightsLoader, @@ -309,16 +314,22 @@ class Cohere2VisionMultiModalProcessor( info=Cohere2VisionProcessingInfo, dummy_inputs=Cohere2VisionDummyInputsBuilder, ) -class Cohere2VisionForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP): +class Cohere2VisionForConditionalGeneration( + nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant +): hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ "model.vision_tower.": "vision_tower.", "model.multi_modal_projector.": "multi_modal_projector.", "model.language_model.": "language_model.model.", - "lm_head.": "language_model.lm_head.", } ) + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config: Cohere2VisionConfig = vllm_config.model_config.hf_config diff --git a/vllm/model_executor/models/cohere_moe.py b/vllm/model_executor/models/cohere_moe.py new file mode 100644 index 0000000000..587a18d5ae --- /dev/null +++ b/vllm/model_executor/models/cohere_moe.py @@ -0,0 +1,528 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Iterable +from itertools import islice + +import torch +from torch import nn +from transformers import CohereConfig + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import CacheConfig, VllmConfig +from vllm.distributed import ( + get_pp_group, + get_tensor_model_parallel_world_size, +) +from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.platforms import current_platform +from vllm.sequence import IntermediateTensors + +from .commandr import LayerNorm +from .interfaces import SupportsPP, SupportsQuant +from .utils import ( + extract_layer_index, + is_pp_missing_parameter, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_prefix, +) + + +@torch.compile(backend=current_platform.simple_compile_backend) +def token_choice_with_bias( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, +): + """Sigmoid -> top-k (-> renormalize) custom routing for CohereMoe.""" + assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" + + scores = gating_output.float().sigmoid() + topk_weights, topk_ids = torch.topk(scores, k=topk, dim=-1, sorted=False) + + if renormalize: + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) + + return topk_weights.to(torch.float32), topk_ids.to(torch.int32) + + +class CohereMoeMLP(nn.Module): + """Cohere MLP used as shared experts in the MoE block.""" + + def __init__( + self, + config: CohereConfig, + intermediate_size: int | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = ( + intermediate_size + if intermediate_size is not None + else config.intermediate_size + ) + self.gate_up_proj = MergedColumnParallelLinear( + self.hidden_size, + [self.intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + self.intermediate_size, + self.hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.down_proj", + ) + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class CohereMoeAttention(nn.Module): + """Cohere MoE attention with sliding-window interleave.""" + + def __init__( + self, + config: CohereConfig, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + tp_size = get_tensor_model_parallel_world_size() + self.config = config + self.layer_idx = extract_layer_index(prefix) + self.hidden_size = config.hidden_size + self.total_num_heads = config.num_attention_heads + self.num_heads = self.total_num_heads // tp_size + self.head_dim = getattr( + config, "head_dim", self.hidden_size // self.total_num_heads + ) + self.total_num_kv_heads = config.num_key_value_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.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 = getattr( + config, "model_max_length", None + ) or getattr(config, "max_position_embeddings", 8192) + self.qkv_proj = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + self.rotary_emb = get_rope( + self.head_dim, + max_position=self.max_position_embeddings, + rope_parameters=config.rope_parameters, + is_neox_style=False, + ) + + self.sliding_window = None + layer_types = getattr(config, "layer_types", None) + if ( + layer_types is not None + and layer_types[self.layer_idx] == "sliding_attention" + ): + self.sliding_window = config.sliding_window + + 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=self.sliding_window, + prefix=f"{prefix}.attn", + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + if self.sliding_window: + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v) + output, _ = self.o_proj(attn_output) + return output + + +class CohereMoe(nn.Module): + """Tensor-parallel MoE block for CohereMoe with shared experts.""" + + def __init__( + self, + config: CohereConfig, + params_dtype: torch.dtype | None = None, + quant_config: QuantizationConfig | None = None, + tp_size: int | None = None, + prefix: str = "", + ): + super().__init__() + self.hidden_size = config.hidden_size + self.tp_size = get_tensor_model_parallel_world_size() + + if self.tp_size > config.num_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than " + f"the number of experts {config.num_experts}." + ) + + if ( + hasattr(config, "expert_selection_fn") + and config.expert_selection_fn == "sigmoid" + ): + self.custom_routing_function = token_choice_with_bias + else: + self.custom_routing_function = None + + self.gate = ReplicatedLinear( + config.hidden_size, + config.num_experts, + bias=False, + params_dtype=params_dtype, + quant_config=None, + prefix=f"{prefix}.gate", + ) + + if hasattr(config, "num_shared_experts") and config.num_shared_experts > 0: + self.shared_experts = CohereMoeMLP( + config=config, + intermediate_size=config.intermediate_size * config.num_shared_experts, + quant_config=quant_config, + prefix=f"{prefix}.shared_experts", + ) + self.shared_expert_combination_strategy = getattr( + config, "shared_expert_combination_strategy", "sum" + ) + assert self.shared_expert_combination_strategy in ("average", "sum"), ( + "shared_expert_combination_strategy must be one of ['average', 'sum']" + ) + else: + self.shared_experts = None + self.shared_expert_combination_strategy = None + + self.experts = FusedMoE( + num_experts=config.num_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + params_dtype=params_dtype, + renormalize=getattr(config, "norm_topk_prob", True), + quant_config=quant_config, + tp_size=tp_size, + prefix=f"{prefix}.experts", + custom_routing_function=self.custom_routing_function, + shared_experts=self.shared_experts, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + orig_shape = hidden_states.shape + hidden_states = hidden_states.view(-1, self.hidden_size) + router_logits, _ = self.gate(hidden_states) + # FusedMoE handles shared expert overlap internally and returns + # shared_output + routed_output when shared_experts is set. + final_hidden_states = self.experts(hidden_states, router_logits) + if self.shared_expert_combination_strategy == "average": + final_hidden_states = final_hidden_states / 2 + return final_hidden_states.view(orig_shape) + + +class CohereMoeDecoderLayer(nn.Module): + def __init__( + self, + config: CohereConfig, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + + self.self_attn = CohereMoeAttention( + config, + cache_config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + ) + self.mlp = CohereMoe( + config=config, quant_config=quant_config, prefix=f"{prefix}.mlp" + ) + self.input_layernorm = LayerNorm( + param_shape=(config.hidden_size,), eps=config.layer_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + residual = hidden_states + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states_attention = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ) + hidden_states_mlp = self.mlp(hidden_states) + + hidden_states = residual + hidden_states_attention + hidden_states_mlp + return hidden_states, residual + + +@support_torch_compile +class CohereMoeModel(nn.Module): + """Transformer decoder for CohereMoe.""" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + + self.config = config + self.vocab_size = config.vocab_size + self.org_vocab_size = config.vocab_size + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, config.hidden_size + ) + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: CohereMoeDecoderLayer( + config, cache_config, quant_config, prefix=prefix + ), + prefix=f"{prefix}.layers", + ) + self.norm = LayerNorm( + param_shape=(config.hidden_size,), eps=config.layer_norm_eps + ) + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], config.hidden_size + ) + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.get_input_embeddings(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + for layer in islice(self.layers, self.start_layer, self.end_layer): + hidden_states, residual = layer(positions, hidden_states, residual) + if not get_pp_group().is_last_rank: + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class CohereMoeForCausalLM(nn.Module, SupportsPP, SupportsQuant): + is_text_generation_model = True + + packed_modules_mapping = { + "qkv_proj": [ + "q_proj", + "k_proj", + "v_proj", + ], + "gate_up_proj": [ + "gate_proj", + "up_proj", + ], + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + assert getattr(config, "tie_word_embeddings", True) + self.unpadded_vocab_size = config.vocab_size + self.quant_config = quant_config + self.logits_scale = config.logit_scale + self.logits_processor = LogitsProcessor( + self.unpadded_vocab_size, config.vocab_size, scale=self.logits_scale + ) + self.model = CohereMoeModel( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.get_input_embeddings(input_ids) + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.get_input_embeddings(input_ids) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds) + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + return self.logits_processor(self.model.embed_tokens, hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + expert_params_mapping = FusedMoE.make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.num_experts, + ) + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + + if self.quant_config is not None and ( + scale_name := self.quant_config.get_cache_scale(name) + ): + param = params_dict[scale_name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + loaded_weight = ( + loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0] + ) + weight_loader(param, loaded_weight) + loaded_params.add(scale_name) + continue + + for param_name, shard_name, shard_id in stacked_params_mapping: + if shard_name not in name: + continue + if "mlp.experts" in name: + continue + name = name.replace(shard_name, param_name) + if name.endswith(".bias") and name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + if is_pp_missing_parameter(name, self): + continue + if ( + name.endswith(".bias") or name.endswith("_bias") + ) and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + shard_id=shard_id, + expert_id=expert_id, + ) + break + else: + if "lm_head.weight" in name: + continue + if ( + name.endswith(".bias") or name.endswith("_bias") + ) and name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + return loaded_params diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 4369d715e6..9a9884807e 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -89,6 +89,7 @@ _TEXT_GENERATION_MODELS = { "ChatGLMForConditionalGeneration": ("chatglm", "ChatGLMForCausalLM"), "CohereForCausalLM": ("commandr", "CohereForCausalLM"), "Cohere2ForCausalLM": ("commandr", "CohereForCausalLM"), + "CohereMoeForCausalLM": ("cohere_moe", "CohereMoeForCausalLM"), "CwmForCausalLM": ("llama", "LlamaForCausalLM"), "DbrxForCausalLM": ("dbrx", "DbrxForCausalLM"), "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),