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.
- optimum/rbln/__init__.py +17 -0
- optimum/rbln/__version__.py +1 -1
- optimum/rbln/diffusers/__init__.py +0 -1
- optimum/rbln/diffusers/models/autoencoder_kl.py +3 -3
- optimum/rbln/diffusers/models/controlnet.py +7 -3
- optimum/rbln/diffusers/models/unet_2d_condition.py +5 -5
- optimum/rbln/diffusers/pipelines/controlnet/multicontrolnet.py +23 -146
- 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 +19 -1
- optimum/rbln/modeling_base.py +162 -18
- optimum/rbln/transformers/__init__.py +8 -0
- optimum/rbln/transformers/cache_utils.py +111 -0
- optimum/rbln/transformers/generation/utils.py +0 -2
- optimum/rbln/transformers/models/__init__.py +3 -0
- optimum/rbln/transformers/models/bart/bart_architecture.py +0 -5
- optimum/rbln/transformers/models/clip/modeling_clip.py +1 -1
- optimum/rbln/transformers/models/decoderonly/__init__.py +36 -0
- optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py +516 -0
- optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py +464 -0
- optimum/rbln/transformers/models/gemma/__init__.py +24 -0
- optimum/rbln/transformers/models/gemma/gemma_architecture.py +123 -0
- optimum/rbln/transformers/models/gemma/modeling_gemma.py +67 -0
- optimum/rbln/transformers/models/gpt2/gpt2_architecture.py +201 -166
- optimum/rbln/transformers/models/gpt2/modeling_gpt2.py +10 -257
- optimum/rbln/transformers/models/llama/llama_architecture.py +3 -610
- optimum/rbln/transformers/models/llama/modeling_llama.py +12 -440
- 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 +10 -325
- optimum/rbln/transformers/models/mistral/__init__.py +24 -0
- optimum/rbln/transformers/models/mistral/mistral_architecture.py +29 -0
- optimum/rbln/transformers/models/mistral/modeling_mistral.py +68 -0
- optimum/rbln/transformers/models/wav2vec2/modeling_wav2vec2.py +1 -1
- 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 +131 -0
- optimum/rbln/transformers/utils/__init__.py +0 -0
- optimum/rbln/transformers/utils/rbln_quantization.py +109 -0
- optimum/rbln/utils/import_utils.py +1 -4
- optimum/rbln/utils/runtime_utils.py +2 -1
- {optimum_rbln-0.1.7.dist-info → optimum_rbln-0.1.9.dist-info}/METADATA +11 -5
- {optimum_rbln-0.1.7.dist-info → optimum_rbln-0.1.9.dist-info}/RECORD +48 -35
- optimum/rbln/transformers/models/llama/llama_architecture_cb.py +0 -764
- {optimum_rbln-0.1.7.dist-info → optimum_rbln-0.1.9.dist-info}/WHEEL +0 -0
- {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
|