optimum-rbln 0.1.0__py3-none-any.whl

optimum/rbln/transformers/models/t5/t5_architecture.py

@@ -0,0 +1,439 @@
+# Copyright 2024 Rebellions Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+from typing import TYPE_CHECKING, Optional, Tuple
+
+import torch
+from torch import nn
+from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPastAndCrossAttentions
+from transformers.models.t5.configuration_t5 import T5Config
+from transformers.models.t5.modeling_t5 import (
+    T5Attention,
+    T5Block,
+    T5LayerCrossAttention,
+    T5LayerSelfAttention,
+    T5Stack,
+)
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+if TYPE_CHECKING:
+    from transformers import T5ForConditionalGeneration
+
+
+class T5Encoder(T5Stack):
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        position_bias: torch.Tensor,
+    ) -> BaseModelOutput:
+        hidden_states = self.embed_tokens(input_ids)
+        extended_attention_mask = self.invert_attention_mask(attention_mask)
+        position_bias = position_bias + extended_attention_mask
+        for i, layer_module in enumerate(self.block):
+            layer_outputs = _T5Block.forward(
+                layer_module,
+                hidden_states,
+                position_bias=position_bias,
+            )
+            hidden_states = layer_outputs[0]
+        hidden_states = self.final_layer_norm(hidden_states)
+        return BaseModelOutput(last_hidden_state=hidden_states)
+
+
+class T5Decoder(T5Stack):
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        encoder_attention_mask: torch.Tensor,
+        past_key_values: torch.Tensor,
+        position_bias: torch.Tensor,
+        encoder_decoder_position_bias: torch.Tensor,
+        cache_position: torch.Tensor,
+    ) -> BaseModelOutputWithPastAndCrossAttentions:
+        hidden_states = self.embed_tokens(input_ids)
+        extended_attention_mask = self.invert_attention_mask(attention_mask)
+        encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+
+        position_bias = position_bias + extended_attention_mask
+        encoder_decoder_position_bias = encoder_decoder_position_bias + encoder_extended_attention_mask
+
+        present_key_value_states = ()
+        for layer_module, past_key_value in zip(self.block, past_key_values):
+            layer_outputs = _T5Block.forward(
+                layer_module,
+                hidden_states,
+                position_bias=position_bias,
+                encoder_hidden_states=encoder_hidden_states,
+                encoder_decoder_position_bias=encoder_decoder_position_bias,
+                past_key_value=past_key_value,
+                cache_position=cache_position,
+            )
+            hidden_states, present_key_value_state = layer_outputs[:2]
+            present_key_value_states = present_key_value_states + (present_key_value_state,)
+
+        hidden_states = self.final_layer_norm(hidden_states)
+
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=present_key_value_states,
+        )
+
+
+class T5EncoderWrapper(torch.nn.Module):
+    def __init__(self, model: "T5ForConditionalGeneration"):
+        super().__init__()
+        self.config = model.config
+        self.model = model
+        self.encoder = model.encoder
+        self.decoder = model.decoder
+        self.default_max_length = getattr(self.config, "n_positions", None) or getattr(
+            self.config, "max_position_embeddings", None
+        )
+        self.encoder_max_length = None
+        self.decoder_max_length = None
+        self.decoder_batch_size = 1
+
+    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        encoder_batch_size = input_ids.shape[0]
+        decoder_batch_size = self.decoder_batch_size
+        decoder_max_length = self.decoder_max_length or self.default_max_length
+        encoder_max_length = self.encoder_max_length or self.default_max_length
+
+        attn_layer = self.encoder.block[0].layer[0].SelfAttention
+        encoder_position_bias = T5Attention.compute_bias(attn_layer, encoder_max_length, encoder_max_length)
+        encoder_outputs = T5Encoder.forward(self.encoder, input_ids, attention_mask, encoder_position_bias)
+
+        attn_layer = self.decoder.block[0].layer[0].SelfAttention
+        decoder_position_bias = T5Attention.compute_bias(attn_layer, decoder_max_length, decoder_max_length)
+        decoder_position_bias = decoder_position_bias[:, :, :1]
+
+        attn_layer = self.decoder.block[0].layer[1].EncDecAttention
+        encoder_decoder_position_bias = torch.zeros(1, attn_layer.n_heads, 1, encoder_max_length)
+
+        dummy_past_key_value = []
+        for i in range(self.config.num_layers):
+            pkv_self_attn_key = torch.zeros(
+                decoder_batch_size, self.config.num_heads, decoder_max_length, self.config.d_kv
+            )
+            pkv_self_attn_value = torch.zeros(
+                decoder_batch_size, self.config.num_heads, decoder_max_length, self.config.d_kv
+            )
+            pkv_cross_attn_key = torch.zeros(
+                encoder_batch_size, self.config.num_heads, encoder_max_length, self.config.d_kv
+            )
+            pkv_cross_attn_value = torch.zeros(
+                encoder_batch_size, self.config.num_heads, encoder_max_length, self.config.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(decoder_batch_size, decoder_max_length, dtype=torch.int64)
+        decoder_attention_mask[:, :1] = 1
+
+        # Since first step of decoder has different graph to further step of it,
+        # here we merges decoder into its corresponding encoder.
+        # TODO(jongho): Separate first-step-decoder.
+        decoder_outputs = T5Decoder.forward(
+            self.decoder,
+            input_ids=torch.zeros(decoder_batch_size, 1, dtype=torch.int64),
+            attention_mask=decoder_attention_mask,
+            position_bias=decoder_position_bias,
+            encoder_decoder_position_bias=encoder_decoder_position_bias,
+            encoder_hidden_states=encoder_outputs.last_hidden_state,
+            encoder_attention_mask=attention_mask,
+            past_key_values=dummy_past_key_value,
+            cache_position=torch.tensor(0, dtype=torch.int32),
+        )
+
+        past_key_values = decoder_outputs.past_key_values
+
+        cross_kv_cache = []
+        for i in range(self.model.config.num_layers):
+            cross_kv_cache.append(past_key_values[i][2])
+            cross_kv_cache.append(past_key_values[i][3])
+        cross_kv_cache = torch.stack(cross_kv_cache, dim=0)
+
+        return cross_kv_cache
+
+
+class T5DecoderWrapper(torch.nn.Module):
+    def __init__(self, model: "T5ForConditionalGeneration"):
+        super().__init__()
+        self.config = model.config
+        self.model = model
+        self.encoder = model.encoder
+        self.decoder = model.decoder
+        self.default_max_length = getattr(self.config, "n_positions", None) or getattr(
+            self.config, "max_position_embeddings", None
+        )
+        self.encoder_max_length = None
+        self.decoder_max_length = None
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        encoder_attention_mask: torch.Tensor,
+        cache_position: torch.Tensor,
+        self_kv_cache: torch.Tensor,
+        cross_kv_cache: torch.Tensor,
+    ) -> Tuple[torch.Tensor]:
+        # cache_position : step 0부터
+        # attention_mask : 1개가 색칠된것부터 ([0:cache_position+1])
+        num_layers = self.model.config.num_layers
+        encoder_max_length = self.encoder_max_length or self.default_max_length
+        decoder_max_length = self.decoder_max_length or self.default_max_length
+
+        kv_cache = ()
+        for i in range(0, 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],
+                ),
+            )
+
+        attn_layer = self.model.decoder.block[0].layer[0].SelfAttention
+        _decoder_position_bias = T5Attention.compute_bias(attn_layer, decoder_max_length, decoder_max_length)
+        decoder_position_bias = _decoder_position_bias[:, :, cache_position].unsqueeze(2)
+
+        attn_layer = self.model.decoder.block[0].layer[1].EncDecAttention
+        encoder_decoder_position_bias = torch.zeros(1, attn_layer.n_heads, 1, encoder_max_length)
+
+        decoder_outputs = T5Decoder.forward(
+            self.model.decoder,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            encoder_hidden_states=1,
+            encoder_attention_mask=encoder_attention_mask,
+            position_bias=decoder_position_bias,
+            encoder_decoder_position_bias=encoder_decoder_position_bias,
+            past_key_values=kv_cache,
+            cache_position=cache_position,
+        )
+
+        past_key_values = decoder_outputs.past_key_values
+        sequence_output = decoder_outputs[0]
+        if self.model.config.tie_word_embeddings:
+            # Rescale output before projecting on vocab
+            # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/transformer.py#L586
+            sequence_output = sequence_output * (self.model.model_dim**-0.5)
+        lm_logits = self.model.lm_head(sequence_output)
+
+        self_kv_cache = []
+        for i in range(self.model.config.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 lm_logits, self_kv_cache
+
+
+class _T5Attention(T5Attention):
+    def __init__(self, config: T5Config, has_relative_attention_bias=False):
+        super().__init__(config, has_relative_attention_bias)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_value_states: Tuple[torch.Tensor] = None,
+        position_bias: torch.Tensor = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        cache_position: Optional[torch.Tensor] = None,  # 현재 cache sequence 길이
+        is_self_attn: Optional[bool] = None,
+    ) -> Tuple[torch.Tensor]:
+        batch_size = hidden_states.shape[0]
+        cross_batch_size = key_value_states.shape[0] if not is_self_attn and cache_position == 0 else None
+
+        def shape(states, batch_size):
+            """projection"""
+            return states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2)
+
+        def unshape(states, batch_size):
+            """reshape"""
+            return states.transpose(1, 2).contiguous().view(batch_size, -1, self.inner_dim)
+
+        query_states = shape(self.q(hidden_states), batch_size)  # (batch_size, n_heads, seq_length, dim_per_head)
+
+        # projection
+        if is_self_attn:
+            key_states = shape(self.k(hidden_states), batch_size)
+            value_states = shape(self.v(hidden_states), batch_size)
+            if past_key_value is not None:
+                # decoder self attn
+                cache_k = past_key_value[0].slice_scatter(
+                    key_states, dim=2, start=cache_position, end=cache_position + 1
+                )
+                cache_v = past_key_value[1].slice_scatter(
+                    value_states, dim=2, start=cache_position, end=cache_position + 1
+                )
+                past_key_value = (cache_k, cache_v)
+                key_states, value_states = past_key_value
+
+        else:
+            # cross-attn
+            if cache_position == 0:
+                key_states = shape(self.k(key_value_states), cross_batch_size)
+                value_states = shape(self.v(key_value_states), cross_batch_size)
+                past_key_value = key_states, value_states
+            else:
+                key_states = past_key_value[0]
+                value_states = past_key_value[1]
+
+        # compute scores
+        scores = torch.matmul(query_states, key_states.transpose(3, 2))
+        scores += position_bias
+
+        attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(
+            scores
+        )  # (batch_size, n_heads, seq_length, key_length)
+
+        attn_output = unshape(torch.matmul(attn_weights, value_states), batch_size)  # (batch_size, seq_length, dim)
+        attn_output = self.o(attn_output)
+
+        outputs = (attn_output,) + (past_key_value,)
+        return outputs
+
+
+class _T5LayerSelfAttention(T5LayerSelfAttention):
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_bias: torch.Tensor = None,
+        past_key_value: Tuple[torch.Tensor] = None,
+        cache_position: Optional[torch.Tensor] = None,
+    ):
+        normed_hidden_states = self.layer_norm(hidden_states)
+        attention_output = _T5Attention.forward(
+            self.SelfAttention,
+            hidden_states=normed_hidden_states,
+            position_bias=position_bias,
+            past_key_value=past_key_value,
+            cache_position=cache_position,
+            is_self_attn=True,
+        )
+
+        # Residual Connection
+        hidden_states = hidden_states + self.dropout(attention_output[0])
+        outputs = (hidden_states,) + attention_output[1:]  # add attentions if we output them
+        return outputs
+
+
+class _T5LayerCrossAttention(T5LayerCrossAttention):
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_value_states: torch.Tensor,
+        position_bias: torch.Tensor = None,
+        past_key_value: Tuple[torch.Tensor] = None,
+        cache_position: Optional[torch.Tensor] = None,
+    ):
+        normed_hidden_states = self.layer_norm(hidden_states)
+        attention_output = _T5Attention.forward(
+            self.EncDecAttention,
+            hidden_states=normed_hidden_states,
+            key_value_states=key_value_states,
+            position_bias=position_bias,
+            past_key_value=past_key_value,
+            cache_position=cache_position,
+            is_self_attn=False,
+        )
+
+        # Residual connection
+        layer_output = hidden_states + self.dropout(attention_output[0])
+        outputs = (layer_output,) + attention_output[1:]  # add attentions if we output them
+        return outputs
+
+
+class _T5Block(T5Block):
+    def forward(
+        self,
+        hidden_states,
+        position_bias=None,
+        encoder_hidden_states=None,
+        encoder_decoder_position_bias=None,
+        past_key_value=None,
+        cache_position=None,
+    ):
+        if past_key_value is not None:
+            if not self.is_decoder:
+                logger.warning("`past_key_values` is passed to the encoder. Please make sure this is intended.")
+            expected_num_past_key_values = 2 if encoder_hidden_states is None else 4
+
+            if len(past_key_value) != expected_num_past_key_values:
+                raise ValueError(
+                    f"There should be {expected_num_past_key_values} past states. "
+                    f"{'2 (past / key) for cross attention. ' if expected_num_past_key_values == 4 else ''}"
+                    f"Got {len(past_key_value)} past key / value states"
+                )
+
+            self_attn_past_key_value = past_key_value[:2]
+            if self_attn_past_key_value == (None, None):
+                self_attn_past_key_value = None
+
+            cross_attn_past_key_value = past_key_value[2:]
+        else:
+            self_attn_past_key_value, cross_attn_past_key_value = None, None
+
+        self_attention_outputs = _T5LayerSelfAttention.forward(
+            self.layer[0],
+            hidden_states=hidden_states,
+            position_bias=position_bias,
+            past_key_value=self_attn_past_key_value,
+            cache_position=cache_position,
+        )
+
+        hidden_states, present_key_value_state = self_attention_outputs[:2]
+
+        do_cross_attention = self.is_decoder and encoder_hidden_states is not None
+        if do_cross_attention:
+            cross_attention_outputs = _T5LayerCrossAttention.forward(
+                self.layer[1],
+                hidden_states,
+                key_value_states=encoder_hidden_states,
+                position_bias=encoder_decoder_position_bias,
+                past_key_value=cross_attn_past_key_value,
+                cache_position=cache_position,
+            )
+            hidden_states = cross_attention_outputs[0]
+            # Combine self attn and cross attn key value states
+            if present_key_value_state is not None:
+                # print(present_key_value_state.shape)
+                present_key_value_state = present_key_value_state + cross_attention_outputs[1]
+
+        # Apply Feed Forward layer
+        hidden_states = self.layer[-1](hidden_states)
+
+        outputs = (hidden_states,)
+        outputs = outputs + (present_key_value_state,)
+
+        return outputs

optimum/rbln/transformers/models/wav2vec2/__init__.py

@@ -0,0 +1,24 @@
+# Copyright 2024 Rebellions Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+from .modeling_wav2vec2 import RBLNWav2Vec2ForCTC

optimum/rbln/transformers/models/wav2vec2/modeling_wav2vec2.py

@@ -0,0 +1,121 @@
+# Copyright 2024 Rebellions Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+import logging
+from typing import TYPE_CHECKING, Optional, Union
+
+import torch
+from transformers import AutoModelForMaskedLM, PretrainedConfig, Wav2Vec2ForCTC
+from transformers.modeling_outputs import CausalLMOutput
+
+from ....modeling_base import RBLNModel
+from ....modeling_config import RBLNConfig, RBLNRuntimeConfig
+
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from transformers import (
+        AutoFeatureExtractor,
+        AutoProcessor,
+        AutoTokenizer,
+        PretrainedConfig,
+    )
+
+
+class _Wav2Vec2(torch.nn.Module):
+    def __init__(self, model: "Wav2Vec2ForCTC"):
+        super().__init__()
+        self.model = model
+
+    def forward(self, input_values):
+        output = self.model.wav2vec2(input_values=input_values)
+        return self.model.lm_head(output[0])
+
+
+class RBLNWav2Vec2ForCTC(RBLNModel):
+    """
+    Wav2Vec2 Model with a `language modeling` head on top for Connectionist Temporal Classification (CTC).
+
+    This model inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model.
+
+    It implements the methods to convert a pre-trained Wav2Vec2 model into a RBLN Wav2Vec2 model by:
+    - transferring the checkpoint weights of the original into an optimized RBLN graph,
+    - compiling the resulting graph using the RBLN compiler.
+    """
+
+    model_type = "rbln_model"
+    main_input_name = "input_values"
+    auto_model_class = AutoModelForMaskedLM
+
+    @classmethod
+    def wrap_model_if_needed(cls, model: torch.nn.Module) -> torch.nn.Module:
+        return _Wav2Vec2(model).eval()
+
+    @classmethod
+    def _get_rbln_config(
+        cls,
+        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
+        model_config: "PretrainedConfig",
+        rbln_max_seq_len: Optional[int] = None,
+        rbln_batch_size: Optional[int] = None,
+    ) -> RBLNConfig:
+        meta = {}
+
+        if rbln_max_seq_len is None:
+            for tokenizer in preprocessors:
+                if hasattr(tokenizer, "model_max_length"):
+                    rbln_max_seq_len = tokenizer.model_max_length
+                    break
+            if rbln_max_seq_len is None:
+                raise ValueError("`rbln_max_seq_len` should be specified!")
+
+        meta["rbln_max_seq_len"] = rbln_max_seq_len
+
+        if rbln_batch_size is None:
+            rbln_batch_size = 1
+
+        input_info = [
+            (
+                "input_values",
+                [
+                    rbln_batch_size,
+                    rbln_max_seq_len,
+                ],
+                "float32",
+            ),
+        ]
+
+        rbln_runtime_config = RBLNRuntimeConfig(input_info=input_info, batch_size=rbln_batch_size)
+
+        rbln_config = RBLNConfig.from_rbln_runtime_configs(
+            [rbln_runtime_config],
+            _rbln_meta=meta,
+        )
+
+        return rbln_config
+
+    def forward(self, input_values: "torch.Tensor", **kwargs):
+        outputs = super().forward(input_values, **kwargs)
+        return CausalLMOutput(logits=outputs)

optimum/rbln/transformers/models/whisper/__init__.py

@@ -0,0 +1,24 @@
+# Copyright 2024 Rebellions Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+from .modeling_whisper import RBLNWhisperForConditionalGeneration