optimum-rbln 0.1.7__py3-none-any.whl → 0.1.9__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 (49) hide show
  1. optimum/rbln/__init__.py +17 -0
  2. optimum/rbln/__version__.py +1 -1
  3. optimum/rbln/diffusers/__init__.py +0 -1
  4. optimum/rbln/diffusers/models/autoencoder_kl.py +3 -3
  5. optimum/rbln/diffusers/models/controlnet.py +7 -3
  6. optimum/rbln/diffusers/models/unet_2d_condition.py +5 -5
  7. optimum/rbln/diffusers/pipelines/controlnet/multicontrolnet.py +23 -146
  8. optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet.py +107 -59
  9. optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py +106 -54
  10. optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py +130 -71
  11. optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py +131 -72
  12. optimum/rbln/modeling_alias.py +19 -1
  13. optimum/rbln/modeling_base.py +162 -18
  14. optimum/rbln/transformers/__init__.py +8 -0
  15. optimum/rbln/transformers/cache_utils.py +111 -0
  16. optimum/rbln/transformers/generation/utils.py +0 -2
  17. optimum/rbln/transformers/models/__init__.py +3 -0
  18. optimum/rbln/transformers/models/bart/bart_architecture.py +0 -5
  19. optimum/rbln/transformers/models/clip/modeling_clip.py +1 -1
  20. optimum/rbln/transformers/models/decoderonly/__init__.py +36 -0
  21. optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py +516 -0
  22. optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py +464 -0
  23. optimum/rbln/transformers/models/gemma/__init__.py +24 -0
  24. optimum/rbln/transformers/models/gemma/gemma_architecture.py +123 -0
  25. optimum/rbln/transformers/models/gemma/modeling_gemma.py +67 -0
  26. optimum/rbln/transformers/models/gpt2/gpt2_architecture.py +201 -166
  27. optimum/rbln/transformers/models/gpt2/modeling_gpt2.py +10 -257
  28. optimum/rbln/transformers/models/llama/llama_architecture.py +3 -610
  29. optimum/rbln/transformers/models/llama/modeling_llama.py +12 -440
  30. optimum/rbln/transformers/models/midm/hf_hub_cached/midm_bitext_tokenization.py +2 -1
  31. optimum/rbln/transformers/models/midm/hf_hub_cached/modeling_midm.py +0 -4
  32. optimum/rbln/transformers/models/midm/midm_architecture.py +160 -357
  33. optimum/rbln/transformers/models/midm/modeling_midm.py +10 -325
  34. optimum/rbln/transformers/models/mistral/__init__.py +24 -0
  35. optimum/rbln/transformers/models/mistral/mistral_architecture.py +29 -0
  36. optimum/rbln/transformers/models/mistral/modeling_mistral.py +68 -0
  37. optimum/rbln/transformers/models/wav2vec2/modeling_wav2vec2.py +1 -1
  38. optimum/rbln/transformers/models/whisper/whisper_architecture.py +0 -6
  39. optimum/rbln/transformers/models/xlm_roberta/__init__.py +24 -0
  40. optimum/rbln/transformers/models/xlm_roberta/modeling_xlm_roberta.py +131 -0
  41. optimum/rbln/transformers/utils/__init__.py +0 -0
  42. optimum/rbln/transformers/utils/rbln_quantization.py +109 -0
  43. optimum/rbln/utils/import_utils.py +1 -4
  44. optimum/rbln/utils/runtime_utils.py +2 -1
  45. {optimum_rbln-0.1.7.dist-info → optimum_rbln-0.1.9.dist-info}/METADATA +11 -5
  46. {optimum_rbln-0.1.7.dist-info → optimum_rbln-0.1.9.dist-info}/RECORD +48 -35
  47. optimum/rbln/transformers/models/llama/llama_architecture_cb.py +0 -764
  48. {optimum_rbln-0.1.7.dist-info → optimum_rbln-0.1.9.dist-info}/WHEEL +0 -0
  49. {optimum_rbln-0.1.7.dist-info → optimum_rbln-0.1.9.dist-info}/licenses/LICENSE +0 -0
@@ -0,0 +1,464 @@
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 glob
24
+ import logging
25
+ from abc import ABC
26
+ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
27
+
28
+ import rebel # noqa: F401
29
+ import torch # noqa: F401
30
+ from safetensors.torch import load_file
31
+ from transformers import AutoConfig, AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
32
+ from transformers.modeling_outputs import CausalLMOutputWithPast
33
+ from transformers.modeling_utils import no_init_weights
34
+
35
+ from ....modeling_base import RBLNModel
36
+ from ....modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNConfig, RBLNRuntimeConfig
37
+ from ....utils.runtime_utils import RBLNPytorchRuntime
38
+ from ...utils.rbln_quantization import replace_quantized_linear_layers
39
+
40
+
41
+ logger = logging.getLogger(__name__)
42
+
43
+ if TYPE_CHECKING:
44
+ from transformers import (
45
+ AutoFeatureExtractor,
46
+ AutoProcessor,
47
+ AutoTokenizer,
48
+ PretrainedConfig,
49
+ )
50
+
51
+ SUPPORTED_QUANTIZATIONS = {
52
+ "rbln": [
53
+ "w4a16",
54
+ ],
55
+ }
56
+
57
+
58
+ class RBLNRuntimeModel(RBLNPytorchRuntime):
59
+ mandatory_members = ["main_input_name"]
60
+
61
+
62
+ class RBLNDecoderOnlyModelForCausalLM(RBLNModel, ABC):
63
+ """
64
+ The DecoderOnly Model transformer with a language modeling head (linear layer) on top.
65
+ This model inherits from [`RBLNMultiModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
66
+
67
+ A class to convert and run pre-trained transformers based DecoderOnlyForCausalLM model on RBLN devices.
68
+ It implements the methods to convert a pre-trained transformers DecoderOnlyForCausalLM model into a RBLN transformer model by:
69
+ - transferring the checkpoint weights of the original into an optimized RBLN graph,
70
+ - compiling the resulting graph using the RBLN compiler.
71
+ """
72
+
73
+ main_input_name = "input_ids"
74
+ auto_model_class = AutoModelForCausalLM
75
+
76
+ def __post_init__(self, **kwargs):
77
+ self.batch_size = self.rbln_config.meta["rbln_batch_size"]
78
+ self.max_seq_len = self.rbln_config.meta["rbln_max_seq_len"]
79
+ self.prefill_chunk_size = self.rbln_config.meta["rbln_prefill_chunk_size"]
80
+
81
+ self.prefill_attention_mask = torch.zeros(1, 1, self.prefill_chunk_size, self.max_seq_len, dtype=torch.int64)
82
+ self.causal_mask = 1 - torch.triu(
83
+ torch.ones(1, 1, self.prefill_chunk_size, self.prefill_chunk_size), diagonal=1
84
+ )
85
+ self.dec_attn_mask_init = torch.zeros(1, 1, 1, self.max_seq_len, dtype=torch.int64)
86
+ self.dec_attn_mask = torch.zeros(self.batch_size, 1, 1, self.max_seq_len, dtype=torch.int64)
87
+ self.prefill_decoder = RBLNRuntimeModel(runtime=self.model[0], main_input_name="input_ids")
88
+ self.decoder = RBLNRuntimeModel(runtime=self.model[1], main_input_name="input_ids")
89
+
90
+ @classmethod
91
+ def get_quantized_model(
92
+ cls,
93
+ model_id: str,
94
+ use_auth_token: Optional[Union[bool, str]] = None,
95
+ revision: Optional[str] = None,
96
+ force_download: bool = False,
97
+ cache_dir: Optional[str] = None,
98
+ subfolder: str = "",
99
+ local_files_only: bool = False,
100
+ trust_remote_code: bool = False,
101
+ rbln_config_kwargs: Optional[Dict[str, Any]] = None,
102
+ rbln_constructor_kwargs: Optional[Dict[str, Any]] = None,
103
+ **kwargs,
104
+ ):
105
+ kwargs = cls.update_kwargs(kwargs)
106
+
107
+ config = AutoConfig.from_pretrained(
108
+ model_id,
109
+ use_auth_token=use_auth_token,
110
+ revision=revision,
111
+ force_download=force_download,
112
+ cache_dir=cache_dir,
113
+ trust_remote_code=trust_remote_code,
114
+ **kwargs,
115
+ )
116
+
117
+ with no_init_weights():
118
+ model = AutoModelForCausalLM.from_config(config)
119
+ replace_quantized_linear_layers(model)
120
+
121
+ state_dict = {}
122
+ for safetensor_file in glob.glob(f"{model_id}/*.safetensors"):
123
+ partial_state_dict = load_file(safetensor_file)
124
+ state_dict.update(partial_state_dict)
125
+
126
+ n_layer = kwargs.get("num_hidden_layers", None)
127
+ if n_layer is not None:
128
+ keys_to_delete = []
129
+ for key in state_dict.keys():
130
+ parts = key.split(".")
131
+ if len(parts) > 2 and parts[2].isdigit():
132
+ layer_num = int(parts[2])
133
+ if layer_num >= n_layer:
134
+ keys_to_delete.append(key)
135
+
136
+ for key in keys_to_delete:
137
+ del state_dict[key]
138
+
139
+ model.load_state_dict(state_dict)
140
+ return model
141
+
142
+ @classmethod
143
+ def get_pytorch_model(
144
+ cls,
145
+ *args,
146
+ **kwargs,
147
+ ) -> "PreTrainedModel":
148
+ rbln_config_kwargs = kwargs.get("rbln_config_kwargs", {})
149
+ rbln_quantization = rbln_config_kwargs.get("rbln_quantization", None)
150
+
151
+ if rbln_quantization is not None and rbln_quantization["format"] == "rbln":
152
+ model = cls.get_quantized_model(*args, **kwargs)
153
+ else:
154
+ model = super().get_pytorch_model(*args, **kwargs)
155
+
156
+ return model
157
+
158
+ @classmethod
159
+ @torch.inference_mode()
160
+ def get_compiled_model(cls, model: "PreTrainedModel", rbln_config: RBLNConfig):
161
+ wrapped_model = cls.wrap_model_if_needed(model, rbln_config)
162
+
163
+ prefill_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][0]
164
+ dec_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][1]
165
+
166
+ def get_scripted_model():
167
+ # This function is nested to dealloc the example inputs before compilation.
168
+ prefill_example_inputs = prefill_rbln_runtime_config.get_dummy_inputs(fill=0)
169
+ dec_example_inputs = dec_rbln_runtime_config.get_dummy_inputs(fill=4)
170
+
171
+ batch_index = 3
172
+ dec_example_inputs[batch_index].fill_(-1) # fill batch_position -1 to indicate it is decoder.
173
+
174
+ prefill_scripted_model = torch.jit.trace(
175
+ wrapped_model, prefill_example_inputs, check_trace=False, _store_inputs=False
176
+ )
177
+ dec_scripted_model = torch.jit.trace(
178
+ wrapped_model, dec_example_inputs, check_trace=False, _store_inputs=False
179
+ )
180
+ return prefill_scripted_model, dec_scripted_model
181
+
182
+ prefill_scripted_model, dec_scripted_model = get_scripted_model()
183
+
184
+ prefill_ir = rebel.torchscript_to_ir(
185
+ prefill_scripted_model,
186
+ input_names=[v[0] for v in prefill_rbln_runtime_config.input_info],
187
+ )
188
+ dec_ir = rebel.torchscript_to_ir(
189
+ dec_scripted_model,
190
+ input_names=[v[0] for v in dec_rbln_runtime_config.input_info],
191
+ )
192
+
193
+ # Caching prefill_decoder/decoder I/O
194
+ cache_index_offset = 4
195
+ connections = [
196
+ (prefill_ir.outputs[1 + i], prefill_ir.inputs[cache_index_offset + i])
197
+ for i in range(model.config.num_hidden_layers * 2)
198
+ ]
199
+
200
+ compiled_model = rebel.compile(
201
+ prefill_ir,
202
+ dec_ir,
203
+ connections=connections,
204
+ fusion=prefill_rbln_runtime_config.fusion,
205
+ npu=prefill_rbln_runtime_config.npu,
206
+ tensor_parallel_size=prefill_rbln_runtime_config.tensor_parallel_size,
207
+ use_weight_sharing=True,
208
+ )
209
+ return compiled_model
210
+
211
+ @classmethod
212
+ def _get_rbln_config(
213
+ cls,
214
+ preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
215
+ model_config: "PretrainedConfig",
216
+ rbln_max_seq_len: Optional[int] = None,
217
+ rbln_batch_size: Optional[int] = None,
218
+ rbln_quantization: Optional[Dict[str, str]] = None,
219
+ **kwargs,
220
+ ) -> RBLNConfig:
221
+ meta = {}
222
+
223
+ prefill_chunk_size = 128
224
+ if rbln_max_seq_len is None:
225
+ rbln_max_seq_len = getattr(model_config, "max_position_embeddings", None) or getattr(
226
+ model_config, "n_positions", None
227
+ )
228
+ if rbln_max_seq_len is None:
229
+ raise ValueError("`rbln_max_seq_len` should be specified.")
230
+ rbln_batch_size = 1 if rbln_batch_size is None else rbln_batch_size
231
+
232
+ meta["rbln_max_seq_len"] = rbln_max_seq_len
233
+ meta["rbln_batch_size"] = rbln_batch_size
234
+ meta["rbln_prefill_chunk_size"] = prefill_chunk_size
235
+
236
+ num_attention_heads = getattr(model_config, "n_head", None) or getattr(model_config, "num_attention_heads")
237
+ num_key_value_heads = getattr(model_config, "num_key_value_heads", None) or num_attention_heads
238
+ num_hidden_layers = getattr(model_config, "n_layer", None) or getattr(model_config, "num_hidden_layers")
239
+ head_dim = getattr(model_config, "head_dim", None) or model_config.hidden_size // num_attention_heads
240
+
241
+ if rbln_quantization is not None:
242
+ q_format = rbln_quantization.get("format", None)
243
+ q_precision = rbln_quantization.get("precision", None)
244
+
245
+ if q_format not in SUPPORTED_QUANTIZATIONS.keys() or q_precision not in SUPPORTED_QUANTIZATIONS[q_format]:
246
+ raise ValueError(
247
+ f'rbln_quantization="{rbln_quantization}" is not a supported quantization format or precesion, '
248
+ f"Possible: {SUPPORTED_QUANTIZATIONS}"
249
+ )
250
+ meta["rbln_quantization"] = rbln_quantization
251
+
252
+ def get_input_info(
253
+ batch_size,
254
+ query_length,
255
+ ):
256
+ input_info = [
257
+ ("input_ids", [batch_size, query_length], "int64"),
258
+ ("attention_mask", [batch_size, 1, query_length, rbln_max_seq_len], "int64"),
259
+ (
260
+ "cache_position",
261
+ [batch_size, query_length],
262
+ "int32",
263
+ ),
264
+ ("batch_position", [], "int16"),
265
+ ]
266
+
267
+ input_info.extend(
268
+ [
269
+ (
270
+ f"past_key_values_{i}",
271
+ [
272
+ rbln_batch_size,
273
+ num_key_value_heads,
274
+ rbln_max_seq_len,
275
+ head_dim,
276
+ ],
277
+ "float32",
278
+ )
279
+ for i in range(num_hidden_layers * 2)
280
+ ]
281
+ )
282
+
283
+ return input_info
284
+
285
+ prefill_input_info = get_input_info(
286
+ batch_size=1,
287
+ query_length=prefill_chunk_size,
288
+ )
289
+ dec_input_info = get_input_info(
290
+ batch_size=rbln_batch_size,
291
+ query_length=1,
292
+ )
293
+
294
+ prefill_rbln_runtime_config = RBLNRuntimeConfig(input_info=prefill_input_info)
295
+ dec_rbln_runtime_config = RBLNRuntimeConfig(input_info=dec_input_info)
296
+
297
+ dec_rbln_runtime_config.batch_size = rbln_batch_size
298
+
299
+ rbln_config = RBLNConfig.from_rbln_runtime_configs(
300
+ [prefill_rbln_runtime_config, dec_rbln_runtime_config],
301
+ _rbln_meta=meta,
302
+ )
303
+
304
+ return rbln_config
305
+
306
+ @classmethod
307
+ def _create_runtimes(
308
+ cls, compiled_models: List[rebel.RBLNCompiledModel], rbln_device_map: Dict[str, int]
309
+ ) -> List[rebel.Runtime]:
310
+ device_val = rbln_device_map[DEFAULT_COMPILED_MODEL_NAME]
311
+ return [
312
+ compiled_models[0].create_runtime(input_info_index=0, tensor_type="pt", device=device_val),
313
+ compiled_models[0].create_runtime(input_info_index=1, tensor_type="pt", device=device_val),
314
+ ]
315
+
316
+ def get_decoder(self):
317
+ return self.decoder
318
+
319
+ def can_generate(self):
320
+ return True
321
+
322
+ def _reorder_cache(self, past_key_values, beam_idx):
323
+ raise NotImplementedError
324
+
325
+ # args input_ids, past_key_values and attention_mask are updated by _update_model_kwargs_for_generation() in _greedy_search() in GenerationMixin
326
+ def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs):
327
+ batch_size = input_ids.shape[0]
328
+
329
+ # FIXME past_key_values is just carriier variable for past_cached_length
330
+ # torch.tensor((4,1),dtype=torch.int32) which refers a past_cached_length of each batch
331
+ past_cached_length = past_key_values
332
+ if past_cached_length is None:
333
+ l_input_ids = []
334
+ cache_positions = []
335
+ past_cached_length = torch.zeros((batch_size, 1), dtype=torch.int32)
336
+ for i in range(batch_size):
337
+ input_id = input_ids[i]
338
+ input_id = input_id[attention_mask[i] == 1]
339
+ valid_len = input_id.shape[-1]
340
+ cache_position = torch.arange(0, valid_len, dtype=torch.int32)
341
+ past_cached_length[i] = valid_len
342
+ l_input_ids.append(input_id.unsqueeze(0))
343
+ cache_positions.append(cache_position.unsqueeze(0))
344
+
345
+ input_ids = l_input_ids
346
+ else:
347
+ input_ids = input_ids[:, -1:]
348
+ cache_positions = past_cached_length
349
+ past_cached_length = past_cached_length + 1
350
+
351
+ model_inputs = {
352
+ "input_ids": input_ids,
353
+ "cache_position": cache_positions,
354
+ "past_cached_length": past_cached_length,
355
+ }
356
+
357
+ return model_inputs
358
+
359
+ def forward(
360
+ self,
361
+ input_ids: torch.LongTensor = None,
362
+ cache_position: Union[List[torch.Tensor], torch.Tensor] = None, # vllm keyword argument
363
+ batch_idx: Optional[int] = None,
364
+ past_cached_length: Optional[torch.Tensor] = None, # past_cached_length
365
+ **kwargs,
366
+ ) -> Tuple[torch.FloatTensor]:
367
+ # prefll & hf generate
368
+ if isinstance(cache_position, list):
369
+ logits = []
370
+ for batch_idx, (input_id, cache_pos) in enumerate(zip(input_ids, cache_position)):
371
+ logit = self._forward_prefill(input_ids=input_id, cache_position=cache_pos, batch_idx=batch_idx)
372
+ logits.append(logit)
373
+ logits = torch.cat(logits, dim=0)
374
+ # prefill & vllm step
375
+ elif cache_position.shape[-1] > 1:
376
+ logits = self._forward_prefill(input_ids=input_ids, cache_position=cache_position, batch_idx=batch_idx)
377
+ # common decoder
378
+ else:
379
+ logits = self._forward_decoder(input_ids=input_ids, cache_position=cache_position)
380
+
381
+ return CausalLMOutputWithPast(
382
+ logits=logits,
383
+ past_key_values=past_cached_length, # past_cached_length
384
+ )
385
+
386
+ def _forward_prefill(
387
+ self,
388
+ input_ids: torch.LongTensor = None,
389
+ cache_position: torch.Tensor = None, # torch.tensor(,dtype=int32) (1,64) // (4,1)
390
+ batch_idx: int = None,
391
+ ) -> torch.FloatTensor:
392
+ if batch_idx is None or batch_idx >= self.batch_size:
393
+ raise RuntimeError(
394
+ f"Invalid batch_idx ({batch_idx}). It must be a non-null value less than the batch size ({self.batch_size})."
395
+ )
396
+
397
+ out_buffers = [
398
+ torch.empty(
399
+ size=[
400
+ 1,
401
+ self.prefill_chunk_size,
402
+ self.config.vocab_size,
403
+ ],
404
+ dtype=torch.float32,
405
+ device="cpu",
406
+ ),
407
+ torch.empty(size=[], dtype=torch.int16, device="cpu"),
408
+ ]
409
+
410
+ query_length = input_ids.shape[1]
411
+ attention_mask = self.prefill_attention_mask.clone()
412
+ for step in range(0, query_length, self.prefill_chunk_size):
413
+ if step + self.prefill_chunk_size > query_length:
414
+ input_ids = torch.nn.functional.pad(input_ids, (0, step + self.prefill_chunk_size - query_length))
415
+ cache_position = torch.cat(
416
+ [
417
+ cache_position,
418
+ torch.arange(
419
+ query_length,
420
+ step + self.prefill_chunk_size,
421
+ dtype=torch.int32,
422
+ ).unsqueeze(0),
423
+ ],
424
+ dim=-1,
425
+ )
426
+
427
+ sliced_input_ids = input_ids[:, step : step + self.prefill_chunk_size]
428
+ sliced_cache_positions = cache_position[:, step : step + self.prefill_chunk_size]
429
+
430
+ if step >= self.prefill_chunk_size:
431
+ attention_mask[:, :, :, step - self.prefill_chunk_size : step] = 1
432
+ attention_mask[:, :, :, step : step + self.prefill_chunk_size] = self.causal_mask
433
+
434
+ logits, _ = self.prefill_decoder(
435
+ sliced_input_ids.contiguous(),
436
+ attention_mask.contiguous(),
437
+ sliced_cache_positions.contiguous(),
438
+ torch.tensor(batch_idx, dtype=torch.int16),
439
+ out=out_buffers,
440
+ )
441
+ logits = logits[:, query_length % self.prefill_chunk_size - 1].unsqueeze(1)
442
+
443
+ self.dec_attn_mask[batch_idx] = self.dec_attn_mask_init.clone()
444
+ self.dec_attn_mask[batch_idx, :, :, :query_length] = 1
445
+
446
+ return logits
447
+
448
+ def _forward_decoder(
449
+ self, input_ids: torch.LongTensor = None, cache_position: torch.Tensor = None
450
+ ) -> torch.FloatTensor:
451
+ batch_size = input_ids.shape[0]
452
+
453
+ for b_idx in range(batch_size):
454
+ decoding_step = cache_position[b_idx].item()
455
+ self.dec_attn_mask[b_idx, :, :, decoding_step] = 1
456
+
457
+ logits, _ = self.decoder(
458
+ input_ids.contiguous(),
459
+ self.dec_attn_mask.contiguous(),
460
+ cache_position.contiguous(),
461
+ torch.tensor(0, dtype=torch.int16),
462
+ )
463
+
464
+ 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,123 @@
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(
43
+ {
44
+ "wrapper": GemmaModel.forward,
45
+ "model": DecoderOnlyDecoderLayer.forward,
46
+ "decoder_layer": DecoderOnlyAttention.forward,
47
+ }
48
+ )
49
+ return forward_dict
50
+
51
+
52
+ class GemmaModel:
53
+ def forward(
54
+ self,
55
+ input_ids: torch.LongTensor = None,
56
+ attention_mask: Optional[torch.Tensor] = None,
57
+ position_ids: Optional[torch.LongTensor] = None,
58
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
59
+ batch_ids: Optional[torch.LongTensor] = None,
60
+ inputs_embeds: Optional[torch.FloatTensor] = None,
61
+ use_cache: Optional[bool] = True,
62
+ output_attentions: Optional[bool] = False,
63
+ output_hidden_states: Optional[bool] = False,
64
+ forward_dict: Optional[Dict[str, classmethod]] = None,
65
+ rotary_pos_emb=None,
66
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
67
+ # embed positions
68
+ inputs_embeds = self.embed_tokens(input_ids)
69
+ hidden_states = inputs_embeds
70
+
71
+ ##### GEMMA change from llama#####
72
+ hidden_states = hidden_states * (self.config.hidden_size**0.5)
73
+
74
+ attention_mask = (1 - attention_mask) * torch.finfo(torch.float16).min
75
+
76
+ # get cos,sin vector
77
+ cos, sin = rotary_pos_emb(inputs_embeds, attention_mask.shape[-1])
78
+ cos, sin = slice_and_unsqueeze_cos_sin(cos, sin, position_ids)
79
+
80
+ # decoder layers
81
+ all_hidden_states = () if output_hidden_states else None
82
+ all_self_attns = () if output_attentions else None
83
+
84
+ for layer_idx, decoder_layer in enumerate(self.layers):
85
+ if output_hidden_states:
86
+ all_hidden_states += (hidden_states,)
87
+ layer_outputs = forward_dict["model"](
88
+ decoder_layer,
89
+ hidden_states,
90
+ layer_idx,
91
+ attention_mask=attention_mask,
92
+ position_ids=position_ids,
93
+ past_key_value=past_key_values,
94
+ output_attentions=output_attentions,
95
+ use_cache=use_cache,
96
+ batch_ids=batch_ids,
97
+ cos=cos,
98
+ sin=sin,
99
+ forward_dict=forward_dict,
100
+ )
101
+
102
+ hidden_states = layer_outputs[0]
103
+
104
+ updated_cache = layer_outputs[2 if output_attentions else 1]
105
+
106
+ if output_attentions:
107
+ all_self_attns += (layer_outputs[1],)
108
+
109
+ hidden_states = self.norm(hidden_states)
110
+
111
+ # add hidden states from the last decoder layer
112
+ if output_hidden_states:
113
+ all_hidden_states += (hidden_states,)
114
+
115
+ # convert RebelDynamicCache to legacy Tuple[Tuple[torch.Tensor]]
116
+ next_cache = updated_cache.to_legacy_cache()
117
+
118
+ return BaseModelOutputWithPast(
119
+ last_hidden_state=hidden_states,
120
+ past_key_values=next_cache,
121
+ hidden_states=all_hidden_states,
122
+ attentions=all_self_attns,
123
+ )
@@ -0,0 +1,67 @@
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 TYPE_CHECKING, Any, Callable
27
+
28
+ from transformers import GemmaForCausalLM
29
+
30
+ from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
31
+ from .gemma_architecture import GemmaWrapper
32
+
33
+
34
+ if TYPE_CHECKING:
35
+ from transformers import PreTrainedModel
36
+
37
+ from ....modeling_config import RBLNConfig
38
+
39
+ logger = logging.getLogger(__name__)
40
+
41
+
42
+ class RBLNGemmaForCausalLM(RBLNDecoderOnlyModelForCausalLM):
43
+ """
44
+ The Gemma Model transformer with a language modeling head (linear layer) on top.
45
+ This model inherits from [`RBLNMultiModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
46
+
47
+ A class to convert and run pre-trained transformers based GemmaForCausalLM model on RBLN devices.
48
+ It implements the methods to convert a pre-trained transformers GemmaForCausalLM model into a RBLN transformer model by:
49
+ - transferring the checkpoint weights of the original into an optimized RBLN graph,
50
+ - compiling the resulting graph using the RBLN compiler.
51
+ """
52
+
53
+ @classmethod
54
+ def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
55
+ rbln_max_seq_len = rbln_config.meta["rbln_max_seq_len"]
56
+ return GemmaWrapper(model, rbln_max_seq_len).eval()
57
+
58
+ def __getattr__(self, __name: str) -> Any:
59
+ def redirect(func):
60
+ return lambda *pargs, **kwargs: func(self, *pargs, **kwargs)
61
+
62
+ val = getattr(GemmaForCausalLM, __name)
63
+
64
+ if isinstance(val, Callable) and "self" in set(inspect.signature(val).parameters):
65
+ return redirect(val)
66
+
67
+ return val