optimum-rbln 0.1.0__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 (51) hide show
  1. optimum/rbln/__init__.py +115 -0
  2. optimum/rbln/__version__.py +1 -0
  3. optimum/rbln/diffusers/__init__.py +64 -0
  4. optimum/rbln/diffusers/models/__init__.py +26 -0
  5. optimum/rbln/diffusers/models/autoencoder_kl.py +313 -0
  6. optimum/rbln/diffusers/models/controlnet.py +180 -0
  7. optimum/rbln/diffusers/models/unet_2d_condition.py +352 -0
  8. optimum/rbln/diffusers/pipelines/__init__.py +30 -0
  9. optimum/rbln/diffusers/pipelines/controlnet/__init__.py +24 -0
  10. optimum/rbln/diffusers/pipelines/controlnet/multicontrolnet.py +266 -0
  11. optimum/rbln/diffusers/pipelines/stable_diffusion/__init__.py +26 -0
  12. optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_controlnet_img2img.py +731 -0
  13. optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +106 -0
  14. optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +116 -0
  15. optimum/rbln/diffusers/pipelines/stable_diffusion_xl/__init__.py +2 -0
  16. optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py +109 -0
  17. optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py +111 -0
  18. optimum/rbln/modeling.py +0 -0
  19. optimum/rbln/modeling_alias.py +49 -0
  20. optimum/rbln/modeling_base.py +645 -0
  21. optimum/rbln/modeling_config.py +169 -0
  22. optimum/rbln/modeling_seq2seq.py +469 -0
  23. optimum/rbln/transformers/__init__.py +59 -0
  24. optimum/rbln/transformers/generation/__init__.py +24 -0
  25. optimum/rbln/transformers/generation/streamers.py +122 -0
  26. optimum/rbln/transformers/models/__init__.py +28 -0
  27. optimum/rbln/transformers/models/bart/__init__.py +24 -0
  28. optimum/rbln/transformers/models/bart/bart_architecture.py +377 -0
  29. optimum/rbln/transformers/models/clip/__init__.py +24 -0
  30. optimum/rbln/transformers/models/clip/modeling_clip.py +116 -0
  31. optimum/rbln/transformers/models/gpt2/__init__.py +24 -0
  32. optimum/rbln/transformers/models/gpt2/gpt2_architecture.py +253 -0
  33. optimum/rbln/transformers/models/gpt2/modeling_gpt2.py +700 -0
  34. optimum/rbln/transformers/models/llama/__init__.py +24 -0
  35. optimum/rbln/transformers/models/llama/llama_architecture.py +607 -0
  36. optimum/rbln/transformers/models/llama/modeling_llama.py +409 -0
  37. optimum/rbln/transformers/models/t5/__init__.py +24 -0
  38. optimum/rbln/transformers/models/t5/t5_architecture.py +439 -0
  39. optimum/rbln/transformers/models/wav2vec2/__init__.py +24 -0
  40. optimum/rbln/transformers/models/wav2vec2/modeling_wav2vec2.py +121 -0
  41. optimum/rbln/transformers/models/whisper/__init__.py +24 -0
  42. optimum/rbln/transformers/models/whisper/modeling_whisper.py +374 -0
  43. optimum/rbln/transformers/models/whisper/whisper_architecture.py +406 -0
  44. optimum/rbln/utils/__init__.py +25 -0
  45. optimum/rbln/utils/import_utils.py +28 -0
  46. optimum/rbln/utils/runtime_utils.py +71 -0
  47. optimum/rbln/utils/save_utils.py +92 -0
  48. optimum_rbln-0.1.0.dist-info/METADATA +144 -0
  49. optimum_rbln-0.1.0.dist-info/RECORD +51 -0
  50. optimum_rbln-0.1.0.dist-info/WHEEL +4 -0
  51. optimum_rbln-0.1.0.dist-info/licenses/LICENSE +201 -0
@@ -0,0 +1,409 @@
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 # noqa: I001
25
+ import logging
26
+ from pathlib import Path
27
+ from tempfile import TemporaryDirectory
28
+ from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
29
+
30
+ import torch # noqa: F401
31
+ import rebel # noqa: F401
32
+
33
+ from optimum.exporters import TasksManager
34
+ from transformers import AutoModelForCausalLM, LlamaForCausalLM, PretrainedConfig, AutoConfig
35
+ from transformers.modeling_outputs import CausalLMOutputWithPast
36
+
37
+ from ....modeling_base import RBLNBaseModel
38
+ from ....modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNConfig, RBLNRuntimeConfig
39
+ from ....utils.runtime_utils import RBLNPytorchRuntime
40
+ from ....utils.save_utils import maybe_save_preprocessors
41
+ from .llama_architecture import (
42
+ LlamaWrapper,
43
+ wrap_llama,
44
+ unwrap_llama,
45
+ )
46
+
47
+
48
+ logger = logging.getLogger(__name__)
49
+
50
+ if TYPE_CHECKING:
51
+ from transformers import (
52
+ AutoFeatureExtractor,
53
+ AutoProcessor,
54
+ AutoTokenizer,
55
+ PretrainedConfig,
56
+ )
57
+
58
+
59
+ class RBLNRuntimeModel(RBLNPytorchRuntime):
60
+ mandatory_members = ["main_input_name"]
61
+
62
+ # RBLN_Runtimemodule
63
+ def forward(
64
+ self,
65
+ input_ids: torch.LongTensor = None,
66
+ attention_mask: torch.LongTensor = None,
67
+ cache_position: torch.Tensor = None,
68
+ **kwargs: Dict[str, Any],
69
+ ):
70
+ logits = super().forward(
71
+ input_ids=input_ids,
72
+ attention_mask=attention_mask,
73
+ cache_position=cache_position,
74
+ )
75
+ return logits
76
+
77
+
78
+ class RBLNLlamaForCausalLM(RBLNBaseModel):
79
+ """
80
+ The Llama Model transformer with a language modeling head (linear layer) on top.
81
+ This model inherits from [`RBLNBaseModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
82
+
83
+ A class to convert and run pre-trained transformers based LlamaForCausalLM model on RBLN devices.
84
+ It implements the methods to convert a pre-trained transformers LlamaForCausalLM model into a RBLN transformer model by:
85
+ - transferring the checkpoint weights of the original into an optimized RBLN graph,
86
+ - compiling the resulting graph using the RBLN compiler.
87
+ """
88
+
89
+ model_type = "rbln_model"
90
+ main_input_name = "input_ids"
91
+ auto_model_class = AutoModelForCausalLM
92
+
93
+ def __post_init__(self, **kwargs):
94
+
95
+ self.batch_size = self.rbln_config.meta["rbln_batch_size"]
96
+ self.max_seq_len = self.rbln_config.meta["rbln_max_seq_len"]
97
+ self.prefill_chunk_size = self.rbln_config.meta["rbln_prefill_chunk_size"]
98
+
99
+ self.prefill_attention_mask = torch.zeros(
100
+ self.batch_size, 1, self.prefill_chunk_size, self.max_seq_len, dtype=torch.int64
101
+ )
102
+ self.causal_mask = 1 - torch.triu(
103
+ torch.ones(self.batch_size, 1, self.prefill_chunk_size, self.prefill_chunk_size), diagonal=1
104
+ )
105
+
106
+ self.prefill_decoder = RBLNRuntimeModel(runtime=self.runtimes[0], main_input_name="input_ids")
107
+ self.decoder = RBLNRuntimeModel(runtime=self.runtimes[1], main_input_name="input_ids")
108
+ self.past_cached_length = 0
109
+
110
+ @classmethod
111
+ @torch.no_grad()
112
+ def _export(
113
+ cls,
114
+ model_id: str,
115
+ config: "PretrainedConfig",
116
+ use_auth_token: Optional[Union[bool, str]] = None,
117
+ revision: Optional[str] = None,
118
+ force_download: bool = False,
119
+ cache_dir: Optional[str] = None,
120
+ subfolder: str = "",
121
+ local_files_only: bool = False,
122
+ trust_remote_code: bool = False,
123
+ **kwargs,
124
+ ) -> "RBLNLlamaForCausalLM":
125
+ task = kwargs.pop("task", None)
126
+ if task is None:
127
+ task = TasksManager.infer_task_from_model(cls.auto_model_class)
128
+
129
+ save_dir = TemporaryDirectory()
130
+ save_dir_path = Path(save_dir.name)
131
+
132
+ def update_configs(kwargs):
133
+ hf_max_position_embeddings = getattr(AutoConfig.from_pretrained(model_id), "max_position_embeddings", None)
134
+ max_seq_len = kwargs.get("rbln_max_seq_len", None)
135
+ if max_seq_len is not None:
136
+ if max_seq_len <= hf_max_position_embeddings:
137
+ kwargs.update({"max_position_embeddings": max_seq_len})
138
+ else:
139
+ raise ValueError("`max_seq_len` should be less or equal than max_position_embeddings!")
140
+
141
+ kwargs.update(
142
+ {
143
+ "torchscript": True,
144
+ "return_dict": False,
145
+ "use_cache": True,
146
+ "torch_dtype": torch.float32,
147
+ "_attn_implementation": "eager",
148
+ }
149
+ )
150
+
151
+ return kwargs
152
+
153
+ kwargs = update_configs(kwargs)
154
+
155
+ rbln_config_kwargs, rbln_constructor_kwargs = cls.pop_rbln_kwargs_from_kwargs(kwargs)
156
+
157
+ origin_mehtods = wrap_llama()
158
+ model: LlamaForCausalLM = TasksManager.get_model_from_task(
159
+ task=task,
160
+ model_name_or_path=model_id,
161
+ subfolder=subfolder,
162
+ revision=revision,
163
+ framework="pt",
164
+ cache_dir=cache_dir,
165
+ use_auth_token=use_auth_token,
166
+ local_files_only=local_files_only,
167
+ force_download=force_download,
168
+ trust_remote_code=trust_remote_code,
169
+ **kwargs,
170
+ )
171
+
172
+ if config is None:
173
+ config = model.config
174
+
175
+ config.save_pretrained(save_dir_path)
176
+ preprocessors = maybe_save_preprocessors(model_id, save_dir_path, src_subfolder=subfolder)
177
+
178
+ # Get compilation arguments
179
+ if rbln_config_kwargs.get("rbln_config", None) is None:
180
+ rbln_config = cls.get_rbln_config(
181
+ preprocessors=preprocessors, model_config=model.config, **rbln_config_kwargs
182
+ )
183
+
184
+ def compile_llama():
185
+ wrapped_model = LlamaWrapper(model).eval()
186
+
187
+ prefill_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][0]
188
+ dec_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][1]
189
+
190
+ prefill_example_inputs = prefill_rbln_runtime_config.get_dummy_inputs(fill=0)
191
+ dec_example_inputs = dec_rbln_runtime_config.get_dummy_inputs(fill=0)
192
+
193
+ prefill_scripted_model = torch.jit.trace(wrapped_model, prefill_example_inputs)
194
+ dec_scripted_model = torch.jit.trace(wrapped_model, dec_example_inputs)
195
+
196
+ prefill_ir = rebel.torchscript_to_ir(
197
+ prefill_scripted_model,
198
+ input_names=[v[0] for v in prefill_rbln_runtime_config.input_info],
199
+ )
200
+ dec_ir = rebel.torchscript_to_ir(
201
+ dec_scripted_model,
202
+ input_names=[v[0] for v in dec_rbln_runtime_config.input_info],
203
+ )
204
+
205
+ # Caching prefill_decoder/decoder I/O
206
+ connections = [
207
+ (prefill_ir.outputs[1 + i], prefill_ir.inputs[3 + i])
208
+ for i in range(model.config.num_hidden_layers * 2)
209
+ ]
210
+
211
+ compiled_model = rebel.compile(
212
+ prefill_ir,
213
+ dec_ir,
214
+ connections=connections,
215
+ fusion=prefill_rbln_runtime_config.fusion,
216
+ npu=prefill_rbln_runtime_config.npu,
217
+ tensor_parallel_size=prefill_rbln_runtime_config.tensor_parallel_size,
218
+ use_weight_sharing=True,
219
+ )
220
+ compiled_model.save(save_dir_path / f"{DEFAULT_COMPILED_MODEL_NAME}.rbln")
221
+
222
+ compile_llama()
223
+ unwrap_llama(origin_mehtods)
224
+
225
+ rbln_config.save(save_dir_path)
226
+
227
+ return cls._from_pretrained(
228
+ model_id=save_dir_path,
229
+ config=config,
230
+ model_save_dir=save_dir,
231
+ **rbln_constructor_kwargs,
232
+ **kwargs,
233
+ )
234
+
235
+ @classmethod
236
+ def _get_rbln_config(
237
+ cls,
238
+ preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
239
+ model_config: "PretrainedConfig",
240
+ rbln_max_seq_len: Optional[int] = None,
241
+ rbln_batch_size: Optional[int] = None,
242
+ ) -> RBLNConfig:
243
+ meta = {}
244
+
245
+ prefill_chunk_size = 128
246
+ if rbln_max_seq_len is None:
247
+ rbln_max_seq_len = getattr(model_config, "max_position_embeddings", None)
248
+
249
+ meta["rbln_max_seq_len"] = rbln_max_seq_len
250
+ meta["rbln_batch_size"] = rbln_batch_size
251
+ meta["rbln_prefill_chunk_size"] = prefill_chunk_size
252
+
253
+ def get_input_info(query_length):
254
+ input_info = [
255
+ ("input_ids", [rbln_batch_size, query_length], "int64"),
256
+ ("attention_mask", [rbln_batch_size, 1, query_length, rbln_max_seq_len], "int64"),
257
+ (
258
+ "cache_position",
259
+ [],
260
+ "int32",
261
+ ),
262
+ ]
263
+ input_info.extend(
264
+ [
265
+ (
266
+ f"past_key_values_{i}",
267
+ [
268
+ rbln_batch_size,
269
+ model_config.num_key_value_heads,
270
+ rbln_max_seq_len,
271
+ model_config.hidden_size // model_config.num_attention_heads,
272
+ ],
273
+ "float32",
274
+ )
275
+ for i in range(model_config.num_hidden_layers * 2)
276
+ ]
277
+ )
278
+ return input_info
279
+
280
+ prefill_input_info = get_input_info(query_length=prefill_chunk_size)
281
+ dec_input_info = get_input_info(query_length=1)
282
+
283
+ prefill_rbln_runtime_config = RBLNRuntimeConfig(input_info=prefill_input_info)
284
+ dec_rbln_runtime_config = RBLNRuntimeConfig(input_info=dec_input_info)
285
+
286
+ dec_rbln_runtime_config.batch_size = rbln_batch_size
287
+
288
+ rbln_config = RBLNConfig.from_rbln_runtime_configs(
289
+ [prefill_rbln_runtime_config, dec_rbln_runtime_config],
290
+ _rbln_meta=meta,
291
+ )
292
+
293
+ return rbln_config
294
+
295
+ def _create_runtimes(self, rbln_device_map: Dict[str, int]) -> List[rebel.Runtime]:
296
+ device_val = rbln_device_map[DEFAULT_COMPILED_MODEL_NAME]
297
+ return [
298
+ self.compiled_models[0].create_runtime(input_info_index=0, tensor_type="pt", device=device_val),
299
+ self.compiled_models[0].create_runtime(input_info_index=1, tensor_type="pt", device=device_val),
300
+ ]
301
+
302
+ def get_decoder(self):
303
+ return self.decoder
304
+
305
+ def can_generate(self):
306
+ return True
307
+
308
+ def __getattr__(self, __name: str) -> Any:
309
+ def redirect(func):
310
+ return lambda *pargs, **kwargs: func(self, *pargs, **kwargs)
311
+
312
+ val = getattr(LlamaForCausalLM, __name)
313
+
314
+ if isinstance(val, Callable) and "self" in set(inspect.signature(val).parameters):
315
+ return redirect(val)
316
+
317
+ return val
318
+
319
+ def _reorder_cache(self, past_key_values, beam_idx):
320
+ raise NotImplementedError
321
+
322
+ # args input_ids, past_key_values and attention_mask are updated by _update_model_kwargs_for_generation() in _greedy_search() in GenerationMixin
323
+ def prepare_inputs_for_generation(self, input_ids, past_key_values=0, attention_mask=None, **kwargs):
324
+ batch_size, hf_input_length = input_ids.shape
325
+ past_cached_length = past_key_values
326
+ query_length = hf_input_length - past_cached_length
327
+
328
+ # In greedy decoding
329
+ if past_key_values == 0:
330
+ self.prompt_length = query_length
331
+ self.prompt_ids = input_ids
332
+ self.prompt_attn_mask = attention_mask.unsqueeze(1).unsqueeze(1).contiguous()
333
+
334
+ attention_mask = torch.zeros(batch_size, 1, self.prefill_chunk_size, self.max_seq_len, dtype=torch.int64)
335
+ cache_position = torch.tensor(0, dtype=torch.int32)
336
+ else:
337
+ attention_mask = torch.nn.functional.pad(attention_mask, (0, self.max_seq_len - hf_input_length))
338
+ attention_mask = attention_mask.reshape(batch_size, 1, 1, -1).contiguous()
339
+ cache_position = torch.tensor(past_cached_length, dtype=torch.int32)
340
+ input_ids = input_ids[:, -1:]
341
+
342
+ model_inputs = {
343
+ "input_ids": input_ids,
344
+ "past_key_values": past_key_values,
345
+ "attention_mask": attention_mask,
346
+ "cache_position": cache_position,
347
+ "query_length": query_length,
348
+ }
349
+
350
+ return model_inputs
351
+
352
+ def forward(
353
+ self,
354
+ input_ids: torch.LongTensor = None,
355
+ attention_mask: Optional[torch.Tensor] = None,
356
+ past_key_values: int = None,
357
+ cache_position: Optional[torch.Tensor] = None,
358
+ query_length: Optional[torch.Tensor] = None,
359
+ **kwargs,
360
+ ) -> Tuple[torch.FloatTensor]:
361
+ if past_key_values is not None:
362
+ past_key_values += query_length
363
+
364
+ # prefill_decoder
365
+ if cache_position == 0:
366
+ while query_length > self.prefill_chunk_size:
367
+ # prepare input_ids & attention_mask
368
+ sliced_input_ids = input_ids[:, cache_position : cache_position + self.prefill_chunk_size].contiguous()
369
+ attention_mask[:, :, :, :cache_position] = 1
370
+ attention_mask[:, :, :, cache_position : cache_position + self.prefill_chunk_size] = self.causal_mask
371
+ attention_mask[:, :, :, : self.prompt_length] *= self.prompt_attn_mask[:, :, :, :]
372
+
373
+ _ = self.prefill_decoder(
374
+ sliced_input_ids,
375
+ attention_mask,
376
+ cache_position,
377
+ )
378
+ # update query_length & cache_position
379
+ query_length -= self.prefill_chunk_size
380
+ cache_position += self.prefill_chunk_size
381
+
382
+ # prepare input_ids & attention_mask
383
+ last_input_ids = input_ids[:, cache_position : cache_position + query_length]
384
+ last_input_ids = torch.nn.functional.pad(last_input_ids, (0, self.prefill_chunk_size - query_length))
385
+
386
+ attention_mask[:, :, :, :cache_position] = 1
387
+ mask_slice = self.causal_mask[:, :, :query_length, :query_length]
388
+ attention_mask[:, :, :query_length, cache_position : cache_position + query_length] = mask_slice
389
+ attention_mask[:, :, :, : self.prompt_length] *= self.prompt_attn_mask[:, :, :, :]
390
+
391
+ outputs = self.prefill_decoder(
392
+ last_input_ids.contiguous(),
393
+ attention_mask.contiguous(),
394
+ cache_position,
395
+ )
396
+
397
+ outputs = outputs[:, query_length - 1].unsqueeze(1)
398
+ # decoder
399
+ else:
400
+ outputs = self.decoder(
401
+ input_ids.contiguous(),
402
+ attention_mask.contiguous(),
403
+ cache_position=cache_position,
404
+ )
405
+
406
+ return CausalLMOutputWithPast(
407
+ logits=outputs,
408
+ past_key_values=past_key_values,
409
+ )
@@ -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 .t5_architecture import T5DecoderWrapper, T5EncoderWrapper