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.
- optimum/rbln/__init__.py +14 -0
- optimum/rbln/__version__.py +1 -1
- optimum/rbln/diffusers/__init__.py +0 -1
- optimum/rbln/diffusers/models/controlnet.py +3 -0
- optimum/rbln/diffusers/models/unet_2d_condition.py +2 -2
- optimum/rbln/diffusers/pipelines/controlnet/multicontrolnet.py +22 -144
- optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet.py +107 -59
- optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py +106 -54
- optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py +130 -71
- optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py +131 -72
- optimum/rbln/modeling_alias.py +14 -0
- optimum/rbln/modeling_base.py +110 -0
- optimum/rbln/transformers/__init__.py +6 -0
- optimum/rbln/transformers/cache_utils.py +111 -0
- optimum/rbln/transformers/generation/utils.py +0 -2
- optimum/rbln/transformers/models/__init__.py +2 -0
- optimum/rbln/transformers/models/bart/bart_architecture.py +0 -5
- optimum/rbln/transformers/models/decoderonly/__init__.py +36 -0
- optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py +515 -0
- optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py +349 -0
- optimum/rbln/transformers/models/gemma/__init__.py +24 -0
- optimum/rbln/transformers/models/gemma/gemma_architecture.py +116 -0
- optimum/rbln/transformers/models/gemma/modeling_gemma.py +61 -0
- optimum/rbln/transformers/models/gpt2/gpt2_architecture.py +201 -166
- optimum/rbln/transformers/models/gpt2/modeling_gpt2.py +56 -220
- optimum/rbln/transformers/models/llama/llama_architecture.py +3 -610
- optimum/rbln/transformers/models/llama/modeling_llama.py +8 -442
- optimum/rbln/transformers/models/midm/hf_hub_cached/midm_bitext_tokenization.py +2 -1
- optimum/rbln/transformers/models/midm/hf_hub_cached/modeling_midm.py +0 -4
- optimum/rbln/transformers/models/midm/midm_architecture.py +160 -357
- optimum/rbln/transformers/models/midm/modeling_midm.py +40 -272
- optimum/rbln/transformers/models/whisper/whisper_architecture.py +0 -6
- optimum/rbln/transformers/models/xlm_roberta/__init__.py +24 -0
- optimum/rbln/transformers/models/xlm_roberta/modeling_xlm_roberta.py +125 -0
- {optimum_rbln-0.1.7.dist-info → optimum_rbln-0.1.8.dist-info}/METADATA +2 -3
- {optimum_rbln-0.1.7.dist-info → optimum_rbln-0.1.8.dist-info}/RECORD +38 -30
- optimum/rbln/transformers/models/llama/llama_architecture_cb.py +0 -764
- {optimum_rbln-0.1.7.dist-info → optimum_rbln-0.1.8.dist-info}/WHEEL +0 -0
- {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
|