From d682968aa9fcd7e7a78218b548c52fc198a87a6c Mon Sep 17 00:00:00 2001 From: Tiezhen WANG <38108242+xianbaoqian@users.noreply.github.com> Date: Thu, 18 Jun 2026 21:51:00 +0800 Subject: [PATCH] [Model] Remove BambaForCausalLM (#45990) Signed-off-by: Xianbao QIAN Co-authored-by: Claude --- docs/models/supported_models.md | 1 - .../models/language/generation/test_hybrid.py | 11 +- tests/models/registry.py | 4 - vllm/model_executor/models/bamba.py | 517 ------------------ vllm/model_executor/models/registry.py | 2 +- 5 files changed, 6 insertions(+), 529 deletions(-) delete mode 100644 vllm/model_executor/models/bamba.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 0826ec7d57..e67bc197d3 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -374,7 +374,6 @@ th { | `BailingMoeForCausalLM` | Ling | `inclusionAI/Ling-lite-1.5`, `inclusionAI/Ling-plus`, etc. | ✅︎ | ✅︎ | | `BailingMoeV2ForCausalLM` | Ling | `inclusionAI/Ling-mini-2.0`, etc. | ✅︎ | ✅︎ | | `BailingMoeV2_5ForCausalLM` | Ling | `inclusionAI/Ling-2.5-1T`, `inclusionAI/Ring-2.5-1T` | | ✅︎ | -| `BambaForCausalLM` | Bamba | `ibm-ai-platform/Bamba-9B-fp8`, `ibm-ai-platform/Bamba-9B` | ✅︎ | ✅︎ | | `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. | ✅︎ | ✅︎ | diff --git a/tests/models/language/generation/test_hybrid.py b/tests/models/language/generation/test_hybrid.py index cd89ca284d..0f19c1038e 100644 --- a/tests/models/language/generation/test_hybrid.py +++ b/tests/models/language/generation/test_hybrid.py @@ -37,7 +37,6 @@ HYBRID_MODELS = [ "ai21labs/Jamba-tiny-dev", "pfnet/plamo-2-1b", "Zyphra/Zamba2-1.2B-instruct", - "hmellor/tiny-random-BambaForCausalLM", "ibm-granite/granite-4.0-tiny-preview", "tiiuae/Falcon-H1-0.5B-Base", "LiquidAI/LFM2-1.2B", @@ -439,7 +438,7 @@ def _get_vLLM_output( return outs, vllm_model -@pytest.mark.parametrize("model", [HYBRID_MODELS[0], HYBRID_MODELS[3]]) +@pytest.mark.parametrize("model", [HYBRID_MODELS[0]]) @pytest.mark.parametrize("max_tokens", [64]) @pytest.mark.parametrize("n_repetitions", [2]) # If num_logprobs is set to -1, then the stringent version @@ -503,7 +502,7 @@ def test_apc_single_prompt( ) -@pytest.mark.parametrize("model", [HYBRID_MODELS[0], HYBRID_MODELS[3]]) +@pytest.mark.parametrize("model", [HYBRID_MODELS[0]]) @pytest.mark.parametrize("max_tokens", [64]) @pytest.mark.parametrize("n_repetitions", [2]) # If num_logprobs is set to -1, then the stringent version @@ -584,7 +583,7 @@ def test_apc_single_prompt_block_align_alignment( ) -@pytest.mark.parametrize("model", [HYBRID_MODELS[0], HYBRID_MODELS[3]]) +@pytest.mark.parametrize("model", [HYBRID_MODELS[0]]) @pytest.mark.parametrize("max_tokens", [64]) @pytest.mark.parametrize("n_repetitions", [2]) # If num_logprobs is set to -1, then the stringent version @@ -653,7 +652,7 @@ def test_apc_multiple_prompts_all_cached_outputs( ) -@pytest.mark.parametrize("model", [HYBRID_MODELS[0], HYBRID_MODELS[3]]) +@pytest.mark.parametrize("model", [HYBRID_MODELS[0]]) @pytest.mark.parametrize("max_tokens", [64]) @pytest.mark.parametrize("n_repetitions", [2]) # If num_logprobs is set to -1, then the stringent version @@ -738,7 +737,7 @@ def test_apc_multiple_prompts_block_align_alignment( ) -@pytest.mark.parametrize("model", [HYBRID_MODELS[0], HYBRID_MODELS[3]]) +@pytest.mark.parametrize("model", [HYBRID_MODELS[0]]) @pytest.mark.parametrize("max_tokens", [64]) @pytest.mark.parametrize("n_repetitions", [2]) # If num_logprobs is set to -1, then the stringent version diff --git a/tests/models/registry.py b/tests/models/registry.py index 29d46860d9..ec2c52db56 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -223,10 +223,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "BailingMoeV2_5ForCausalLM": _HfExamplesInfo( "inclusionAI/Ring-2.5-1T", trust_remote_code=True ), - "BambaForCausalLM": _HfExamplesInfo( - "ibm-ai-platform/Bamba-9B-v1", - extras={"tiny": "hmellor/tiny-random-BambaForCausalLM"}, - ), "BloomForCausalLM": _HfExamplesInfo( "bigscience/bloom-560m", {"1b": "bigscience/bloomz-1b1"} ), diff --git a/vllm/model_executor/models/bamba.py b/vllm/model_executor/models/bamba.py deleted file mode 100644 index d220b22dda..0000000000 --- a/vllm/model_executor/models/bamba.py +++ /dev/null @@ -1,517 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Inference-only Bamba model.""" - -# Added by the IBM Team, 2024 -from collections.abc import Iterable - -import torch -from torch import nn -from transformers import BambaConfig - -from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, ModelConfig, VllmConfig -from vllm.distributed import get_tensor_model_parallel_world_size -from vllm.distributed.parallel_state import get_pp_group -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 ( - MergedColumnParallelLinear, - QKVParallelLinear, - RowParallelLinear, -) -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2 -from vllm.model_executor.layers.mamba.mamba_utils import ( - MambaStateCopyFunc, - MambaStateCopyFuncCalculator, - MambaStateDtypeCalculator, - MambaStateShapeCalculator, -) -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 default_weight_loader -from vllm.sequence import IntermediateTensors - -from .interfaces import ( - HasInnerState, - IsHybrid, - SupportsLoRA, - SupportsMambaPrefixCaching, - SupportsPP, - SupportsQuant, -) -from .utils import ( - AutoWeightsLoader, - is_pp_missing_parameter, - make_empty_intermediate_tensors_factory, - make_layers, - maybe_prefix, -) - - -class BambaMLP(nn.Module): - def __init__( - self, - config: BambaConfig, - quant_config: QuantizationConfig | None = None, - bias: bool = False, - prefix: str = "", - ) -> None: - super().__init__() - self.gate_up_proj = MergedColumnParallelLinear( - input_size=config.hidden_size, - output_sizes=[config.intermediate_size] * 2, - bias=bias, - quant_config=quant_config, - prefix=f"{prefix}.gate_up_proj", - ) - self.down_proj = RowParallelLinear( - input_size=config.intermediate_size, - output_size=config.hidden_size, - bias=bias, - quant_config=quant_config, - prefix=f"{prefix}.down_proj", - ) - if config.hidden_act != "silu": - raise ValueError( - f"Unsupported activation: {config.hidden_act}. " - "Only silu is supported for now." - ) - self.act_fn = SiluAndMul() - - def forward(self, x): - x, _ = self.gate_up_proj(x) - x = self.act_fn(x) - x, _ = self.down_proj(x) - return x - - -class BambaMixerDecoderLayer(nn.Module): - def __init__( - self, - config: BambaConfig, - layer_idx: int, - model_config: ModelConfig | None = None, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__() - self.config = config - self.mamba = MambaMixer2( - hidden_size=config.hidden_size, - ssm_state_size=config.mamba_d_state, - conv_kernel_size=config.mamba_d_conv, - intermediate_size=config.mamba_expand * config.hidden_size, - use_conv_bias=config.mamba_conv_bias, - use_bias=config.mamba_proj_bias, - n_groups=config.mamba_n_groups, - num_heads=config.mamba_n_heads, - head_dim=config.mamba_d_head, - rms_norm_eps=config.rms_norm_eps, - activation=config.hidden_act, - model_config=model_config, - cache_config=cache_config, - quant_config=quant_config, - prefix=f"{prefix}.mixer", - ) - - self.feed_forward = BambaMLP( - config, quant_config=quant_config, prefix=f"{prefix}.feed_forward" - ) - self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.pre_ff_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - residual: torch.Tensor | None, - **kwargs, - ): - if residual is None: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - else: - hidden_states, residual = self.input_layernorm(hidden_states, residual) - - output = self.mamba(hidden_states) - # Fully Connected - hidden_states, residual = self.pre_ff_layernorm(output, residual) - hidden_states = self.feed_forward(hidden_states) - return hidden_states, residual - - -class BambaAttentionDecoderLayer(nn.Module): - def __init__( - self, - config: BambaConfig, - layer_idx: int, - model_config: ModelConfig | None = None, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__() - max_position_embeddings = getattr(config, "max_position_embeddings", 8192) - 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 // 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 - - rotary_dim = getattr(config, "attn_rotary_emb", self.head_dim) - config.rope_parameters["partial_rotary_factor"] = rotary_dim / self.head_dim - - self.rotary_emb = get_rope( - head_size=self.head_dim, - max_position=max_position_embeddings, - rope_parameters=config.rope_parameters, - is_neox_style=True, - dtype=torch.get_default_dtype(), # see impl of get_rope - ) - - 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", - ) - - 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", - ) - - self.feed_forward = BambaMLP( - config, quant_config=quant_config, prefix=f"{prefix}.feed_forward" - ) - self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.pre_ff_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - - def self_attention( - 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, k = self.rotary_emb(positions, q, k) - attn_output = self.attn(q, k, v) - output, _ = self.o_proj(attn_output) - return output - - 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.input_layernorm(hidden_states) - else: - hidden_states, residual = self.input_layernorm(hidden_states, residual) - - hidden_states = self.self_attention( - positions=positions, - hidden_states=hidden_states, - ) - # Fully Connected - hidden_states, residual = self.pre_ff_layernorm(hidden_states, residual) - hidden_states = self.feed_forward(hidden_states) - return hidden_states, residual - - -ALL_DECODER_LAYER_TYPES = { - "attention": BambaAttentionDecoderLayer, - "mamba": BambaMixerDecoderLayer, -} - - -@support_torch_compile -class BambaModel(nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - - config: BambaConfig = vllm_config.model_config.hf_config - model_config = vllm_config.model_config - cache_config = vllm_config.cache_config - quant_config = vllm_config.quant_config - - self.config = config - - self.vocab_size = config.vocab_size - - self.embed_tokens = VocabParallelEmbedding( - self.vocab_size, - config.hidden_size, - ) - - def get_layer(prefix: str): - layer_idx = int(prefix.rsplit(".", 1)[1]) - layer_class = ALL_DECODER_LAYER_TYPES[config.layers_block_type[layer_idx]] - return layer_class( - config, - layer_idx, - model_config, - cache_config, - quant_config=quant_config, - prefix=prefix, - ) - - self.start_layer, self.end_layer, self.layers = make_layers( - config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers" - ) - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( - ["hidden_states", "residual"], config.hidden_size - ) - - self.final_layernorm = 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"] - - residual = None - for i, layer in enumerate(self.layers): - hidden_states, residual = layer( - 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.final_layernorm(hidden_states, residual) - return hidden_states - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("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), - ] - - 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 "A_log" in name: - name = name.replace("A_log", "A") - - if ".self_attn." in name: - name = name.replace(".self_attn", "") - - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - # Skip layers on other devices. - 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: - # Skip loading extra bias for GPTQ models. - 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 = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - - -class BambaForCausalLM( - nn.Module, - HasInnerState, - SupportsLoRA, - SupportsPP, - IsHybrid, - SupportsQuant, - SupportsMambaPrefixCaching, -): - packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - "gate_up_proj": ["up_proj", "down_proj"], - } - - # LoRA specific attributes - embedding_modules = { - "embed_tokens": "input_embeddings", - "lm_head": "output_embeddings", - } - - @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_expand * hf_config.hidden_size - - return MambaStateShapeCalculator.mamba2_state_shape( - intermediate_size=intermediate_size, - tp_world_size=parallel_config.tensor_parallel_size, - n_groups=hf_config.mamba_n_groups, - num_heads=hf_config.mamba_n_heads, - head_dim=hf_config.mamba_d_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 __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - config = vllm_config.model_config.hf_config - self.vllm_config = vllm_config - self.model_config = vllm_config.model_config - - scheduler_config = vllm_config.scheduler_config - self.quant_config = vllm_config.quant_config - - super().__init__() - self.config = config - self.scheduler_config = scheduler_config - self.model = BambaModel( - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") - ) - - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - prefix=maybe_prefix(prefix, "lm_head"), - ) - - self.logits_processor = LogitsProcessor(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 - - 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) - return loader.load_weights(weights) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 6d35b978c0..f6286439e6 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -84,7 +84,6 @@ _TEXT_GENERATION_MODELS = { "BailingMoeForCausalLM": ("bailing_moe", "BailingMoeForCausalLM"), "BailingMoeV2ForCausalLM": ("bailing_moe", "BailingMoeV2ForCausalLM"), "BailingMoeV2_5ForCausalLM": ("bailing_moe_linear", "BailingMoeV25ForCausalLM"), - "BambaForCausalLM": ("bamba", "BambaForCausalLM"), "BloomForCausalLM": ("bloom", "BloomForCausalLM"), "ChatGLMModel": ("chatglm", "ChatGLMForCausalLM"), "ChatGLMForConditionalGeneration": ("chatglm", "ChatGLMForCausalLM"), @@ -733,6 +732,7 @@ _PREVIOUSLY_SUPPORTED_MODELS = { "MllamaForConditionalGeneration": "0.10.2", "XverseForCausalLM": "0.23.0", "Dots1ForCausalLM": "0.23.0", + "BambaForCausalLM": "0.23.0", } _OOT_SUPPORTED_MODELS = {