From 84d26b9ee3dde1991442539dcf5fc5b26aa08f7e Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Fri, 24 Jul 2026 21:54:24 +0100 Subject: [PATCH] [Model] Remove Plamo2 (#49729) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- .buildkite/test-amd.yaml | 6 +- .buildkite/test_areas/models_language.yaml | 2 - docs/design/custom_op.md | 2 - docs/models/supported_models.md | 1 - docs/usage/v1_guide.md | 2 +- tests/distributed/test_pipeline_parallel.py | 1 - .../models/language/generation/test_hybrid.py | 2 - tests/models/registry.py | 4 - tests/quantization/test_experts_int8.py | 2 +- vllm/config/compilation.py | 1 - vllm/config/model.py | 3 +- vllm/model_executor/models/plamo2.py | 992 ------------------ vllm/model_executor/models/registry.py | 2 +- 13 files changed, 6 insertions(+), 1014 deletions(-) delete mode 100644 vllm/model_executor/models/plamo2.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index b8046fda9f5..b67acbc40b3 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -81,10 +81,8 @@ # the above test.) Also run if model initialization test file is modified. # # * [Language Models Tests (Extra Standard) %N]: Shard slow subset of standard language models tests. Only run when model # # source is modified, or when specified test files are modified. # -# * [Language Models Tests (Hybrid) %N]: Install fast path packages for testing against transformers (mamba, conv1d) and to # -# run plamo2 model in vLLM. # -# * [Language Models Test (Extended Generation)]: Install fast path packages for testing against transformers (mamba, conv1d) # -# and to run plamo2 model in vLLM. # +# * [Language Models Tests (Hybrid) %N]: Install fast path packages for testing against transformers (mamba, conv1d). # +# * [Language Models Test (Extended Generation)]: Install fast path packages for testing against transformers (mamba, conv1d). # # * [Multi-Modal Models (Standard) 1-4]: # # - Do NOT remove `VLLM_WORKER_MULTIPROC_METHOD=spawn` setting as ROCm requires this for certain models to function. # # * [Transformers Nightly Models]: Whisper needs `VLLM_WORKER_MULTIPROC_METHOD=spawn` to avoid deadlock. # diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index 6124cb023dd..2d3f14d1715 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -63,7 +63,6 @@ steps: - tests/models/language/generation commands: # Install fast path packages for testing against transformers - # Note: also needed to run plamo2 model in vLLM - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0' - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' # Shard the hybrid language model tests that are numerically stable on Hopper. @@ -105,7 +104,6 @@ steps: - tests/models/language/generation commands: # Install fast path packages for testing against transformers - # Note: also needed to run plamo2 model in vLLM - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0' - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' diff --git a/docs/design/custom_op.md b/docs/design/custom_op.md index d2557a2281c..f99c8043fa0 100644 --- a/docs/design/custom_op.md +++ b/docs/design/custom_op.md @@ -122,8 +122,6 @@ For example: --8<-- "vllm/model_executor/layers/mamba/mamba_mixer2.py:mixer2_gated_rms_norm" ---8<-- "vllm/model_executor/models/plamo2.py:plamo2_mamba_mixer" - --8<-- "vllm/model_executor/layers/mamba/short_conv.py:short_conv" ``` diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 98f9cb5b65d..f1292d7a988 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -435,7 +435,6 @@ th { | `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ | | `Phi3ForCausalLM` | Phi-4, Phi-3 | `microsoft/Phi-4-mini-instruct`, `microsoft/Phi-4`, `microsoft/Phi-3-mini-4k-instruct`, `microsoft/Phi-3-mini-128k-instruct`, `microsoft/Phi-3-medium-128k-instruct`, etc. | ✅︎ | ✅︎ | | `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ | -| `Plamo2ForCausalLM` | PLaMo2 | `pfnet/plamo-2-1b`, `pfnet/plamo-2-8b`, etc. | ✅ | ✅︎ | | `Plamo3ForCausalLM` | PLaMo3 | `pfnet/plamo-3-nict-2b-base`, `pfnet/plamo-3-nict-8b-base`, etc. | ✅ | ✅︎ | | `Qwen2ForCausalLM` | QwQ, Qwen2 | `Qwen/QwQ-32B-Preview`, `Qwen/Qwen2-7B-Instruct`, `Qwen/Qwen2-7B`, etc. | ✅︎ | ✅︎ | | `Qwen2MoeForCausalLM` | Qwen2MoE | `Qwen/Qwen1.5-MoE-A2.7B`, `Qwen/Qwen1.5-MoE-A2.7B-Chat`, etc. | ✅︎ | ✅︎ | diff --git a/docs/usage/v1_guide.md b/docs/usage/v1_guide.md index 5613d5ba4e8..4ef2b0b5b6a 100644 --- a/docs/usage/v1_guide.md +++ b/docs/usage/v1_guide.md @@ -126,7 +126,7 @@ Models using selective state-space mechanisms instead of standard transformer at Models that use Mamba-2 and Mamba-1 layers (e.g., `Mamba2ForCausalLM`, `MambaForCausalLM`, `FalconMambaForCausalLM`) are supported. Hybrid models that combine Mamba-2 and Mamba-1 layers with standard attention layers are also supported (e.g., -`Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`, `Plamo2ForCausalLM`). +`Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`). Hybrid models with mechanisms different to Mamba are also supported (e.g, `Lfm2ForCausalLM`). diff --git a/tests/distributed/test_pipeline_parallel.py b/tests/distributed/test_pipeline_parallel.py index fd4a1d0f570..22e0321da60 100644 --- a/tests/distributed/test_pipeline_parallel.py +++ b/tests/distributed/test_pipeline_parallel.py @@ -121,7 +121,6 @@ TEXT_GENERATION_MODELS = { "ibm/PowerMoE-3b": PPTestSettings.fast(), "internlm/internlm2-chat-7b": PPTestSettings.fast(), "ai21labs/Jamba-tiny-dev": PPTestSettings.fast(), - "pfnet/plamo-2-1b": PPTestSettings.fast(), "pfnet/plamo-3-nict-2b-base": PPTestSettings.fast(), "meta-llama/Llama-3.2-1B-Instruct": PPTestSettings.detailed(), # Tests TransformersForCausalLM diff --git a/tests/models/language/generation/test_hybrid.py b/tests/models/language/generation/test_hybrid.py index 3fee0662bc6..768b26c8688 100644 --- a/tests/models/language/generation/test_hybrid.py +++ b/tests/models/language/generation/test_hybrid.py @@ -35,7 +35,6 @@ SSM_MODELS = [ HYBRID_MODELS = [ "ai21labs/Jamba-tiny-dev", - "pfnet/plamo-2-1b", "Zyphra/Zamba2-1.2B-instruct", "ibm-granite/granite-4.0-tiny-preview", "tiiuae/Falcon-H1-0.5B-Base", @@ -50,7 +49,6 @@ HYBRID_MODELS_REQUIRING_CHUNKED_PREFILL = { FULL_CUDA_GRAPH_MODELS = [ "ai21labs/Jamba-tiny-dev", - "pfnet/plamo-2-1b", "Zyphra/Zamba2-1.2B-instruct", ] diff --git a/tests/models/registry.py b/tests/models/registry.py index 4362e134c76..65cdc84dec8 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -484,10 +484,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "PhiMoEForCausalLM": _HfExamplesInfo( "microsoft/Phi-3.5-MoE-instruct", trust_remote_code=True ), - "Plamo2ForCausalLM": _HfExamplesInfo( - "pfnet/plamo-2-1b", - trust_remote_code=True, - ), "Plamo3ForCausalLM": _HfExamplesInfo( "pfnet/plamo-3-nict-2b-base", trust_remote_code=True, diff --git a/tests/quantization/test_experts_int8.py b/tests/quantization/test_experts_int8.py index 7cdb135fa07..6119c1c8468 100644 --- a/tests/quantization/test_experts_int8.py +++ b/tests/quantization/test_experts_int8.py @@ -12,7 +12,7 @@ from tests.quantization.utils import is_quant_method_supported from ..models.registry import HF_EXAMPLE_MODELS -MODELS = ["ai21labs/Jamba-tiny-random", "pfnet/plamo-2-1b"] +MODELS = ["ai21labs/Jamba-tiny-random"] @pytest.mark.skipif( diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 810e97c4cad..8142004748a 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -768,7 +768,6 @@ class CompilationConfig: "vllm::mamba_mixer", "vllm::short_conv", "vllm::linear_attention", - "vllm::plamo2_mamba_mixer", "vllm::qwen_gdn_attention_core", "vllm::gdn_attention_core_xpu", "vllm::olmo_hybrid_gdn_full_forward", diff --git a/vllm/config/model.py b/vllm/config/model.py index 1f62351c526..a65d5d5a94f 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1448,7 +1448,7 @@ class ModelConfig: """ Returns the mamba chunk size if it exists """ - # used by e.g. Bamba, FalconH1, Granite, PLaMo2 + # used by e.g. Bamba, FalconH1, Granite chunk_size = getattr(self.hf_text_config, "mamba_chunk_size", None) if chunk_size is None: # used by e.g. Mamba2, NemotronH, Zamba @@ -2022,7 +2022,6 @@ _FLOAT16_NOT_SUPPORTED_MODELS = { "gemma2": "Numerical instability. Please use bfloat16 or float32 instead.", "gemma3": "Numerical instability. Please use bfloat16 or float32 instead.", "gemma3_text": "Numerical instability. Please use bfloat16 or float32 instead.", - "plamo2": "Numerical instability. Please use bfloat16 or float32 instead.", "glm4": "Numerical instability. Please use bfloat16 or float32 instead.", } diff --git a/vllm/model_executor/models/plamo2.py b/vllm/model_executor/models/plamo2.py deleted file mode 100644 index 5fd925cf0be..00000000000 --- a/vllm/model_executor/models/plamo2.py +++ /dev/null @@ -1,992 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Inference-only PLaMo2 model.""" - -from collections.abc import Iterable -from itertools import islice -from typing import TYPE_CHECKING - -import torch -from torch import nn -from transformers import PretrainedConfig - -from vllm.compilation.decorators import support_torch_compile -from vllm.config import VllmConfig, get_current_vllm_config -from vllm.distributed import divide, get_tensor_model_parallel_world_size -from vllm.distributed.parallel_state import get_pp_group -from vllm.forward_context import ForwardContext, get_forward_context -from vllm.model_executor.custom_op import PluggableLayer -from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import ( - ColumnParallelLinear, - MergedColumnParallelLinear, - QKVParallelLinear, - RowParallelLinear, -) -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.mamba.abstract import MambaBase -from vllm.model_executor.layers.mamba.mamba_utils import ( - MambaStateCopyFunc, - MambaStateCopyFuncCalculator, - MambaStateDtypeCalculator, - MambaStateShapeCalculator, - is_conv_state_dim_first, -) -from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( - causal_conv1d_fn, - causal_conv1d_update, -) -from vllm.model_executor.layers.mamba.ops.ssd_combined import ( - mamba_chunk_scan_combined_varlen, -) -from vllm.model_executor.layers.mamba.ops.ssu_dispatch import selective_state_update -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 ( - ParallelLMHead, - VocabParallelEmbedding, -) -from vllm.model_executor.model_loader.weight_utils import ( - composed_weight_loader, - default_weight_loader, - sharded_weight_loader, -) -from vllm.model_executor.models.interfaces import ( - HasInnerState, - IsHybrid, - SupportsLoRA, - SupportsPP, -) -from vllm.model_executor.models.utils import ( - AutoWeightsLoader, - is_pp_missing_parameter, - make_empty_intermediate_tensors_factory, - make_layers, - maybe_prefix, -) -from vllm.model_executor.utils import set_weight_attrs -from vllm.platforms import current_platform -from vllm.sequence import IntermediateTensors -from vllm.utils.torch_utils import direct_register_custom_op -from vllm.v1.attention.backend import AttentionMetadata -from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata -from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum - -# Only used for type hinting. -if TYPE_CHECKING: - - class Plamo2Config(PretrainedConfig): # type: ignore - model_type: str = "plamo2" - - hidden_size: int - num_hidden_layers: int - rms_norm_eps: float - # Attention - num_attention_heads: int - hidden_size_per_head: int - num_key_value_heads: int - # Mamba - mamba_d_state: int - mamba_d_conv: int - mamba_num_heads: int - mamba_step: int - # MLP - intermediate_size: int - # Tokenizer - vocab_size: int - - -def is_mamba(config: "Plamo2Config", i: int) -> bool: - assert config.mamba_step > 1 - - if config.num_hidden_layers <= (config.mamba_step // 2): - # use attention in last layer - return i != config.num_hidden_layers - 1 - return (i % config.mamba_step) != (config.mamba_step // 2) - - -# Adapted from: -# vllm.model_executor.layers.mamba.mamba_mixer2.MambaMixer2 -# transformers.models.mamba.modeling_mamba.MambaMixer -# --8<-- [start:plamo2_mamba_mixer] -@PluggableLayer.register("plamo2_mamba_mixer") -class Plamo2MambaMixer(MambaBase, PluggableLayer): - # --8<-- [end:plamo2_mamba_mixer] - - def __init__(self, vllm_config: VllmConfig, *, prefix: str = "", **kwargs) -> None: - super().__init__() - self.config = vllm_config.model_config.hf_config - self.cache_config = vllm_config.cache_config - self.model_config = vllm_config.model_config - self.quant_config = vllm_config.quant_config - self.is_lora_enabled = bool(vllm_config.lora_config) - self.hidden_size = self.config.hidden_size - self.ssm_state_size = self.config.mamba_d_state - self.conv_kernel_size = self.config.mamba_d_conv - self.intermediate_size = ( - self.config.mamba_num_heads * self.config.hidden_size_per_head - ) - self.tp_size = get_tensor_model_parallel_world_size() - self.head_dim = self.config.hidden_size_per_head - self.num_heads = self.config.mamba_num_heads - self.time_step_rank = max(64, self.hidden_size // 16) - self.conv1d = ColumnParallelLinear( - input_size=self.conv_kernel_size, - output_size=self.intermediate_size, - bias=False, - prefix=f"{prefix}.conv1d", - return_bias=False, - ) - # unsqueeze to fit conv1d weights shape into the linear weights shape. - # Can't do this in `weight_loader` since it already exists in - # `ColumnParallelLinear` and `set_weight_attrs` - # doesn't allow to override it - self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1) - - self.in_proj = MergedColumnParallelLinear( - self.hidden_size, - [self.intermediate_size] * 2, - bias=False, - quant_config=self.quant_config, - prefix=f"{prefix}.in_proj", - return_bias=False, - ) - # selective projection used to make dt, B and C input dependent - self.bcdt_proj = RowParallelLinear( - self.intermediate_size, - self.time_step_rank + self.ssm_state_size * 2, - bias=False, - quant_config=self.quant_config, - prefix=f"{prefix}.bcdt_proj", - return_bias=False, - ) - # time step projection (discretization) - - # In the forward we need to apply dt_proj without the bias, - # as the bias is added in the selective scan kernel. - self.dt_proj = ColumnParallelLinear( - self.time_step_rank, - self.num_heads, - bias=False, - quant_config=self.quant_config, - prefix=f"{prefix}.dt_proj", - return_bias=False, - ) - - self.A = nn.Parameter( - torch.empty( - divide(self.num_heads, self.tp_size), - dtype=torch.float32, - ) - ) - self.D = nn.Parameter(torch.ones(divide(self.num_heads, self.tp_size))) - self.dt_bias = nn.Parameter(torch.ones(divide(self.num_heads, self.tp_size))) - - set_weight_attrs(self.D, {"weight_loader": sharded_weight_loader(0)}) - a_weight_loader = composed_weight_loader( - sharded_weight_loader(0), lambda x: -torch.exp(x.float()) - ) - set_weight_attrs(self.A, {"weight_loader": a_weight_loader}) - set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)}) - - self.out_proj = RowParallelLinear( - self.intermediate_size, - self.hidden_size, - bias=False, - input_is_parallel=True, - quant_config=self.quant_config, - prefix=f"{prefix}.out_proj", - return_bias=False, - ) - # The activation function is fixed to SiLU. - self.activation = "silu" - - self.dt_norm = RMSNorm(self.time_step_rank, eps=self.config.rms_norm_eps) - self.B_norm = RMSNorm(self.ssm_state_size, eps=self.config.rms_norm_eps) - self.C_norm = RMSNorm(self.ssm_state_size, eps=self.config.rms_norm_eps) - - self.chunk_size = self.config.mamba_chunk_size - - compilation_config = get_current_vllm_config().compilation_config - if prefix in compilation_config.static_forward_context: - raise ValueError(f"Duplicate layer name: {prefix}") - compilation_config.static_forward_context[prefix] = self - # The tuple is (conv_state, ssm_state) - self.kv_cache = (torch.tensor([]), torch.tensor([])) - assert self.chunk_size != -1, "chunk_size must be set for v1" - - self.prefix = prefix - - def _project_ssm_parameters(self, hidden_states): - if self.is_lora_enabled: - # Lora kernel requires contiguous tensor. - ssm_parameters = self.bcdt_proj(hidden_states.contiguous()) - else: - ssm_parameters = self.bcdt_proj(hidden_states) - B, C, time_step = torch.split( - ssm_parameters, - [self.ssm_state_size, self.ssm_state_size, self.time_step_rank], - dim=-1, - ) - - # vllm._custom_ops.rms_norm requires contiguous input tensors. - time_step = self.dt_norm(time_step.contiguous()) - B = self.B_norm(B.contiguous()) - C = self.C_norm(C.contiguous()) - dt = self.dt_proj(time_step) - return B, C, dt - - def forward( - self, - hidden_states: torch.Tensor, - output: torch.Tensor, - **kwargs, - ): - torch.ops.vllm.plamo2_mamba_mixer( - hidden_states, - output, - self.prefix, - ) - - def forward_impl( - self, - hidden_states: torch.Tensor, - output: torch.Tensor, - **kwargs, - ): - forward_context = get_forward_context() - # attn_metadata contains metadata necessary for the mamba2 triton - # kernels to operate in continuous batching and in chunked prefill - # modes; they are computed at top-level model forward since they - # stay the same and reused for all mamba layers in the same iteration - attn_metadata: AttentionMetadata = forward_context.attn_metadata - - if attn_metadata is not None: - assert isinstance(attn_metadata, dict) - attn_metadata = attn_metadata[self.prefix] - assert isinstance(attn_metadata, Mamba2AttentionMetadata) - self_kv_cache = self.kv_cache - # conv_state = (..., dim, width-1) yet contiguous along 'dim' - # conv_state must be (..., dim, width-1) for the conv kernels. - # DS layout stores it that way directly; SD layout needs a transpose. - conv_state = ( - self_kv_cache[0] - if is_conv_state_dim_first() - else self_kv_cache[0].transpose(-1, -2) - ) - ssm_state = self_kv_cache[1] - state_indices_tensor_p = attn_metadata.state_indices_tensor_p - state_indices_tensor_d = attn_metadata.state_indices_tensor_d - has_initial_states_p = attn_metadata.has_initial_states_p - prep_initial_states = attn_metadata.prep_initial_states - chunk_size = attn_metadata.chunk_size - seq_idx_p = attn_metadata.seq_idx_p - query_start_loc_p = attn_metadata.query_start_loc_p - cu_chunk_seqlen_p = attn_metadata.cu_chunk_seqlen_p - last_chunk_indices_p = attn_metadata.last_chunk_indices_p - - # 1. Gated MLP's linear projection - projected_states = self.in_proj(hidden_states) - gate, hidden_states = projected_states.chunk(2, dim=-1) - - # 2. Convolution sequence transformation - conv_weights = self.conv1d.weight.view( - self.conv1d.weight.size(0), self.conv1d.weight.size(2) - ) - - if attn_metadata is None: - # profile run - hidden_states = ( - hidden_states.transpose(0, 1).clone().transpose(0, 1) - ).contiguous() - output[:] = self.out_proj(hidden_states) - return - - num_prefills = attn_metadata.num_prefills # request count - num_decodes = attn_metadata.num_decode_tokens # token count (=request) - num_prefill_tokens = attn_metadata.num_prefill_tokens # token count - has_prefill = num_prefills > 0 - has_decode = num_decodes > 0 - num_actual_tokens = num_prefill_tokens + num_decodes - - # Separate prefill and decode by splitting varlen input - # Split along token dimension - hidden_states_d, hidden_states_p = torch.split( - hidden_states[:num_actual_tokens], - [num_decodes, num_prefill_tokens], - dim=0, - ) - gate_d, gate_p = torch.split( - gate[:num_actual_tokens], [num_decodes, num_prefill_tokens], dim=0 - ) - # Preallocate output tensor to avoid memcpy cost for merging prefill - # and decode outputs - preallocated_ssm_out = torch.empty( - [ - num_prefill_tokens + num_decodes, - (self.num_heads // self.tp_size) * self.head_dim, - ], - dtype=hidden_states.dtype, - device=hidden_states.device, - ) - preallocated_ssm_out_d, preallocated_ssm_out_p = torch.split( - preallocated_ssm_out, - [num_decodes, num_prefill_tokens], - dim=0, - ) - - # Process prefill requests - if has_prefill: - # 2. Convolution sequence transformation - # - "cache_indices" updates the conv_state cache in positions - # pointed to by "state_indices_tensor_p" - x = hidden_states_p.transpose(0, 1) # this is the form that causal-conv see - hidden_states_p = causal_conv1d_fn( - x, - conv_weights, - self.conv1d.bias, - activation=self.activation, - conv_states=conv_state, - has_initial_state=has_initial_states_p, - cache_indices=state_indices_tensor_p, - metadata=attn_metadata, - query_start_loc=query_start_loc_p, - ) - hidden_states_p = hidden_states_p.transpose(0, 1) - hidden_states_p = hidden_states_p[:num_prefill_tokens] - # In some instances, the following `bcdt_proj` op - # requires contiguous inputs - # (e.g. if the Marlin kernel is used). - hidden_states_p = hidden_states_p.contiguous() - - B, C, dt = self._project_ssm_parameters(hidden_states_p) - - # 3. State Space Model sequence transformation - initial_states = None - if has_initial_states_p is not None and prep_initial_states: - # making a copy of the states - initial_states = torch.where( - has_initial_states_p[:, None, None, None], - ssm_state[state_indices_tensor_p], - 0, - ) - - varlen_state = mamba_chunk_scan_combined_varlen( - hidden_states_p.view( - num_prefill_tokens, self.num_heads // self.tp_size, self.head_dim - ), - dt, - self.A, - B.view(num_prefill_tokens, 1, -1), - C.view(num_prefill_tokens, 1, -1), - chunk_size=chunk_size, - D=self.D, - z=gate_p.view( - num_prefill_tokens, self.num_heads // self.tp_size, self.head_dim - ), - dt_bias=self.dt_bias, - seq_idx=seq_idx_p, - cu_seqlens=query_start_loc_p, - cu_chunk_seqlens=cu_chunk_seqlen_p, - last_chunk_indices=last_chunk_indices_p, - initial_states=initial_states, - dt_softplus=True, - dt_limit=(0.0, float("inf")), - out=preallocated_ssm_out_p.view(num_prefill_tokens, -1, self.head_dim), - state_dtype=ssm_state.dtype, - ) - - # update ssm states - # - varlen state is a (batch, nheads, headdim, dstate) tensor - ssm_state[state_indices_tensor_p] = varlen_state - - # Process decode requests - if has_decode: - # 2. Convolution sequence transformation - hidden_states_d = causal_conv1d_update( - hidden_states_d, - conv_state, - conv_weights, - self.conv1d.bias, - self.activation, - conv_state_indices=state_indices_tensor_d, - ) - - # ROCm: Ensure contiguous tensor for bcdt_proj linear layer. - # causal_conv1d_update returns a non-contiguous view (stride 8192 - # instead of 4096 for shape [batch, 4096]), causing incorrect GEMM - # results when batch > 1 on ROCm. - if current_platform.is_rocm(): - hidden_states_d = hidden_states_d.contiguous() - - B, C, dt = self._project_ssm_parameters(hidden_states_d) - - # 3. State Space Model sequence transformation - A = self.A[:, None, ...][:, :, None].expand( - -1, self.head_dim, self.config.mamba_d_state - ) - dt = dt[:, :, None].expand(-1, -1, self.head_dim) - dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim) - D = self.D[:, None, ...].expand(-1, self.head_dim) - B = B.unsqueeze(1) - C = C.unsqueeze(1) - hidden_states_d = hidden_states_d.view( - -1, self.num_heads // self.tp_size, self.head_dim - ) - - # - the hidden is reshaped into (bs, num_heads, head_dim) - # - ssm_state's slots will be selected - # using state_indices_tensor_d - - # NOTE: final output is an in-place update of out tensor - selective_state_update( - ssm_state, - hidden_states_d, - dt, - A, - B, - C, - D, - dt_bias, - z=gate_d.reshape(num_decodes, -1, self.head_dim), - dt_softplus=True, - state_batch_indices=state_indices_tensor_d, - out=preallocated_ssm_out_d.view(num_decodes, -1, self.head_dim), - ) - - # 4. Final linear projection - output[:num_actual_tokens] = self.out_proj(preallocated_ssm_out) - - def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]: - assert self.model_config is not None - assert self.cache_config is not None - return MambaStateDtypeCalculator.mamba2_state_dtype( - self.model_config.dtype, - self.cache_config.mamba_cache_dtype, - self.cache_config.mamba_ssm_cache_dtype, - ) - - def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]: - return MambaStateShapeCalculator.mamba2_state_shape( - intermediate_size=self.intermediate_size, - tp_world_size=get_tensor_model_parallel_world_size(), - n_groups=0, - num_heads=self.num_heads, - head_dim=self.head_dim, - state_size=self.ssm_state_size, - conv_kernel=self.conv_kernel_size, - ) - - @property - def mamba_type(self) -> MambaAttentionBackendEnum: - return MambaAttentionBackendEnum.MAMBA2 - - -def plamo2_mamba_mixer( - hidden_states: torch.Tensor, - output: torch.Tensor, - layer_name: str, -) -> None: - forward_context: ForwardContext = get_forward_context() - self = forward_context.no_compile_layers[layer_name] - self.forward_impl(hidden_states=hidden_states, output=output) - - -def plamo2_mamba_mixer_fake( - hidden_states: torch.Tensor, - output: torch.Tensor, - layer_name: str, -) -> None: - return - - -direct_register_custom_op( - op_name="plamo2_mamba_mixer", - op_func=plamo2_mamba_mixer, - mutates_args=["output"], - fake_impl=plamo2_mamba_mixer_fake, -) - - -class DenseMLP(nn.Module): - def __init__( - self, - config: "Plamo2Config", - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__() - self.hidden_size = config.hidden_size - self.intermediate_size = config.intermediate_size - self.gate_up_proj = MergedColumnParallelLinear( - self.hidden_size, - [self.intermediate_size] * 2, - bias=False, - prefix=f"{prefix}.gate_up_proj", - quant_config=quant_config, - return_bias=False, - ) - self.act = SiluAndMul() - self.down_proj = RowParallelLinear( - self.intermediate_size, - self.hidden_size, - bias=False, - prefix=f"{prefix}.down_proj", - quant_config=quant_config, - return_bias=False, - ) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - h = self.gate_up_proj(hidden_states) - h = self.act(h) - return self.down_proj(h) - - -class Plamo2AttentionMixer(nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs) -> None: - super().__init__() - config = vllm_config.model_config.hf_config - cache_config = vllm_config.cache_config - quant_config = vllm_config.quant_config - self.hidden_size = config.hidden_size - tp_size = get_tensor_model_parallel_world_size() - self.total_num_heads = config.num_attention_heads - assert self.total_num_heads % tp_size == 0 - self.num_heads = self.total_num_heads // tp_size - self.total_num_kv_heads = config.num_key_value_heads - if self.total_num_kv_heads >= tp_size: - # Number of KV heads is greater than TP size, so we partition - # the KV heads across multiple tensor parallel GPUs. - assert self.total_num_kv_heads % tp_size == 0 - else: - # Number of KV heads is less than TP size, so we replicate - # the KV heads across multiple tensor parallel GPUs. - 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 = config.hidden_size_per_head - 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.qkv_proj = QKVParallelLinear( - config.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, - config.hidden_size, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.o_proj", - ) - - max_position = config.max_position_embeddings - if hasattr(vllm_config.model_config, "max_model_len") and isinstance( - vllm_config.model_config.max_model_len, int - ): - max_position = min(max_position, vllm_config.model_config.max_model_len) - - self.rotary_emb = get_rope( - self.head_dim, - max_position=max_position, - rope_parameters=config.rope_parameters, - ) - self.q_norm = RMSNorm(config.hidden_size_per_head, eps=config.rms_norm_eps) - self.q_norm.weight = torch.nn.Parameter( - torch.ones((self.num_heads, config.hidden_size_per_head)) - ) - set_weight_attrs( - self.q_norm.weight, {"weight_loader": sharded_weight_loader(0)} - ) - self.k_norm = RMSNorm(config.hidden_size_per_head, eps=config.rms_norm_eps) - self.k_norm.weight = torch.nn.Parameter( - torch.ones((self.num_kv_heads, config.hidden_size_per_head)) - ) - # Tensor-parallelism shards the K norm weights to the tp ranks - # in a head-wise manner. This approach does not work if there is only - # a single KV head, as is the case for PLaMo 2-1B. - if self.total_num_kv_heads != 1: - set_weight_attrs( - self.k_norm.weight, {"weight_loader": sharded_weight_loader(0)} - ) - - self.attn = Attention( - self.num_heads, - self.head_dim, - self.scaling, - num_kv_heads=self.num_kv_heads, - cache_config=cache_config, - prefix=f"{prefix}.attn", - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - **kwargs, - ) -> torch.Tensor: - qkv, _ = self.qkv_proj(hidden_states) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - - q_shape = q.shape - q = q.reshape(q_shape[:-1] + self.q_norm.weight.shape) - q = self.q_norm.forward_native(q).reshape(q_shape) - k_shape = k.shape - k = k.reshape(k_shape[:-1] + self.k_norm.weight.shape) - k = self.k_norm.forward_native(k).reshape(k_shape) - - q, k = self.rotary_emb(positions, q, k) - attn_output = self.attn(q, k, v) - output, _ = self.o_proj(attn_output) - return output - - -class Plamo2DecoderLayer(nn.Module): - def __init__( - self, vllm_config: VllmConfig, layer_idx: int, prefix: str = "", **kwargs - ) -> None: - super().__init__() - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - - self.is_mamba = is_mamba(config, layer_idx) - if self.is_mamba: - self.mixer = Plamo2MambaMixer( - vllm_config=vllm_config, prefix=f"{prefix}.mixer" - ) - else: - self.mixer = Plamo2AttentionMixer( - vllm_config=vllm_config, prefix=f"{prefix}.mixer" - ) - - self.mlp = DenseMLP( - config=config, quant_config=quant_config, prefix=f"{prefix}.mlp" - ) - self.pre_mixer_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_mixer_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.pre_mlp_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_mlp_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - residual: torch.Tensor | None, - **kwargs, - ): - if residual is None: - residual = hidden_states - hidden_states = self.pre_mixer_norm(hidden_states) - else: - hidden_states, residual = self.pre_mixer_norm(hidden_states, residual) - - if self.is_mamba: - # Plamo2MambaMixer writes output to this tensor - output = torch.empty_like(hidden_states) - mixer_kwargs = { - "output": output, - } - else: - mixer_kwargs = { - "positions": positions, - } - hidden_states = self.mixer( - hidden_states=hidden_states, - **mixer_kwargs, - ) - if self.is_mamba: - hidden_states = output - hidden_states = self.post_mixer_norm(hidden_states) - # Fully Connected - hidden_states, residual = self.pre_mlp_norm(hidden_states, residual) - hidden_states = self.mlp(hidden_states) - hidden_states = self.post_mlp_norm(hidden_states) - return hidden_states, residual - - -class Plamo2Decoder(torch.nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: - super().__init__() - config = vllm_config.model_config.hf_config - extra_kwargs = {"is_lora_enabled": bool(vllm_config.lora_config)} - - def get_layer(prefix: str): - layer_idx = int(prefix.rsplit(".", 1)[1]) - return Plamo2DecoderLayer( - vllm_config=vllm_config, - layer_idx=layer_idx, - prefix=prefix, - **extra_kwargs, - ) - - self.start_layer, self.end_layer, self.layers = make_layers( - config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers" - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - residual: torch.Tensor | None, - ) -> torch.Tensor: - for layer in islice(self.layers, self.start_layer, self.end_layer): - hidden_states, residual = layer( - positions=positions, - hidden_states=hidden_states, - residual=residual, - ) - return hidden_states, residual - - -@support_torch_compile -class Plamo2Model(torch.nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - - config = vllm_config.model_config.hf_config - - self.config = config - self.vocab_size = config.vocab_size - - self.embed_tokens = VocabParallelEmbedding( - self.vocab_size, - config.hidden_size, - prefix=f"{prefix}.embed_tokens", - ) - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( - ["hidden_states", "residual"], config.hidden_size - ) - self.layers = Plamo2Decoder(vllm_config=vllm_config, prefix=f"{prefix}.layers") - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - - 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 = None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor: - if get_pp_group().is_first_rank: - if inputs_embeds is not None: - hidden_states = inputs_embeds - else: - hidden_states = self.embed_input_ids(input_ids) - residual = None - else: - assert intermediate_tensors is not None - hidden_states = intermediate_tensors["hidden_states"] - residual = intermediate_tensors["residual"] - - hidden_states, residual = self.layers( - positions=positions, - hidden_states=hidden_states, - residual=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 - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - # Update the weight names to be compatible with the vllm version - # of the model. - # Do not change the order of the replacements. - replacements = { - # Rename incompatible weight names. - ".A_log": ".A", - ".B_norm_weight": ".B_norm.weight", - ".C_norm_weight": ".C_norm.weight", - ".dt_norm_weight": ".dt_norm.weight", - ".q_weight": ".q_norm.weight", - ".k_weight": ".k_norm.weight", - } - # Apply replacements based on the defined mappings - for old, new in replacements.items(): - if old in name: - name = name.replace(old, new) - - # Reshape the in_proj weights to match the shape expected - # by MergedColumnParallelLinear. - # This works both for unquantized weights and - # for quantized weights. - # In the quantized case, the weights are already transposed. - # Also, in addition to the quantized weights, - # the zero points and scales have to be reshaped as well. - # Packing should not be affected by this. - if ( - ".mixer.in_proj.weight" in name - or "mixer.in_proj.qweight" in name - or "mixer.in_proj.scales" in name - or "mixer.in_proj.qzeros" in name - ): - if "mixer.in_proj.weight" in name: - loaded_weight = loaded_weight.transpose(0, 1) - # for weight: - # loaded_weight.shape[0] == self.config.hidden_size - # for qweight: - # loaded_weight.shape[0] == self.config.hidden_size // param.pack_factor # noqa - # for scales and qzeros: - # loaded_weight.shape[0] == self.config.hidden_size // self.vllm_config.quant_config.group_size # noqa - loaded_weight = loaded_weight.reshape( - loaded_weight.shape[0], self.config.mamba_num_heads, -1 - ) - gate_weight, hidden_states_weight = loaded_weight.chunk(2, dim=-1) - gate_weight = gate_weight.reshape(loaded_weight.shape[0], -1) - hidden_states_weight = hidden_states_weight.reshape( - loaded_weight.shape[0], -1 - ) - loaded_weight = torch.cat([gate_weight, hidden_states_weight], dim=-1) - if "mixer.in_proj.weight" in name: - loaded_weight = loaded_weight.transpose(0, 1) - - # Offset parameter with vllm's RMSNorm haven't been supported yet. - if ".pre_mixer_norm" in name: - loaded_weight += 1.0 - elif ".post_mixer_norm" in name: - loaded_weight += 1.0 / 5 - elif ".pre_mlp_norm" in name: - loaded_weight += 1.0 - elif ".post_mlp_norm" in name: - loaded_weight += 1.0 / (5**1.5) - elif name == "norm.weight": - loaded_weight += 1.0 - - # Skip layers on other devices. - if is_pp_missing_parameter(name, self): - 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 - - -class Plamo2ForCausalLM( - torch.nn.Module, HasInnerState, SupportsLoRA, SupportsPP, IsHybrid -): - packed_modules_mapping = { - "qkv_proj": ["qkv_proj"], - "gate_up_proj": ["gate_up_proj"], - "in_proj": ["in_proj"], - } - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: - super().__init__() - config = vllm_config.model_config.hf_config - scheduler_config = vllm_config.scheduler_config - - self.config = config - self.vllm_config = vllm_config - self.model_config = vllm_config.model_config - self.scheduler_config = scheduler_config - - # ModelConfig.get_head_size assumes head_dim is set or calculated as - # hidden_size // num_attention_heads. However, this is not always - # the case for PLaMo2, as indicated by the FIXME comment. - self.config.head_dim = self.config.hidden_size_per_head - - self.model = Plamo2Model( - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") - ) - self.vocab_size = self.config.vocab_size - self.lm_head = ParallelLMHead( - self.vocab_size, - self.config.hidden_size, - prefix=f"{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.config.vocab_size - ) - 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.embed_input_ids(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - **kwargs, - ): - hidden_states = self.model( - input_ids, positions, intermediate_tensors, inputs_embeds - ) - return hidden_states - - @classmethod - def get_mamba_state_dtype_from_config( - cls, - vllm_config: "VllmConfig", - ) -> tuple[torch.dtype, torch.dtype]: - return MambaStateDtypeCalculator.mamba2_state_dtype( - vllm_config.model_config.dtype, - vllm_config.cache_config.mamba_cache_dtype, - vllm_config.cache_config.mamba_ssm_cache_dtype, - ) - - @classmethod - def get_mamba_state_shape_from_config( - cls, - vllm_config: "VllmConfig", - ) -> tuple[tuple[int, int], tuple[int, int, int]]: - """Calculate shapes for Mamba's convolutional and state caches. - Args: - vllm_config: vLLM config - Returns: - Tuple containing: - - conv_state_shape: Shape for convolutional state cache - - temporal_state_shape: Shape for state space model cache - """ - parallel_config = vllm_config.parallel_config - hf_config = vllm_config.model_config.hf_config - intermediate_size = hf_config.mamba_num_heads * hf_config.hidden_size_per_head - - return MambaStateShapeCalculator.mamba2_state_shape( - intermediate_size=intermediate_size, - tp_world_size=parallel_config.tensor_parallel_size, - n_groups=0, - num_heads=hf_config.mamba_num_heads, - head_dim=hf_config.hidden_size_per_head, - state_size=hf_config.mamba_d_state, - conv_kernel=hf_config.mamba_d_conv, - ) - - @classmethod - def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: - return MambaStateCopyFuncCalculator.mamba2_state_copy_func() - - def compute_logits( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor | None: - logits = self.logits_processor(self.lm_head, hidden_states) - return logits - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) - return loader.load_weights(weights) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index a4fd44fc2cd..947e0d043dd 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -191,7 +191,6 @@ _TEXT_GENERATION_MODELS = { "PhiForCausalLM": ("phi", "PhiForCausalLM"), "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"), "PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"), - "Plamo2ForCausalLM": ("plamo2", "Plamo2ForCausalLM"), "Plamo3ForCausalLM": ("plamo3", "Plamo3ForCausalLM"), "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"), "Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"), @@ -761,6 +760,7 @@ _PREVIOUSLY_SUPPORTED_MODELS = { "TeleChatForCausalLM": "0.25.0", "PersimmonForCausalLM": "0.25.0", "FuyuForCausalLM": "0.25.0", + "Plamo2ForCausalLM": "0.26.0", } _OOT_SUPPORTED_MODELS = {