add falcon mamba support

Author: eaidova

optimum/exporters/openvino/model_configs.py

@@ -77,6 +77,7 @@
     DeciLMModelPatcher,
     DeepseekPatcher,
     FalconModelPatcher,
+    FalconMambaPatcher,
     FluxTransfromerModelPatcher,
     Gemma2ModelPatcher,
     GptBigCodeModelPatcher,
@@ -3010,3 +3011,13 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs):
                 )
 
         return dummy_inputs
+
+
+@register_in_tasks_manager(
+    "falcon-mamba", *["text-generation", "text-generation-with-past"], library_name="transformers"
+)
+class FalconMambaOpenVINOConfig(MambaOpenVINOConfig):
+    def patch_model_for_export(
+        self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None
+    ):
+        return FalconMambaPatcher(self, model, model_kwargs)

optimum/exporters/openvino/model_patcher.py

@@ -4510,6 +4510,75 @@ def mamba_mixer_forward(
     return contextualized_states
 
 
+def falcon_mamba_mixer_forward(
+    self,
+    input_states,
+    cache_params=None,
+    cache_position: Optional[torch.LongTensor] = None,
+    attention_mask: Optional[torch.LongTensor] = None,
+):
+    from transformers.models.falcon_mamba.modeling_falcon_mamba import rms_forward
+
+    batch_size, seq_len, _ = input_states.shape
+    dtype = input_states.dtype
+    # 1. Gated MLP's linear projection
+    projected_states = self.in_proj(input_states).transpose(1, 2)  # [batch, 2 * intermediate_size, seq_len]
+    hidden_states, gate = projected_states.chunk(2, dim=1)
+
+    if attention_mask is not None:
+        hidden_states = hidden_states * attention_mask.unsqueeze(1)
+
+    # 2. Convolution sequence transformation
+    if cache_params is not None:
+        ssm_state = cache_params.ssm_states[self.layer_idx].clone()
+        ssm_state = ssm_state.to(hidden_states.device)
+        # use `cache_position.shape[0]` to check whether we are in prefill
+        # stage, it's equivalent to check `cache_position[0] == 0`, which
+        # breaks dynamo fullgraph constraints
+        hidden_states, conv_state = self.conv_sequence_transform(
+            hidden_states, cache_position, cache_params.conv_states[self.layer_idx]
+        )
+        cache_params.conv_states[self.layer_idx] = conv_state
+    else:
+        ssm_state = torch.zeros(
+            (batch_size, self.intermediate_size, self.ssm_state_size), device=hidden_states.device, dtype=dtype
+        )
+        hidden_states = self.act(self.conv1d(hidden_states)[..., :seq_len])  # [batch, intermediate_size, seq_len]
+
+    if attention_mask is not None:
+        hidden_states = hidden_states * attention_mask.unsqueeze(1)
+
+    # 3. State Space Model sequence transformation
+    # 3.a. Selection:  [batch, seq_len, self.time_step_rank + self.ssm_state_size * 2]
+    ssm_parameters = self.x_proj(hidden_states.transpose(1, 2))
+    time_step, B, C = torch.split(
+        ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1
+    )
+
+    B = rms_forward(B, variance_epsilon=self.rms_eps)
+    C = rms_forward(C, variance_epsilon=self.rms_eps)
+    time_step = rms_forward(time_step, variance_epsilon=self.rms_eps)
+    discrete_time_step = self.dt_proj(time_step)  # [batch, seq_len, intermediate_size]
+
+    discrete_time_step = torch.nn.functional.softplus(discrete_time_step)  # [batch, intermediate_size, seq_len]
+    A = -torch.exp(self.A_log.float())
+    B = B.float()
+    D = self.D.float()
+
+    scan_output, ssm_state = self.selective_scan(
+        ssm_state, hidden_states.float().transpose(1, 2), discrete_time_step, A, B, C, D
+    )
+    scan_output = scan_output.transpose(1, 2)
+    scan_output = scan_output * self.act(gate)
+
+    if cache_params is not None:
+        cache_params.ssm_states[self.layer_idx].copy_(ssm_state)
+
+    # 4. Final linear projection
+    contextualized_states = self.out_proj(scan_output.transpose(1, 2))  # [batch, seq_len, hidden_size]
+    return contextualized_states
+
+
 class MambaPatcher(ModelPatcher):
     def __init__(
         self,
@@ -4684,3 +4753,22 @@ def __exit__(self, exc_type, exc_value, traceback):
         self._model.forward = self._model.__orig_forward
         for layer in self._model.backbone.layers:
             layer.mixer.forward = layer.mixer._orig_forward
+
+
+class FalconMambaPatcher(MambaPatcher):
+    def __enter__(self):
+        super().__enter__()
+        selective_scan = SelectiveScan()
+
+        for layer in self._model.backbone.layers:
+            layer.mixer.selective_scan = selective_scan
+            layer.mixer._orig_forward = layer.mixer.forward
+            layer.mixer.forward = types.MethodType(falcon_mamba_mixer_forward, layer.mixer)
+            conv_transform = ConvSequenceTransform(
+                layer.mixer.conv_kernel_size,
+                layer.mixer.use_conv_bias,
+                layer.mixer.conv1d,
+                layer.mixer.act,
+                layer.mixer.conv1d.bias,
+            )
+            layer.mixer.conv_sequence_transform = torch.jit.script(conv_transform)

optimum/exporters/openvino/stateful.py

@@ -297,10 +297,13 @@ def patch_stateful_ssm(config, ov_model):
     build_state_initializer(ov_model, batch_dim)
 
 
+SSM_MODELS = ["mamba", "falcon-mamba"]
+
+
 def patch_stateful(config: PretrainedConfig, ov_model: ov.Model):
     if config.is_encoder_decoder and model_has_input_output_name(ov_model, "encoder_hidden_states"):
         return patch_stateful_encoder_decoder(config, ov_model)
-    if config.model_type == "mamba":
+    if config.model_type.replace("_", "-") in SSM_MODELS:
         return patch_stateful_ssm(config, ov_model)
     return patch_stateful_decoder(config, ov_model)
 

optimum/exporters/openvino/utils.py

@@ -229,6 +229,8 @@ def get_submodels(model):
     "qwen2-5-vl",
 ]
 
+SSM_MODELS = ["mamba", "falcon-mamba"]
+
 
 def save_config(config, save_dir):
     try:

optimum/intel/openvino/modeling_decoder.py

@@ -14,9 +14,9 @@
 import copy
 import logging
 import os
-from dataclasses import dataclass
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
+from dataclasses import dataclass
 
 import numpy as np
 import openvino
@@ -37,6 +37,7 @@
 
 from ...exporters.openvino import ensure_stateful_is_available, main_export, patch_stateful
 from ...exporters.openvino.stateful import model_has_state
+from ...exporters.openvino.utils import SSM_MODELS
 from ..utils.import_utils import compare_versions, is_nncf_available, is_transformers_version
 from ..utils.modeling_utils import MULTI_QUERY_ATTN_MODELS
 from .configuration import (
@@ -858,7 +859,7 @@ def _from_pretrained(
             init_cls = OVBloomForCausalLM
         elif model_type == "gpt-bigcode":
             init_cls = OVGPTBigCodeForCausalLM
-        elif model_type == "mamba":
+        elif model_type in SSM_MODELS:
             init_cls = OVMambaForCausalLM
         else:
             init_cls = cls
@@ -1138,17 +1139,13 @@ def forward(
                 self._past_length = 0
 
         ssm_states, conv_states = [], []
-        print(inputs.keys())
-
         self.request.start_async(inputs, share_inputs=True)
         self.request.wait()
         logits = torch.from_numpy(self.request.get_tensor("logits").data).to(self.device)
 
         if self.stateful:
             self._past_length += input_ids.shape[1]
         else:
-            print(self.ssm_cache_output_names)
-            print(self.conv_cache_output_names)
             ssm_states = [self.request.get_tensor(key).data for key in self.ssm_cache_output_names]
             conv_states = [self.request.get_tensor(key).data for key in self.conv_cache_output_names]
         cache_params = OVMambaCache(self.config, input_ids.shape[0], conv_states=conv_states, ssm_states=ssm_states)
@@ -1159,7 +1156,6 @@ def _update_model_kwargs_for_generation(
         self, outputs: ModelOutput, model_kwargs: Dict[str, Any], num_new_tokens: int = 1, **kwargs
     ) -> Dict[str, Any]:
         model_kwargs["cache_params"] = outputs.get("cache_params", None)
-        print(model_kwargs["cache_params"])
         if (
             model_kwargs.get("use_cache", True)
             and "cache_position" in model_kwargs