optimum-rbln 0.1.7__py3-none-any.whl → 0.1.8__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (39) hide show
  1. optimum/rbln/__init__.py +14 -0
  2. optimum/rbln/__version__.py +1 -1
  3. optimum/rbln/diffusers/__init__.py +0 -1
  4. optimum/rbln/diffusers/models/controlnet.py +3 -0
  5. optimum/rbln/diffusers/models/unet_2d_condition.py +2 -2
  6. optimum/rbln/diffusers/pipelines/controlnet/multicontrolnet.py +22 -144
  7. optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet.py +107 -59
  8. optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py +106 -54
  9. optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py +130 -71
  10. optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py +131 -72
  11. optimum/rbln/modeling_alias.py +14 -0
  12. optimum/rbln/modeling_base.py +110 -0
  13. optimum/rbln/transformers/__init__.py +6 -0
  14. optimum/rbln/transformers/cache_utils.py +111 -0
  15. optimum/rbln/transformers/generation/utils.py +0 -2
  16. optimum/rbln/transformers/models/__init__.py +2 -0
  17. optimum/rbln/transformers/models/bart/bart_architecture.py +0 -5
  18. optimum/rbln/transformers/models/decoderonly/__init__.py +36 -0
  19. optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py +515 -0
  20. optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py +349 -0
  21. optimum/rbln/transformers/models/gemma/__init__.py +24 -0
  22. optimum/rbln/transformers/models/gemma/gemma_architecture.py +116 -0
  23. optimum/rbln/transformers/models/gemma/modeling_gemma.py +61 -0
  24. optimum/rbln/transformers/models/gpt2/gpt2_architecture.py +201 -166
  25. optimum/rbln/transformers/models/gpt2/modeling_gpt2.py +56 -220
  26. optimum/rbln/transformers/models/llama/llama_architecture.py +3 -610
  27. optimum/rbln/transformers/models/llama/modeling_llama.py +8 -442
  28. optimum/rbln/transformers/models/midm/hf_hub_cached/midm_bitext_tokenization.py +2 -1
  29. optimum/rbln/transformers/models/midm/hf_hub_cached/modeling_midm.py +0 -4
  30. optimum/rbln/transformers/models/midm/midm_architecture.py +160 -357
  31. optimum/rbln/transformers/models/midm/modeling_midm.py +40 -272
  32. optimum/rbln/transformers/models/whisper/whisper_architecture.py +0 -6
  33. optimum/rbln/transformers/models/xlm_roberta/__init__.py +24 -0
  34. optimum/rbln/transformers/models/xlm_roberta/modeling_xlm_roberta.py +125 -0
  35. {optimum_rbln-0.1.7.dist-info → optimum_rbln-0.1.8.dist-info}/METADATA +2 -3
  36. {optimum_rbln-0.1.7.dist-info → optimum_rbln-0.1.8.dist-info}/RECORD +38 -30
  37. optimum/rbln/transformers/models/llama/llama_architecture_cb.py +0 -764
  38. {optimum_rbln-0.1.7.dist-info → optimum_rbln-0.1.8.dist-info}/WHEEL +0 -0
  39. {optimum_rbln-0.1.7.dist-info → optimum_rbln-0.1.8.dist-info}/licenses/LICENSE +0 -0
@@ -0,0 +1,349 @@
1
+ # Copyright 2024 Rebellions Inc.
2
+
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at:
6
+
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ # Portions of this software are licensed under the Apache License,
16
+ # Version 2.0. See the NOTICE file distributed with this work for
17
+ # additional information regarding copyright ownership.
18
+
19
+ # All other portions of this software, including proprietary code,
20
+ # are the intellectual property of Rebellions Inc. and may not be
21
+ # copied, modified, or distributed without prior written permission
22
+ # from Rebellions Inc.
23
+ import logging
24
+ from abc import ABC, abstractmethod
25
+ from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
26
+
27
+ import rebel # noqa: F401
28
+ import torch # noqa: F401
29
+ from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
30
+ from transformers.modeling_outputs import CausalLMOutputWithPast
31
+
32
+ from ....modeling_base import RBLNModel
33
+ from ....modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNConfig, RBLNRuntimeConfig
34
+ from ....utils.runtime_utils import RBLNPytorchRuntime
35
+
36
+
37
+ logger = logging.getLogger(__name__)
38
+
39
+ if TYPE_CHECKING:
40
+ from transformers import (
41
+ AutoFeatureExtractor,
42
+ AutoProcessor,
43
+ AutoTokenizer,
44
+ PretrainedConfig,
45
+ )
46
+
47
+
48
+ class RBLNRuntimeModel(RBLNPytorchRuntime):
49
+ mandatory_members = ["main_input_name"]
50
+
51
+
52
+ class RBLNDecoderOnlyModelForCausalLM(RBLNModel, ABC):
53
+ """
54
+ The DecoderOnly Model transformer with a language modeling head (linear layer) on top.
55
+ This model inherits from [`RBLNMultiModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
56
+
57
+ A class to convert and run pre-trained transformers based DecoderOnlyForCausalLM model on RBLN devices.
58
+ It implements the methods to convert a pre-trained transformers DecoderOnlyForCausalLM model into a RBLN transformer model by:
59
+ - transferring the checkpoint weights of the original into an optimized RBLN graph,
60
+ - compiling the resulting graph using the RBLN compiler.
61
+ """
62
+
63
+ main_input_name = "input_ids"
64
+ auto_model_class = AutoModelForCausalLM
65
+
66
+ def __post_init__(self, **kwargs):
67
+ self.batch_size = self.rbln_config.meta["rbln_batch_size"]
68
+ self.max_seq_len = self.rbln_config.meta["rbln_max_seq_len"]
69
+ self.prefill_chunk_size = self.rbln_config.meta["rbln_prefill_chunk_size"]
70
+
71
+ self.prefill_attention_mask = torch.zeros(1, 1, self.prefill_chunk_size, self.max_seq_len, dtype=torch.int64)
72
+ self.causal_mask = 1 - torch.triu(
73
+ torch.ones(1, 1, self.prefill_chunk_size, self.prefill_chunk_size), diagonal=1
74
+ )
75
+ self.dec_attn_mask_init = torch.zeros(1, 1, 1, self.max_seq_len, dtype=torch.int64)
76
+ self.dec_attn_mask = torch.zeros(self.batch_size, 1, 1, self.max_seq_len, dtype=torch.int64)
77
+ self.prefill_decoder = RBLNRuntimeModel(runtime=self.model[0], main_input_name="input_ids")
78
+ self.decoder = RBLNRuntimeModel(runtime=self.model[1], main_input_name="input_ids")
79
+
80
+ @classmethod
81
+ @abstractmethod
82
+ def wrapping_torch_model(self, model: "PreTrainedModel", rbln_max_seq_len: int):
83
+ pass
84
+
85
+ @classmethod
86
+ @torch.inference_mode()
87
+ def get_compiled_model(cls, model: "PreTrainedModel", rbln_config: RBLNConfig):
88
+ wrapped_model = cls.wrapping_torch_model(model, rbln_config.meta["rbln_max_seq_len"])
89
+
90
+ prefill_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][0]
91
+ dec_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][1]
92
+
93
+ prefill_example_inputs = prefill_rbln_runtime_config.get_dummy_inputs(fill=0)
94
+ dec_example_inputs = dec_rbln_runtime_config.get_dummy_inputs(fill=4)
95
+
96
+ batch_index = 3
97
+ dec_example_inputs[batch_index].fill_(-1) # fill batch_position -1 to indicate it is decoder.
98
+
99
+ prefill_scripted_model = torch.jit.trace(wrapped_model, prefill_example_inputs, check_trace=False)
100
+ dec_scripted_model = torch.jit.trace(wrapped_model, dec_example_inputs, check_trace=False)
101
+
102
+ prefill_ir = rebel.torchscript_to_ir(
103
+ prefill_scripted_model,
104
+ input_names=[v[0] for v in prefill_rbln_runtime_config.input_info],
105
+ )
106
+ dec_ir = rebel.torchscript_to_ir(
107
+ dec_scripted_model,
108
+ input_names=[v[0] for v in dec_rbln_runtime_config.input_info],
109
+ )
110
+
111
+ # Caching prefill_decoder/decoder I/O
112
+ cache_index_offset = 4
113
+ connections = [
114
+ (prefill_ir.outputs[1 + i], prefill_ir.inputs[cache_index_offset + i])
115
+ for i in range(model.config.num_hidden_layers * 2)
116
+ ]
117
+
118
+ compiled_model = rebel.compile(
119
+ prefill_ir,
120
+ dec_ir,
121
+ connections=connections,
122
+ fusion=prefill_rbln_runtime_config.fusion,
123
+ npu=prefill_rbln_runtime_config.npu,
124
+ tensor_parallel_size=prefill_rbln_runtime_config.tensor_parallel_size,
125
+ use_weight_sharing=True,
126
+ )
127
+ return compiled_model
128
+
129
+ @classmethod
130
+ def _get_rbln_config(
131
+ cls,
132
+ preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
133
+ model_config: "PretrainedConfig",
134
+ rbln_max_seq_len: Optional[int] = None,
135
+ rbln_batch_size: Optional[int] = None,
136
+ **kwargs,
137
+ ) -> RBLNConfig:
138
+ meta = {}
139
+
140
+ prefill_chunk_size = 128
141
+ if rbln_max_seq_len is None:
142
+ rbln_max_seq_len = getattr(model_config, "max_position_embeddings", None)
143
+ rbln_batch_size = 1 if rbln_batch_size is None else rbln_batch_size
144
+
145
+ meta["rbln_max_seq_len"] = rbln_max_seq_len
146
+ meta["rbln_batch_size"] = rbln_batch_size
147
+ meta["rbln_prefill_chunk_size"] = prefill_chunk_size
148
+
149
+ def get_input_info(
150
+ batch_size,
151
+ query_length,
152
+ ):
153
+ head_dim = (
154
+ model_config.head_dim
155
+ if hasattr(model_config, "head_dim")
156
+ else model_config.hidden_size // model_config.num_attention_heads
157
+ )
158
+ input_info = [
159
+ ("input_ids", [batch_size, query_length], "int64"),
160
+ ("attention_mask", [batch_size, 1, query_length, rbln_max_seq_len], "int64"),
161
+ (
162
+ "cache_position",
163
+ [batch_size, query_length],
164
+ "int32",
165
+ ),
166
+ ("batch_position", [], "int16"),
167
+ ]
168
+
169
+ input_info.extend(
170
+ [
171
+ (
172
+ f"past_key_values_{i}",
173
+ [
174
+ rbln_batch_size,
175
+ model_config.num_key_value_heads,
176
+ rbln_max_seq_len,
177
+ head_dim,
178
+ ],
179
+ "float32",
180
+ )
181
+ for i in range(model_config.num_hidden_layers * 2)
182
+ ]
183
+ )
184
+
185
+ return input_info
186
+
187
+ prefill_input_info = get_input_info(
188
+ batch_size=1,
189
+ query_length=prefill_chunk_size,
190
+ )
191
+ dec_input_info = get_input_info(
192
+ batch_size=rbln_batch_size,
193
+ query_length=1,
194
+ )
195
+
196
+ prefill_rbln_runtime_config = RBLNRuntimeConfig(input_info=prefill_input_info)
197
+ dec_rbln_runtime_config = RBLNRuntimeConfig(input_info=dec_input_info)
198
+
199
+ dec_rbln_runtime_config.batch_size = rbln_batch_size
200
+
201
+ rbln_config = RBLNConfig.from_rbln_runtime_configs(
202
+ [prefill_rbln_runtime_config, dec_rbln_runtime_config],
203
+ _rbln_meta=meta,
204
+ )
205
+
206
+ return rbln_config
207
+
208
+ @classmethod
209
+ def _create_runtimes(
210
+ cls, compiled_models: List[rebel.RBLNCompiledModel], rbln_device_map: Dict[str, int]
211
+ ) -> List[rebel.Runtime]:
212
+ device_val = rbln_device_map[DEFAULT_COMPILED_MODEL_NAME]
213
+ return [
214
+ compiled_models[0].create_runtime(input_info_index=0, tensor_type="pt", device=device_val),
215
+ compiled_models[0].create_runtime(input_info_index=1, tensor_type="pt", device=device_val),
216
+ ]
217
+
218
+ def get_decoder(self):
219
+ return self.decoder
220
+
221
+ def can_generate(self):
222
+ return True
223
+
224
+ def _reorder_cache(self, past_key_values, beam_idx):
225
+ raise NotImplementedError
226
+
227
+ # args input_ids, past_key_values and attention_mask are updated by _update_model_kwargs_for_generation() in _greedy_search() in GenerationMixin
228
+ def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs):
229
+ batch_size = input_ids.shape[0]
230
+
231
+ # FIXME past_key_values is just carriier variable for past_cached_length
232
+ # torch.tensor((4,1),dtype=torch.int32) which refers a past_cached_length of each batch
233
+ past_cached_length = past_key_values
234
+ if past_cached_length is None:
235
+ l_input_ids = []
236
+ cache_positions = []
237
+ past_cached_length = torch.zeros((batch_size, 1), dtype=torch.int32)
238
+ for i in range(batch_size):
239
+ input_id = input_ids[i]
240
+ input_id = input_id[attention_mask[i] == 1]
241
+ valid_len = input_id.shape[-1]
242
+ cache_position = torch.arange(0, valid_len, dtype=torch.int32)
243
+ past_cached_length[i] = valid_len
244
+ l_input_ids.append(input_id.unsqueeze(0))
245
+ cache_positions.append(cache_position.unsqueeze(0))
246
+
247
+ input_ids = l_input_ids
248
+ else:
249
+ input_ids = input_ids[:, -1:]
250
+ cache_positions = past_cached_length
251
+ past_cached_length = past_cached_length + 1
252
+
253
+ model_inputs = {
254
+ "input_ids": input_ids,
255
+ "cache_position": cache_positions,
256
+ "past_cached_length": past_cached_length,
257
+ }
258
+
259
+ return model_inputs
260
+
261
+ def forward(
262
+ self,
263
+ input_ids: torch.LongTensor = None,
264
+ cache_position: Union[List[torch.Tensor], torch.Tensor] = None, # vllm keyword argument
265
+ batch_idx: Optional[int] = None,
266
+ past_cached_length: Optional[torch.Tensor] = None, # past_cached_length
267
+ **kwargs,
268
+ ) -> Tuple[torch.FloatTensor]:
269
+ # prefll & hf generate
270
+ if isinstance(cache_position, list):
271
+ logits = []
272
+ for batch_idx, (input_id, cache_pos) in enumerate(zip(input_ids, cache_position)):
273
+ logit = self._forward_prefill(input_ids=input_id, cache_position=cache_pos, batch_idx=batch_idx)
274
+ logits.append(logit)
275
+ logits = torch.cat(logits, dim=0)
276
+ # prefill & vllm step
277
+ elif cache_position.shape[-1] > 1:
278
+ logits = self._forward_prefill(input_ids=input_ids, cache_position=cache_position, batch_idx=batch_idx)
279
+ # common decoder
280
+ else:
281
+ logits = self._forward_decoder(input_ids=input_ids, cache_position=cache_position)
282
+
283
+ return CausalLMOutputWithPast(
284
+ logits=logits,
285
+ past_key_values=past_cached_length, # past_cached_length
286
+ )
287
+
288
+ def _forward_prefill(
289
+ self,
290
+ input_ids: torch.LongTensor = None,
291
+ cache_position: torch.Tensor = None, # torch.tensor(,dtype=int32) (1,64) // (4,1)
292
+ batch_idx: int = None,
293
+ ) -> torch.FloatTensor:
294
+ if batch_idx is None or batch_idx >= self.batch_size:
295
+ raise RuntimeError(
296
+ f"Invalid batch_idx ({batch_idx}). It must be a non-null value less than the batch size ({self.batch_size})."
297
+ )
298
+ query_length = input_ids.shape[1]
299
+ attention_mask = self.prefill_attention_mask.clone()
300
+ for step in range(0, query_length, self.prefill_chunk_size):
301
+ if step + self.prefill_chunk_size > query_length:
302
+ input_ids = torch.nn.functional.pad(input_ids, (0, step + self.prefill_chunk_size - query_length))
303
+ cache_position = torch.cat(
304
+ [
305
+ cache_position,
306
+ torch.arange(
307
+ query_length,
308
+ step + self.prefill_chunk_size,
309
+ dtype=torch.int32,
310
+ ).unsqueeze(0),
311
+ ],
312
+ dim=-1,
313
+ )
314
+
315
+ sliced_input_ids = input_ids[:, step : step + self.prefill_chunk_size]
316
+ sliced_cache_positions = cache_position[:, step : step + self.prefill_chunk_size]
317
+ attention_mask[:, :, :, :step] = 1
318
+ attention_mask[:, :, :, step : step + self.prefill_chunk_size] = self.causal_mask
319
+
320
+ logits, _ = self.prefill_decoder(
321
+ sliced_input_ids.contiguous(),
322
+ attention_mask.contiguous(),
323
+ sliced_cache_positions.contiguous(),
324
+ torch.tensor(batch_idx, dtype=torch.int16),
325
+ )
326
+ logits = logits[:, query_length % self.prefill_chunk_size - 1].unsqueeze(1)
327
+
328
+ self.dec_attn_mask[batch_idx] = self.dec_attn_mask_init.clone()
329
+ self.dec_attn_mask[batch_idx, :, :, :query_length] = 1
330
+
331
+ return logits
332
+
333
+ def _forward_decoder(
334
+ self, input_ids: torch.LongTensor = None, cache_position: torch.Tensor = None
335
+ ) -> torch.FloatTensor:
336
+ batch_size = input_ids.shape[0]
337
+
338
+ for b_idx in range(batch_size):
339
+ decoding_step = cache_position[b_idx].item()
340
+ self.dec_attn_mask[b_idx, :, :, decoding_step] = 1
341
+
342
+ logits, _ = self.decoder(
343
+ input_ids.contiguous(),
344
+ self.dec_attn_mask.contiguous(),
345
+ cache_position.contiguous(),
346
+ torch.tensor(0, dtype=torch.int16),
347
+ )
348
+
349
+ return logits
@@ -0,0 +1,24 @@
1
+ # Copyright 2024 Rebellions Inc.
2
+
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at:
6
+
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ # Portions of this software are licensed under the Apache License,
16
+ # Version 2.0. See the NOTICE file distributed with this work for
17
+ # additional information regarding copyright ownership.
18
+
19
+ # All other portions of this software, including proprietary code,
20
+ # are the intellectual property of Rebellions Inc. and may not be
21
+ # copied, modified, or distributed without prior written permission
22
+ # from Rebellions Inc.
23
+
24
+ from .modeling_gemma import RBLNGemmaForCausalLM
@@ -0,0 +1,116 @@
1
+ # Copyright 2024 Rebellions Inc.
2
+
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at:
6
+
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ # Portions of this software are licensed under the Apache License,
16
+ # Version 2.0. See the NOTICE file distributed with this work for
17
+ # additional information regarding copyright ownership.
18
+
19
+ # All other portions of this software, including proprietary code,
20
+ # are the intellectual property of Rebellions Inc. and may not be
21
+ # copied, modified, or distributed without prior written permission
22
+ # from Rebellions Inc.
23
+
24
+ from typing import Dict, List, Optional, Tuple, Union
25
+
26
+ import torch
27
+ from transformers.modeling_outputs import (
28
+ BaseModelOutputWithPast,
29
+ )
30
+
31
+ from ...models.decoderonly import (
32
+ DecoderOnlyAttention,
33
+ DecoderOnlyDecoderLayer,
34
+ DecoderOnlyWrapper,
35
+ slice_and_unsqueeze_cos_sin,
36
+ )
37
+
38
+
39
+ class GemmaWrapper(DecoderOnlyWrapper):
40
+ def get_forward_dict(self):
41
+ forward_dict = {}
42
+ forward_dict.update({"wrapper": GemmaModel.forward, "model": DecoderOnlyDecoderLayer.forward, "decoder_layer": DecoderOnlyAttention.forward,})
43
+ return forward_dict
44
+
45
+ class GemmaModel:
46
+ def forward(
47
+ self,
48
+ input_ids: torch.LongTensor = None,
49
+ attention_mask: Optional[torch.Tensor] = None,
50
+ position_ids: Optional[torch.LongTensor] = None,
51
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
52
+ batch_ids: Optional[torch.LongTensor] = None,
53
+ inputs_embeds: Optional[torch.FloatTensor] = None,
54
+ use_cache: Optional[bool] = True,
55
+ output_attentions: Optional[bool] = False,
56
+ output_hidden_states: Optional[bool] = False,
57
+ forward_dict : Optional[Dict[str, classmethod]] = None,
58
+ rotary_pos_emb=None,
59
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
60
+ # embed positions
61
+ inputs_embeds = self.embed_tokens(input_ids)
62
+ hidden_states = inputs_embeds
63
+
64
+ ##### GEMMA change from llama#####
65
+ hidden_states = hidden_states * (self.config.hidden_size**0.5)
66
+
67
+ attention_mask = (1 - attention_mask) * torch.finfo(torch.float16).min
68
+
69
+ # get cos,sin vector
70
+ cos, sin = rotary_pos_emb(inputs_embeds, attention_mask.shape[-1])
71
+ cos, sin = slice_and_unsqueeze_cos_sin(cos, sin, position_ids)
72
+
73
+ # decoder layers
74
+ all_hidden_states = () if output_hidden_states else None
75
+ all_self_attns = () if output_attentions else None
76
+
77
+ for layer_idx, decoder_layer in enumerate(self.layers):
78
+ if output_hidden_states:
79
+ all_hidden_states += (hidden_states,)
80
+ layer_outputs = forward_dict["model"](
81
+ decoder_layer,
82
+ hidden_states,
83
+ layer_idx,
84
+ attention_mask=attention_mask,
85
+ position_ids=position_ids,
86
+ past_key_value=past_key_values,
87
+ output_attentions=output_attentions,
88
+ use_cache=use_cache,
89
+ batch_ids=batch_ids,
90
+ cos=cos,
91
+ sin=sin,
92
+ forward_dict=forward_dict
93
+ )
94
+
95
+ hidden_states = layer_outputs[0]
96
+
97
+ updated_cache = layer_outputs[2 if output_attentions else 1]
98
+
99
+ if output_attentions:
100
+ all_self_attns += (layer_outputs[1],)
101
+
102
+ hidden_states = self.norm(hidden_states)
103
+
104
+ # add hidden states from the last decoder layer
105
+ if output_hidden_states:
106
+ all_hidden_states += (hidden_states,)
107
+
108
+ # convert RebelDynamicCache to legacy Tuple[Tuple[torch.Tensor]]
109
+ next_cache = updated_cache.to_legacy_cache()
110
+
111
+ return BaseModelOutputWithPast(
112
+ last_hidden_state=hidden_states,
113
+ past_key_values=next_cache,
114
+ hidden_states=all_hidden_states,
115
+ attentions=all_self_attns,
116
+ )
@@ -0,0 +1,61 @@
1
+ # Copyright 2024 Rebellions Inc.
2
+
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at:
6
+
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ # Portions of this software are licensed under the Apache License,
16
+ # Version 2.0. See the NOTICE file distributed with this work for
17
+ # additional information regarding copyright ownership.
18
+
19
+ # All other portions of this software, including proprietary code,
20
+ # are the intellectual property of Rebellions Inc. and may not be
21
+ # copied, modified, or distributed without prior written permission
22
+ # from Rebellions Inc.
23
+
24
+ import inspect
25
+ import logging
26
+ from typing import Any, Callable
27
+
28
+ from transformers import GemmaForCausalLM, PreTrainedModel
29
+
30
+ from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
31
+ from .gemma_architecture import GemmaWrapper
32
+
33
+
34
+ logger = logging.getLogger(__name__)
35
+
36
+
37
+ class RBLNGemmaForCausalLM(RBLNDecoderOnlyModelForCausalLM):
38
+ """
39
+ The Gemma Model transformer with a language modeling head (linear layer) on top.
40
+ This model inherits from [`RBLNMultiModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
41
+
42
+ A class to convert and run pre-trained transformers based GemmaForCausalLM model on RBLN devices.
43
+ It implements the methods to convert a pre-trained transformers GemmaForCausalLM model into a RBLN transformer model by:
44
+ - transferring the checkpoint weights of the original into an optimized RBLN graph,
45
+ - compiling the resulting graph using the RBLN compiler.
46
+ """
47
+
48
+ @classmethod
49
+ def wrapping_torch_model(self, model: "PreTrainedModel", rbln_max_seq_len: int):
50
+ return GemmaWrapper(model, rbln_max_seq_len).eval()
51
+
52
+ def __getattr__(self, __name: str) -> Any:
53
+ def redirect(func):
54
+ return lambda *pargs, **kwargs: func(self, *pargs, **kwargs)
55
+
56
+ val = getattr(GemmaForCausalLM, __name)
57
+
58
+ if isinstance(val, Callable) and "self" in set(inspect.signature(val).parameters):
59
+ return redirect(val)
60
+
61
+ return val