optimum-rbln 0.1.13__py3-none-any.whl → 0.2.0__py3-none-any.whl

optimum/rbln/transformers/models/bart/bart_architecture.py

@@ -21,497 +21,140 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
-from typing import Optional, Tuple
+from typing import Tuple
 
 import torch
 from torch import nn
 from transformers.modeling_attn_mask_utils import (
     _prepare_4d_attention_mask,
-    _prepare_4d_attention_mask_for_sdpa,
-    _prepare_4d_causal_attention_mask,
-    _prepare_4d_causal_attention_mask_for_sdpa,
-)
-from transformers.modeling_outputs import (
-    BaseModelOutputWithPastAndCrossAttentions,
-)
-from transformers.models.bart.modeling_bart import (
-    BartAttention,
-    BartDecoder,
-    BartDecoderLayer,
-    BartForConditionalGeneration,
-    BartSdpaAttention,
 )
 from transformers.utils import logging
 
+from ..seq2seq.seq2seq_architecture import (
+    Seq2SeqDecoder,
+    Seq2SeqDecoderLayer,
+    Seq2SeqDecoderWrapper,
+    Seq2SeqEncoderWrapper,
+    Seq2SeqForConditionalGeneration,
+    Seq2SeqSelfAttention,
+)
+
 
 logger = logging.get_logger(__name__)
 
 
 class BartWrapper:
-    def __init__(self, model):
-        self.encoder = BartEncoderWrapper(model)
+    def __init__(self, model: nn.Module, enc_max_seq_len: int):
+        self.encoder = Seq2SeqEncoderWrapper(model, enc_max_seq_len)
         self.decoder = BartDecoderWrapper(model)
 
 
-class _BartAttention(BartAttention):
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        past_key_value: Tuple[torch.Tensor],
-        attention_mask: torch.Tensor,
-        cache_position: torch.Tensor,
-        batch_index: torch.Tensor,
-        key_value_states: Optional[torch.Tensor] = None,
-    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor]]:
-        bsz, tgt_len, _ = hidden_states.size()
-        is_cross_attention = key_value_states is not None
+class BartDecoderWrapper(Seq2SeqDecoderWrapper):
+    def convert_to_rbln_conditional_generation(self, model: nn.Module):
+        new_layers = []
+        for layer in model.get_decoder().layers:
+            self_attn = BartSelfAttention(layer.self_attn)
+            new_layers.append(BartDecoderLayer(layer, self_attn))
 
-        query_states = self.q_proj(hidden_states) * self.scaling
+        decoder_model = BartDecoder(model.get_decoder(), new_layers)
+        new_model = BartForConditionalGeneration(model, decoder_model)
 
-        if is_cross_attention:
-            is_dummy_decoder = len(key_value_states.shape) > 1
-            if is_dummy_decoder:
-                key_states = self._shape(self.k_proj(key_value_states), -1, bsz)
-                value_states = self._shape(self.v_proj(key_value_states), -1, bsz)
-            else:
-                key_states = past_key_value[0]
-                value_states = past_key_value[1]
-        else:
-            key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
-            value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
-
-        if cache_position.dim() > 0:
-            proj_shape = (bsz, self.num_heads, -1, self.head_dim)
-            query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
-            key_states = key_states.reshape(*proj_shape)
-            value_states = value_states.reshape(*proj_shape)
-
-            all_key_states = []
-            all_value_states = []
-            all_attn_output = []
-            for b in range(bsz):
-                batch_query_states = query_states[b].unsqueeze(0).unsqueeze(2)
-                batch_attention_mask = attention_mask[b].unsqueeze(0).unsqueeze(2)
-                batch_key_states = key_states[b].unsqueeze(0).unsqueeze(2)
-                batch_value_states = value_states[b].unsqueeze(0).unsqueeze(2)
-                if not is_cross_attention:
-                    batch_key_states = (
-                        past_key_value[0][b]
-                        .unsqueeze(0)
-                        .unsqueeze(2)
-                        .slice_scatter(
-                            batch_key_states, dim=-2, start=cache_position[b][0], end=cache_position[b][0] + 1
-                        )
-                    )
-                    batch_value_states = (
-                        past_key_value[1][b]
-                        .unsqueeze(0)
-                        .unsqueeze(2)
-                        .slice_scatter(
-                            batch_value_states, dim=-2, start=cache_position[b][0], end=cache_position[b][0] + 1
-                        )
-                    )
-                attn_weights = torch.matmul(batch_query_states, batch_key_states.transpose(3, 4))
-                attn_weights = attn_weights + batch_attention_mask
-                attn_weights = nn.functional.softmax(attn_weights, dim=-1)
-
-                attn_output = torch.matmul(attn_weights, batch_value_states)
-                attn_output = attn_output.view(1, self.num_heads, tgt_len, self.head_dim)
-                attn_output = attn_output.transpose(1, 2)
-                attn_output = attn_output.reshape(1, tgt_len, self.embed_dim)
-                all_key_states.append(batch_key_states)
-                all_value_states.append(batch_value_states)
-                all_attn_output.append(attn_output)
-            key_states = torch.cat(all_key_states, dim=0).squeeze(2)
-            value_states = torch.cat(all_value_states, dim=0).squeeze(2)
-            attn_output = torch.cat(all_attn_output, dim=0)
+        return new_model
 
-        else:
-            if batch_index is None or batch_index == -1:
-                batch_index = 0
-
-            if not is_cross_attention:
-                key_states = past_key_value[0].slice_scatter(
-                    key_states, dim=2, start=cache_position, end=cache_position + 1
-                )
-                value_states = past_key_value[1].slice_scatter(
-                    value_states, dim=2, start=cache_position, end=cache_position + 1
-                )
-
-            proj_shape = (bsz * self.num_heads, -1, self.head_dim)
-            query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
-            key_states = key_states.reshape(*proj_shape)
-            value_states = value_states.reshape(*proj_shape)
-
-            src_len = key_states.size(1)
-            attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
-            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
-            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
-            attn_weights = nn.functional.softmax(attn_weights, dim=-1)
-
-            attn_output = torch.bmm(attn_weights, value_states)
-            attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
-            attn_output = attn_output.transpose(1, 2)
-            key_states = key_states.unsqueeze(0)
-            value_states = value_states.unsqueeze(0)
-            attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)
-
-        attn_output = self.out_proj(attn_output)
-
-        present_key_value = (key_states, value_states)
-
-        return attn_output, present_key_value
-
-
-class _BartSdpaAttention(BartSdpaAttention):
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        past_key_value: Tuple[torch.Tensor],
-        attention_mask: torch.Tensor,
-        cache_position: torch.Tensor,
-        batch_index: torch.Tensor,
-        key_value_states: Optional[torch.Tensor] = None,
-    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor]]:
-        bsz, tgt_len, _ = hidden_states.size()
-        is_cross_attention = key_value_states is not None
-
-        query_states = self.q_proj(hidden_states)
-
-        if is_cross_attention:
-            is_dummy_decoder = len(key_value_states.shape) > 1
-            if is_dummy_decoder:
-                key_states = self._shape(self.k_proj(key_value_states), -1, bsz)
-                value_states = self._shape(self.v_proj(key_value_states), -1, bsz)
-            else:
-                key_states = past_key_value[0]
-                value_states = past_key_value[1]
-        else:
-            key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
-            value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
-
-        query_states = self._shape(query_states, tgt_len, bsz)
-
-        if (batch_index is None or batch_index == -1) and bsz > 1:
-            all_key_states = []
-            all_value_states = []
-            all_attn_output = []
-
-            for b in range(bsz):
-                batch_query_states = query_states[b].unsqueeze(0)
-                batch_attention_mask = attention_mask[b].unsqueeze(0)
-                batch_key_states = key_states[b].unsqueeze(0)
-                batch_value_states = value_states[b].unsqueeze(0)
-
-                if not is_cross_attention:
-                    batch_key_states = (
-                        past_key_value[0][b]
-                        .unsqueeze(0)
-                        .slice_scatter(
-                            batch_key_states, dim=-2, start=cache_position[b][0], end=cache_position[b][0] + 1
-                        )
-                    )
-                    batch_value_states = (
-                        past_key_value[1][b]
-                        .unsqueeze(0)
-                        .slice_scatter(
-                            batch_value_states, dim=-2, start=cache_position[b][0], end=cache_position[b][0] + 1
-                        )
-                    )
-
-                attn_output = torch.nn.functional.scaled_dot_product_attention(
-                    batch_query_states, batch_key_states, batch_value_states, attn_mask=batch_attention_mask
-                )
-                attn_output = attn_output.transpose(1, 2)
-                attn_output = attn_output.reshape(1, tgt_len, self.embed_dim)
-                all_key_states.append(batch_key_states)
-                all_value_states.append(batch_value_states)
-                all_attn_output.append(attn_output)
-
-            key_states = torch.cat(all_key_states, dim=0)
-            value_states = torch.cat(all_value_states, dim=0)
-            attn_output = torch.cat(all_attn_output, dim=0)
 
+class BartForConditionalGeneration(Seq2SeqForConditionalGeneration):
+    has_rescaling = False
+
+    def __post_init__(self):
+        self.scaling = self.config.d_model**-0.5
+
+
+class BartDecoder(Seq2SeqDecoder):
+    has_pos_emb = True
+
+    def __post_init__(self):
+        self.embed_positions = self._original_mod.embed_positions
+        self.layernorm_embedding = self._original_mod.layernorm_embedding
+        self.embed_scale = getattr(self._original_mod, "embed_scale", None)
+
+    def prepare_attn_mask(self, attention_mask, encoder_attention_mask, **kwargs):
+        attention_mask = attention_mask[:, None, None, :]
+        encoder_attention_mask = _prepare_4d_attention_mask(encoder_attention_mask, torch.float32, tgt_len=1)
+
+        return attention_mask, encoder_attention_mask
+
+    def apply_position_embedding(self, inputs_embeds, cache_position):
+        hidden_all = []
+        for i in range(inputs_embeds.shape[0]):
+            positions_idx = cache_position[i]
+            position_weight = self.embed_positions.weight[2:]
+            position = position_weight[positions_idx]
+            batch_hidden = position + inputs_embeds[i]
+            hidden_all.append(batch_hidden)
+        hidden_states = torch.stack(hidden_all, dim=0)
+
+        hidden_states = self.layernorm_embedding(hidden_states)
+
+        return hidden_states
+
+    def get_embedding(self):
+        if self.embed_scale is not None:
+            return lambda x: self.embed_tokens(x) * self.embed_scale
         else:
-            if batch_index is None or batch_index == -1:
-                batch_index = 0
-
-            if not is_cross_attention:
-                key_states = past_key_value[0].slice_scatter(
-                    key_states, dim=2, start=cache_position, end=cache_position + 1
-                )
-                value_states = past_key_value[1].slice_scatter(
-                    value_states, dim=2, start=cache_position, end=cache_position + 1
-                )
-
-            # need 4d shape (input tensors) for scaled_dot_product_attention
-            attn_output = torch.nn.functional.scaled_dot_product_attention(
-                query_states,
-                key_states,
-                value_states,
-                attn_mask=attention_mask,
-            )
-            attn_output = attn_output.transpose(1, 2)
-            attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)
-
-        attn_output = self.out_proj(attn_output)
-
-        present_key_value = (key_states, value_states)
-
-        return attn_output, present_key_value
-
-
-ATTN_FORWARD_MAP = {"eager": _BartAttention.forward, "sdpa": _BartSdpaAttention.forward}
-
-
-class _BartDecoderLayer(BartDecoderLayer):
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor,
-        encoder_attention_mask: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        past_key_value: Tuple[torch.Tensor],
-        cache_position: torch.Tensor,
-        batch_ids: torch.Tensor,
-        attn_impl: str = "eager",
-    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor]]:
-        # Self Attention Block
-        residual = hidden_states
-        self_attn_past_key_value = past_key_value[:2]
-
-        hidden_states, present_key_value = ATTN_FORWARD_MAP[attn_impl](
-            self.self_attn,
-            hidden_states=hidden_states,
-            past_key_value=self_attn_past_key_value,
-            attention_mask=attention_mask,
-            cache_position=cache_position,
-            batch_index=batch_ids,
-        )
-        hidden_states = residual + hidden_states
-        hidden_states = self.self_attn_layer_norm(hidden_states)
+            return self.embed_tokens
 
-        # Cross-Attention Block
-        residual = hidden_states
-        cross_attn_past_key_value = past_key_value[-2:]
-
-        hidden_states, cross_attn_present_key_value = ATTN_FORWARD_MAP[attn_impl](
-            self.encoder_attn,
-            hidden_states=hidden_states,
-            key_value_states=encoder_hidden_states,
-            past_key_value=cross_attn_past_key_value,
-            attention_mask=encoder_attention_mask,
-            cache_position=cache_position,
-            batch_index=batch_ids,
-        )
-        hidden_states = residual + hidden_states
-        hidden_states = self.encoder_attn_layer_norm(hidden_states)
-        present_key_value = present_key_value + cross_attn_present_key_value
 
-        # Fully Connected Block
+class BartLayerFF(nn.Module):
+    def __init__(self, decoder_layer):
+        super().__init__()
+        self.fc1 = decoder_layer.fc1
+        self.fc2 = decoder_layer.fc2
+        self.activation_fn = decoder_layer.activation_fn
+        self.layer_norm = decoder_layer.final_layer_norm
+
+    def forward(self, hidden_states):
+        # Residual Connection
         residual = hidden_states
         hidden_states = self.activation_fn(self.fc1(hidden_states))
         hidden_states = self.fc2(hidden_states)
         hidden_states = residual + hidden_states
-        hidden_states = self.final_layer_norm(hidden_states)
-
-        return hidden_states, present_key_value
-
-
-class _BartDecoder(BartDecoder):
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        attention_mask: torch.Tensor,
-        encoder_attention_mask: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        past_key_values: torch.Tensor,
-        cache_position: torch.Tensor,
-        batch_ids: torch.Tensor,
-        attn_impl: str = "eager",
-    ):
-        # embedding
-        if hasattr(self, "embed_scale"):
-            inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
-        else:
-            inputs_embeds = self.embed_tokens(input_ids)
+        hidden_states = self.layer_norm(hidden_states)
+        return hidden_states
 
-        if cache_position.dim() == 0:
-            positions_idx = cache_position + self.embed_positions.offset
-            positions = self.embed_positions.weight[positions_idx]
-            hidden_states = inputs_embeds + positions
-        else:
-            hidden_all = []
-            # compiler pattern base dependency -> take + add
-            for i in range(input_ids.shape[0]):
-                # cache position [N,1]
-                positions_idx = cache_position[i]
-                # offset is set 2 in bart embedding
-                position_weight = self.embed_positions.weight[2:]
-                position = position_weight[positions_idx]
-                batch_hidden = position + inputs_embeds[i]
-                hidden_all.append(batch_hidden)
-            hidden_states = torch.stack(hidden_all, dim=0)
 
-        hidden_states = self.layernorm_embedding(hidden_states)
+class BartDecoderLayer(Seq2SeqDecoderLayer):
+    def __post_init__(self):
+        self.self_attn_layer_norm = self._original_mod.self_attn_layer_norm
+        self.encoder_attn = self._original_mod.encoder_attn
+        self.encoder_attn_layer_norm = self._original_mod.encoder_attn_layer_norm
+        self.ff_layer = BartLayerFF(self._original_mod)
 
-        # prepare attn_mask
-        input_shape = input_ids.size()
-        if self._use_sdpa:
-            attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
-                attention_mask, input_shape, inputs_embeds, cache_position
-            )
-            encoder_attention_mask = _prepare_4d_attention_mask_for_sdpa(
-                encoder_attention_mask, torch.float32, tgt_len=input_shape[-1]
-            )
-        else:
-            attention_mask = _prepare_4d_causal_attention_mask(
-                attention_mask, input_shape, inputs_embeds, cache_position
-            )
-            encoder_attention_mask = _prepare_4d_attention_mask(
-                encoder_attention_mask, torch.float32, tgt_len=input_shape[-1]
-            )
-
-        # iterate decoder_layer
-        next_decoder_cache = ()
-        for idx, decoder_layer in enumerate(self.layers):
-            past_key_value = past_key_values[idx]
-            layer_outputs = _BartDecoderLayer.forward(
-                decoder_layer,
-                hidden_states,
-                attention_mask=attention_mask,
-                encoder_hidden_states=encoder_hidden_states,
-                encoder_attention_mask=encoder_attention_mask,
-                past_key_value=past_key_value,
-                cache_position=cache_position,
-                batch_ids=batch_ids,
-                attn_impl=attn_impl,
-            )
-            hidden_states = layer_outputs[0]
-            next_decoder_cache += (layer_outputs[1],)
-
-        return BaseModelOutputWithPastAndCrossAttentions(
-            last_hidden_state=hidden_states,
-            past_key_values=next_decoder_cache,
-        )
-
-
-class BartDecoderWrapper(torch.nn.Module):
-    def __init__(self, model: "BartForConditionalGeneration"):
-        super().__init__()
-        self.config = model.config
-        self.decoder = model.get_decoder()
-        self.num_layers = self.config.decoder_layers
-        self.lm_head = model.lm_head
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        attention_mask: torch.Tensor,
-        encoder_attention_mask: torch.Tensor,
-        cache_position: torch.Tensor,
-        batch_position: torch.Tensor,
-        self_kv_cache: torch.Tensor,
-        cross_kv_cache: torch.Tensor,
-    ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor]]:
-        if input_ids.shape[1] == 1:
-            rbln_batch_position = None
-        else:
-            rbln_batch_position = batch_position
-        # prepare past_key_values
-        kv_cache = ()
-        for i in range(0, self.num_layers * 2, 2):
-            kv_cache = kv_cache + (
-                (
-                    self_kv_cache[i],
-                    self_kv_cache[i + 1],
-                    cross_kv_cache[i],
-                    cross_kv_cache[i + 1],
-                ),
-            )
-        # decode
-        decoder_outputs = _BartDecoder.forward(
-            self.decoder,
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            encoder_attention_mask=encoder_attention_mask,
-            cache_position=cache_position,
-            past_key_values=kv_cache,
-            encoder_hidden_states=torch.tensor([1]),
-            attn_impl=self.config._attn_implementation,
-            batch_ids=rbln_batch_position,
-        )
-        sequence_output = decoder_outputs[0]
-        lm_logits = self.lm_head(sequence_output)
-
-        # get self_kv_cache from ouputs
-        past_key_values = decoder_outputs[1]
-        self_kv_cache = []
-        for i in range(self.num_layers):
-            self_kv_cache.append(past_key_values[i][0])
-            self_kv_cache.append(past_key_values[i][1])
-        self_kv_cache = torch.stack(self_kv_cache, dim=0)
-
-        # return batch_position to keep it as a variable within the graph
-        return lm_logits, self_kv_cache, batch_position
-
-
-class BartEncoderWrapper(torch.nn.Module):
-    def __init__(self, model):
-        super().__init__()
-        self.model = model
-        self.config = model.config
-        self.decoder = model.get_decoder()
-        self.encoder = model.get_encoder()
-        self.num_layers = self.config.encoder_layers
-        self.decoder_max_length = self.config.max_position_embeddings
-        self.encoder_max_length = self.config.max_position_embeddings
-        self.num_heads = self.config.decoder_attention_heads
-        self.d_kv = self.config.d_model // self.num_heads
-
-    def forward(
-        self,
-        input_ids: torch.LongTensor,
-        attention_mask: torch.LongTensor,
-        cross_key_value: torch.Tensor = None,
-        batch_idx: torch.Tensor = None,
-    ) -> Tuple[torch.Tensor]:
-        # 1. run encoder
-        encoder_outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
-        last_hidden_states = encoder_outputs[0]
-
-        # 2. run dummy decoder to get pre-calculated cross-key_values for generation
-        dummy_past_key_value = []
-        for _ in range(self.num_layers):
-            pkv_self_attn_key = torch.zeros(1, self.num_heads, self.decoder_max_length, self.d_kv)
-            pkv_self_attn_value = torch.zeros(1, self.num_heads, self.decoder_max_length, self.d_kv)
-            pkv_cross_attn_key = torch.zeros(1, self.num_heads, self.encoder_max_length, self.d_kv)
-            pkv_cross_attn_value = torch.zeros(1, self.num_heads, self.encoder_max_length, self.d_kv)
-            layer_pkv = (pkv_self_attn_key, pkv_self_attn_value, pkv_cross_attn_key, pkv_cross_attn_value)
-            dummy_past_key_value.append(layer_pkv)
-
-        decoder_attention_mask = torch.zeros(1, self.decoder_max_length, dtype=torch.float32)
-        decoder_attention_mask[:, :1] = 1
-
-        decoder_outputs = _BartDecoder.forward(
-            self.decoder,
-            input_ids=torch.zeros((1, 1), dtype=torch.int64),
-            attention_mask=decoder_attention_mask,
-            encoder_attention_mask=attention_mask,
-            cache_position=torch.tensor(0, dtype=torch.int32),
-            encoder_hidden_states=last_hidden_states,
-            past_key_values=dummy_past_key_value,
-            batch_ids=torch.tensor(0, dtype=torch.int32),
-            attn_impl=self.config._attn_implementation,
-        )
-        first_past_kv = decoder_outputs[1]
-
-        encoder_kv = []
-        for i in range(self.model.config.decoder_layers):
-            encoder_kv.append(first_past_kv[i][2].unsqueeze(0))
-            encoder_kv.append(first_past_kv[i][3].unsqueeze(0))
-        encoder_kv = torch.cat(encoder_kv, dim=0)
-
-        cross_key_value = cross_key_value.slice_scatter(encoder_kv, dim=1, start=batch_idx, end=batch_idx + 1)
-
-        return cross_key_value
+    def pre_self_attn_layer_norm(self, hidden_states):
+        return hidden_states
+
+    def post_self_attn_layer_norm(self, hidden_states):
+        return self.self_attn_layer_norm(hidden_states)
+
+    def pre_cross_attn_layer_norm(self, hidden_states):
+        return hidden_states
+
+    def post_cross_attn_layer_norm(self, hidden_states):
+        return self.encoder_attn_layer_norm(hidden_states)
+
+
+class BartSelfAttention(Seq2SeqSelfAttention):
+    def __post_init__(self):
+        self.q_proj = self._original_mod.q_proj
+        self.k_proj = self._original_mod.k_proj
+        self.v_proj = self._original_mod.v_proj
+        self.out_proj = self._original_mod.out_proj
+        self.num_heads = self._original_mod.num_heads
+        self.head_dim = self._original_mod.embed_dim // self._original_mod.num_heads
+        self.scaling = self.head_dim**-0.5
+        self.attn_decode = torch.ops.rbln_custom_ops.attn_decode
+
+    def projection(self, hidden_states) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        query_states = self.q_proj(hidden_states) * self.scaling
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+        return query_states, key_states, value_states

optimum/rbln/transformers/models/bart/modeling_bart.py

@@ -24,9 +24,9 @@
 import inspect
 from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
 
-from transformers import BartConfig, BartForConditionalGeneration, BartModel, PretrainedConfig
+from transformers import BartForConditionalGeneration, PretrainedConfig, PreTrainedModel
 
-from ....modeling_base import RBLNModel
+from ....modeling import RBLNModel
 from ....modeling_config import RBLNCompileConfig, RBLNConfig
 from ....utils.logging import get_logger
 from ...models.seq2seq import RBLNModelForSeq2SeqLM
@@ -41,9 +41,6 @@ if TYPE_CHECKING:
 
 
 class RBLNBartModel(RBLNModel):
-    original_model_class = BartModel
-    original_config_class = BartConfig
-
     @classmethod
     def _get_rbln_config(
         cls,
@@ -82,7 +79,7 @@ class RBLNBartModel(RBLNModel):
             if rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names"):
                 rbln_model_input_names = cls.rbln_model_input_names
             elif rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names") is False:
-                input_names_order = inspect.signature(cls.original_model_class.forward).parameters.keys()
+                input_names_order = inspect.signature(cls.hf_class.forward).parameters.keys()
                 raise ValueError(
                     "Specify the model input names obtained by the tokenizer via `rbln_model_input_names`, "
                     f"and be sure to make the order of the inputs same as BartModel forward() arguments like ({list(input_names_order)})"
@@ -96,11 +93,12 @@ class RBLNBartModel(RBLNModel):
             for model_input_name in rbln_model_input_names
         ]
 
-        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
+        enc_compile_config = RBLNCompileConfig(input_info=input_info, compiled_model_name="encoder")
+        dec_compile_config = RBLNCompileConfig(input_info=input_info, compiled_model_name="decoder")
 
         rbln_config = RBLNConfig(
             rbln_cls=cls.__name__,
-            compile_cfgs=[rbln_compile_config],
+            compile_cfgs=[enc_compile_config, dec_compile_config],
             rbln_kwargs=rbln_kwargs,
         )
 
@@ -111,7 +109,10 @@ class RBLNBartModel(RBLNModel):
 class RBLNBartForConditionalGeneration(RBLNModelForSeq2SeqLM):
     @classmethod
     def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
-        return BartWrapper(model)
+        enc_max_seq_len = (
+            rbln_config.model_cfg["enc_max_seq_len"] if "enc_max_seq_len" in rbln_config.model_cfg else 1024
+        )
+        return BartWrapper(model, enc_max_seq_len=enc_max_seq_len)
 
     def __getattr__(self, __name: str) -> Any:
         def redirect(func):