optimum-rbln 0.8.1rc0__py3-none-any.whl → 0.8.2__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.
Potentially problematic release.
This version of optimum-rbln might be problematic. Click here for more details.
- optimum/rbln/__init__.py +58 -9
- optimum/rbln/__version__.py +2 -2
- optimum/rbln/configuration_utils.py +24 -5
- optimum/rbln/diffusers/configurations/models/__init__.py +1 -1
- optimum/rbln/diffusers/configurations/models/configuration_autoencoder_kl.py +2 -2
- optimum/rbln/diffusers/configurations/models/configuration_autoencoder_kl_cosmos.py +5 -3
- optimum/rbln/diffusers/configurations/models/configuration_controlnet.py +2 -2
- optimum/rbln/diffusers/configurations/models/configuration_prior_transformer.py +2 -2
- optimum/rbln/diffusers/configurations/models/{configuration_cosmos_transformer.py → configuration_transformer_cosmos.py} +7 -2
- optimum/rbln/diffusers/configurations/models/configuration_transformer_sd3.py +2 -2
- optimum/rbln/diffusers/configurations/models/configuration_unet_2d_condition.py +2 -2
- optimum/rbln/diffusers/configurations/models/configuration_vq_model.py +2 -2
- optimum/rbln/diffusers/configurations/pipelines/configuration_controlnet.py +3 -3
- optimum/rbln/diffusers/configurations/pipelines/configuration_cosmos.py +10 -6
- optimum/rbln/diffusers/configurations/pipelines/configuration_kandinsky2_2.py +4 -4
- optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion.py +2 -2
- optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_3.py +2 -2
- optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_xl.py +2 -2
- optimum/rbln/diffusers/modeling_diffusers.py +4 -5
- optimum/rbln/diffusers/models/__init__.py +3 -13
- optimum/rbln/diffusers/models/autoencoders/autoencoder_kl.py +1 -0
- optimum/rbln/diffusers/models/autoencoders/autoencoder_kl_cosmos.py +1 -0
- optimum/rbln/diffusers/models/autoencoders/vq_model.py +1 -0
- optimum/rbln/diffusers/models/transformers/transformer_cosmos.py +1 -1
- optimum/rbln/diffusers/pipelines/__init__.py +1 -5
- optimum/rbln/diffusers/pipelines/cosmos/configuration_cosmos_guardrail.py +12 -4
- optimum/rbln/diffusers/pipelines/cosmos/cosmos_guardrail.py +4 -26
- optimum/rbln/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py +2 -2
- optimum/rbln/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py +2 -2
- optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py +1 -6
- optimum/rbln/modeling.py +4 -5
- optimum/rbln/modeling_base.py +18 -14
- optimum/rbln/ops/kv_cache_update.py +5 -0
- optimum/rbln/ops/linear.py +7 -0
- optimum/rbln/transformers/__init__.py +60 -0
- optimum/rbln/transformers/configuration_generic.py +4 -4
- optimum/rbln/transformers/modeling_attention_utils.py +252 -0
- optimum/rbln/transformers/modeling_generic.py +1 -4
- optimum/rbln/transformers/models/__init__.py +45 -30
- optimum/rbln/transformers/models/bart/bart_architecture.py +2 -7
- optimum/rbln/transformers/models/blip_2/configuration_blip_2.py +2 -2
- optimum/rbln/transformers/models/blip_2/modeling_blip_2.py +1 -5
- optimum/rbln/transformers/models/clip/configuration_clip.py +14 -3
- optimum/rbln/transformers/models/clip/modeling_clip.py +123 -28
- optimum/rbln/transformers/models/colpali/colpali_architecture.py +1 -4
- optimum/rbln/transformers/models/colpali/configuration_colpali.py +2 -2
- optimum/rbln/transformers/models/colpali/modeling_colpali.py +2 -10
- optimum/rbln/transformers/models/decoderonly/__init__.py +2 -2
- optimum/rbln/transformers/models/decoderonly/configuration_decoderonly.py +214 -45
- optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py +323 -454
- optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py +579 -362
- optimum/rbln/transformers/models/exaone/exaone_architecture.py +17 -42
- optimum/rbln/transformers/models/gemma/__init__.py +2 -2
- optimum/rbln/transformers/models/gemma/configuration_gemma.py +9 -1
- optimum/rbln/transformers/models/gemma/gemma_architecture.py +3 -44
- optimum/rbln/transformers/models/gemma/modeling_gemma.py +22 -1
- optimum/rbln/transformers/models/gemma3/configuration_gemma3.py +21 -9
- optimum/rbln/transformers/models/gemma3/gemma3_architecture.py +9 -63
- optimum/rbln/transformers/models/gemma3/modeling_gemma3.py +200 -292
- optimum/rbln/transformers/models/gpt2/__init__.py +2 -2
- optimum/rbln/transformers/models/gpt2/configuration_gpt2.py +31 -3
- optimum/rbln/transformers/models/gpt2/gpt2_architecture.py +19 -24
- optimum/rbln/transformers/models/gpt2/modeling_gpt2.py +18 -1
- optimum/rbln/transformers/models/idefics3/configuration_idefics3.py +2 -2
- optimum/rbln/transformers/models/idefics3/modeling_idefics3.py +3 -9
- optimum/rbln/transformers/models/llama/__init__.py +2 -2
- optimum/rbln/transformers/models/llama/configuration_llama.py +9 -1
- optimum/rbln/transformers/models/llama/modeling_llama.py +22 -1
- optimum/rbln/transformers/models/llava/__init__.py +16 -0
- optimum/rbln/transformers/models/llava/configuration_llava.py +54 -0
- optimum/rbln/transformers/models/llava/modeling_llava.py +419 -0
- optimum/rbln/transformers/models/llava_next/configuration_llava_next.py +20 -3
- optimum/rbln/transformers/models/llava_next/modeling_llava_next.py +6 -16
- optimum/rbln/transformers/models/midm/midm_architecture.py +14 -22
- optimum/rbln/transformers/models/mistral/__init__.py +2 -2
- optimum/rbln/transformers/models/mistral/configuration_mistral.py +9 -1
- optimum/rbln/transformers/models/mistral/mistral_architecture.py +1 -1
- optimum/rbln/transformers/models/mistral/modeling_mistral.py +26 -3
- optimum/rbln/transformers/models/opt/__init__.py +2 -2
- optimum/rbln/transformers/models/opt/configuration_opt.py +8 -1
- optimum/rbln/transformers/models/opt/modeling_opt.py +41 -1
- optimum/rbln/transformers/models/opt/opt_architecture.py +16 -25
- optimum/rbln/transformers/models/pegasus/__init__.py +17 -0
- optimum/rbln/transformers/models/pegasus/configuration_pegasus.py +34 -0
- optimum/rbln/transformers/models/pegasus/modeling_pegasus.py +69 -0
- optimum/rbln/transformers/models/pegasus/pegasus_architecture.py +161 -0
- optimum/rbln/transformers/models/phi/__init__.py +2 -2
- optimum/rbln/transformers/models/phi/configuration_phi.py +9 -1
- optimum/rbln/transformers/models/phi/modeling_phi.py +10 -1
- optimum/rbln/transformers/models/phi/phi_architecture.py +16 -22
- optimum/rbln/transformers/models/pixtral/__init__.py +16 -0
- optimum/rbln/transformers/models/pixtral/configuration_pixtral.py +43 -0
- optimum/rbln/transformers/models/pixtral/modeling_pixtral.py +315 -0
- optimum/rbln/transformers/models/pixtral/pixtral_architecture.py +73 -0
- optimum/rbln/transformers/models/qwen2/__init__.py +2 -2
- optimum/rbln/transformers/models/qwen2/configuration_qwen2.py +9 -1
- optimum/rbln/transformers/models/qwen2/modeling_qwen2.py +27 -1
- optimum/rbln/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py +3 -3
- optimum/rbln/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py +5 -15
- optimum/rbln/transformers/models/qwen2_5_vl/qwen2_5_vl_architecture.py +1 -4
- optimum/rbln/transformers/models/qwen3/__init__.py +16 -0
- optimum/rbln/transformers/models/qwen3/configuration_qwen3.py +71 -0
- optimum/rbln/transformers/models/qwen3/modeling_qwen3.py +133 -0
- optimum/rbln/transformers/models/qwen3/qwen3_architecture.py +31 -0
- optimum/rbln/transformers/models/seq2seq/configuration_seq2seq.py +2 -12
- optimum/rbln/transformers/models/seq2seq/modeling_seq2seq.py +3 -1
- optimum/rbln/transformers/models/siglip/__init__.py +2 -6
- optimum/rbln/transformers/models/siglip/modeling_siglip.py +2 -2
- optimum/rbln/transformers/models/time_series_transformer/configuration_time_series_transformer.py +2 -2
- optimum/rbln/transformers/models/time_series_transformer/modeling_time_series_transformer.py +3 -5
- optimum/rbln/transformers/models/whisper/configuration_whisper.py +3 -12
- optimum/rbln/transformers/models/whisper/modeling_whisper.py +8 -2
- optimum/rbln/transformers/models/xlm_roberta/__init__.py +2 -8
- optimum/rbln/utils/depreacate_utils.py +16 -0
- optimum/rbln/utils/hub.py +8 -47
- optimum/rbln/utils/runtime_utils.py +31 -5
- {optimum_rbln-0.8.1rc0.dist-info → optimum_rbln-0.8.2.dist-info}/METADATA +1 -1
- {optimum_rbln-0.8.1rc0.dist-info → optimum_rbln-0.8.2.dist-info}/RECORD +120 -103
- {optimum_rbln-0.8.1rc0.dist-info → optimum_rbln-0.8.2.dist-info}/WHEEL +0 -0
- {optimum_rbln-0.8.1rc0.dist-info → optimum_rbln-0.8.2.dist-info}/licenses/LICENSE +0 -0
|
@@ -0,0 +1,315 @@
|
|
|
1
|
+
# Copyright 2025 Rebellions Inc. All rights reserved.
|
|
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
|
+
from pathlib import Path
|
|
16
|
+
from typing import TYPE_CHECKING, Any, Optional, Tuple, Union
|
|
17
|
+
|
|
18
|
+
import rebel
|
|
19
|
+
import torch
|
|
20
|
+
import torch.nn as nn
|
|
21
|
+
from transformers import PixtralVisionConfig, PixtralVisionModel
|
|
22
|
+
from transformers.modeling_outputs import BaseModelOutput
|
|
23
|
+
from transformers.modeling_utils import no_init_weights
|
|
24
|
+
from transformers.models.pixtral.modeling_pixtral import PixtralRMSNorm, PixtralRotaryEmbedding
|
|
25
|
+
|
|
26
|
+
from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
|
|
27
|
+
from ....modeling import RBLNModel
|
|
28
|
+
from ....utils.logging import get_logger
|
|
29
|
+
from ....utils.runtime_utils import RBLNPytorchRuntime
|
|
30
|
+
from .configuration_pixtral import RBLNPixtralVisionModelConfig
|
|
31
|
+
from .pixtral_architecture import PixtralAttention
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
logger = get_logger(__name__)
|
|
35
|
+
|
|
36
|
+
if TYPE_CHECKING:
|
|
37
|
+
from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
|
|
38
|
+
|
|
39
|
+
from ....diffusers.modeling_diffusers import RBLNDiffusionMixin, RBLNDiffusionMixinConfig
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
class RBLNRuntimePixtralVisionModel(RBLNPytorchRuntime):
|
|
43
|
+
mandatory_members = ["main_input_name"]
|
|
44
|
+
|
|
45
|
+
def __init__(
|
|
46
|
+
self,
|
|
47
|
+
runtime: rebel.Runtime,
|
|
48
|
+
config: PixtralVisionConfig,
|
|
49
|
+
rbln_config: RBLNPixtralVisionModelConfig,
|
|
50
|
+
**kwargs: Any,
|
|
51
|
+
) -> None:
|
|
52
|
+
super().__init__(runtime, **kwargs)
|
|
53
|
+
self.patch_positional_embedding = PixtralRotaryEmbedding(config)
|
|
54
|
+
self.patch_size = config.patch_size
|
|
55
|
+
self.image_size = config.image_size
|
|
56
|
+
self.hidden_size = config.hidden_size
|
|
57
|
+
self.max_image_size = rbln_config.max_image_size
|
|
58
|
+
|
|
59
|
+
def forward(
|
|
60
|
+
self,
|
|
61
|
+
pixel_values: torch.Tensor,
|
|
62
|
+
image_sizes: torch.Tensor,
|
|
63
|
+
output_hidden_states: Optional[bool] = None,
|
|
64
|
+
return_dict: Optional[bool] = None,
|
|
65
|
+
**kwargs,
|
|
66
|
+
):
|
|
67
|
+
if pixel_values.shape[2] > self.max_image_size[0] or pixel_values.shape[3] > self.max_image_size[1]:
|
|
68
|
+
raise ValueError("The height() and width of pixel_values can't be larger than max_image_size.")
|
|
69
|
+
|
|
70
|
+
if pixel_values.shape[2] != self.max_image_size[0] or pixel_values.shape[3] != self.max_image_size[1]:
|
|
71
|
+
padded_pixel_values = [
|
|
72
|
+
torch.nn.functional.pad(
|
|
73
|
+
image,
|
|
74
|
+
pad=(
|
|
75
|
+
0,
|
|
76
|
+
self.max_image_size[1] - pixel_values.shape[3],
|
|
77
|
+
0,
|
|
78
|
+
self.max_image_size[0] - pixel_values.shape[2],
|
|
79
|
+
),
|
|
80
|
+
)
|
|
81
|
+
for image in pixel_values
|
|
82
|
+
]
|
|
83
|
+
pixel_values = torch.stack(padded_pixel_values)
|
|
84
|
+
|
|
85
|
+
batch_size, _, H_max, W_max = pixel_values.shape
|
|
86
|
+
H_max_p = H_max // self.patch_size
|
|
87
|
+
W_max_p = W_max // self.patch_size
|
|
88
|
+
|
|
89
|
+
final_hidden_states = None
|
|
90
|
+
|
|
91
|
+
last_hidden_state_list = []
|
|
92
|
+
if output_hidden_states:
|
|
93
|
+
batch_hidden_states_list = []
|
|
94
|
+
|
|
95
|
+
for i in range(batch_size):
|
|
96
|
+
h_patched_original = image_sizes[i, 0] // self.patch_size
|
|
97
|
+
w_patched_original = image_sizes[i, 1] // self.patch_size
|
|
98
|
+
|
|
99
|
+
single_pixel_values = pixel_values[i : i + 1]
|
|
100
|
+
patch_embed = self.patch_conv(single_pixel_values)
|
|
101
|
+
patch_embed_seq = patch_embed[:, :, :h_patched_original, :w_patched_original].flatten(2).transpose(1, 2)
|
|
102
|
+
patch_embed_seq = self.ln_pre(patch_embed_seq)
|
|
103
|
+
patch_embed_seq = nn.functional.pad(
|
|
104
|
+
patch_embed_seq, (0, 0, 0, H_max_p * W_max_p - patch_embed_seq.shape[1]), "constant", value=0
|
|
105
|
+
)
|
|
106
|
+
|
|
107
|
+
max_w_from_config = self.image_size // self.patch_size
|
|
108
|
+
mesh = torch.meshgrid(torch.arange(h_patched_original), torch.arange(w_patched_original), indexing="ij")
|
|
109
|
+
h_grid, v_grid = torch.stack(mesh, dim=-1).reshape(-1, 2).chunk(2, -1)
|
|
110
|
+
ids = h_grid * max_w_from_config + v_grid
|
|
111
|
+
position_ids = ids[:, 0]
|
|
112
|
+
|
|
113
|
+
position_embeddings = self.patch_positional_embedding(patch_embed_seq, position_ids)
|
|
114
|
+
cos = nn.functional.pad(
|
|
115
|
+
position_embeddings[0],
|
|
116
|
+
(0, 0, 0, H_max_p * W_max_p - position_embeddings[0].shape[0]),
|
|
117
|
+
"constant",
|
|
118
|
+
value=0,
|
|
119
|
+
)
|
|
120
|
+
sin = nn.functional.pad(
|
|
121
|
+
position_embeddings[1],
|
|
122
|
+
(0, 0, 0, H_max_p * W_max_p - position_embeddings[1].shape[0]),
|
|
123
|
+
"constant",
|
|
124
|
+
value=0,
|
|
125
|
+
)
|
|
126
|
+
|
|
127
|
+
attention_mask = torch.full(
|
|
128
|
+
(1, patch_embed_seq.shape[-2]), fill_value=torch.finfo(patch_embed_seq.dtype).min
|
|
129
|
+
)
|
|
130
|
+
attention_mask[:, : h_patched_original * w_patched_original] = 0
|
|
131
|
+
|
|
132
|
+
transformer_output = super().forward(patch_embed_seq, attention_mask, cos, sin)
|
|
133
|
+
|
|
134
|
+
last_hidden_state_list.append(transformer_output[0][:, : h_patched_original * w_patched_original, :])
|
|
135
|
+
hidden_states = transformer_output[1:]
|
|
136
|
+
|
|
137
|
+
if output_hidden_states:
|
|
138
|
+
batch_hidden_states_list.append(
|
|
139
|
+
[hidden_state[:, : h_patched_original * w_patched_original, :] for hidden_state in hidden_states]
|
|
140
|
+
)
|
|
141
|
+
|
|
142
|
+
final_last_hidden_state = torch.cat(last_hidden_state_list, dim=1)
|
|
143
|
+
|
|
144
|
+
if output_hidden_states:
|
|
145
|
+
hidden_states = [
|
|
146
|
+
torch.cat(
|
|
147
|
+
[batch_hidden_states[layer_idx] for batch_hidden_states in batch_hidden_states_list],
|
|
148
|
+
dim=1,
|
|
149
|
+
)
|
|
150
|
+
for layer_idx in range(len(batch_hidden_states_list[0]))
|
|
151
|
+
]
|
|
152
|
+
|
|
153
|
+
final_hidden_states = tuple(hidden_states)
|
|
154
|
+
|
|
155
|
+
if not return_dict:
|
|
156
|
+
return tuple(v for v in (final_last_hidden_state, final_hidden_states) if v is not None)
|
|
157
|
+
|
|
158
|
+
# TODO: output_attentions
|
|
159
|
+
return BaseModelOutput(
|
|
160
|
+
last_hidden_state=final_last_hidden_state,
|
|
161
|
+
hidden_states=final_hidden_states,
|
|
162
|
+
)
|
|
163
|
+
|
|
164
|
+
|
|
165
|
+
class _PixtralVisionModel(torch.nn.Module):
|
|
166
|
+
def __init__(self, model: PixtralVisionModel, output_hidden_states: bool):
|
|
167
|
+
super().__init__()
|
|
168
|
+
self.transformer = self.convert_to_rbln_pixtral_vision_model(model)
|
|
169
|
+
self.output_hidden_states = output_hidden_states
|
|
170
|
+
|
|
171
|
+
def convert_to_rbln_pixtral_vision_model(self, model: nn.Module):
|
|
172
|
+
for layer in model.transformer.layers:
|
|
173
|
+
layer.attention = PixtralAttention(layer.attention)
|
|
174
|
+
return model.transformer
|
|
175
|
+
|
|
176
|
+
def forward(self, patch_embeds, attention_mask, position_embeddings_1, position_embeddings_2):
|
|
177
|
+
output = self.transformer(
|
|
178
|
+
inputs_embeds=patch_embeds,
|
|
179
|
+
attention_mask=attention_mask,
|
|
180
|
+
position_embeddings=(position_embeddings_1, position_embeddings_2),
|
|
181
|
+
output_hidden_states=self.output_hidden_states,
|
|
182
|
+
return_dict=False,
|
|
183
|
+
)
|
|
184
|
+
return output
|
|
185
|
+
|
|
186
|
+
|
|
187
|
+
class RBLNPixtralVisionModel(RBLNModel):
|
|
188
|
+
"""
|
|
189
|
+
RBLN optimized Pixtral vision encoder model.
|
|
190
|
+
|
|
191
|
+
This class provides hardware-accelerated inference for Pixtral vision encoders
|
|
192
|
+
on RBLN devices, supporting image encoding for multimodal tasks.
|
|
193
|
+
"""
|
|
194
|
+
|
|
195
|
+
def __post_init__(self, **kwargs):
|
|
196
|
+
artifacts = torch.load(self.model_save_dir / self.subfolder / "torch_artifacts.pth", weights_only=False)
|
|
197
|
+
with no_init_weights():
|
|
198
|
+
self.patch_conv = nn.Conv2d(
|
|
199
|
+
in_channels=self.config.num_channels,
|
|
200
|
+
out_channels=self.config.hidden_size,
|
|
201
|
+
kernel_size=self.config.patch_size,
|
|
202
|
+
stride=self.config.patch_size,
|
|
203
|
+
bias=False,
|
|
204
|
+
)
|
|
205
|
+
self.ln_pre = PixtralRMSNorm(self.config.hidden_size, eps=1e-5)
|
|
206
|
+
self.patch_conv.load_state_dict(artifacts["patch_conv"])
|
|
207
|
+
self.ln_pre.load_state_dict(artifacts["ln_pre"])
|
|
208
|
+
self.model = RBLNRuntimePixtralVisionModel(
|
|
209
|
+
self.model[0],
|
|
210
|
+
main_input_name="pixel_values",
|
|
211
|
+
config=self.config,
|
|
212
|
+
rbln_config=self.rbln_config,
|
|
213
|
+
patch_conv=self.patch_conv,
|
|
214
|
+
ln_pre=self.ln_pre,
|
|
215
|
+
)
|
|
216
|
+
|
|
217
|
+
@classmethod
|
|
218
|
+
def save_torch_artifacts(
|
|
219
|
+
cls,
|
|
220
|
+
model: "PreTrainedModel",
|
|
221
|
+
save_dir_path: Path,
|
|
222
|
+
subfolder: str,
|
|
223
|
+
rbln_config: RBLNModelConfig,
|
|
224
|
+
):
|
|
225
|
+
save_dict = {}
|
|
226
|
+
save_dict["patch_conv"] = model.get_input_embeddings().state_dict()
|
|
227
|
+
save_dict["ln_pre"] = model.ln_pre.state_dict()
|
|
228
|
+
torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
|
|
229
|
+
|
|
230
|
+
@classmethod
|
|
231
|
+
def wrap_model_if_needed(
|
|
232
|
+
cls, model: torch.nn.Module, rbln_config: RBLNPixtralVisionModelConfig
|
|
233
|
+
) -> torch.nn.Module:
|
|
234
|
+
wrapper_cfg = {
|
|
235
|
+
"output_hidden_states": rbln_config.output_hidden_states,
|
|
236
|
+
}
|
|
237
|
+
return _PixtralVisionModel(model, **wrapper_cfg).eval()
|
|
238
|
+
|
|
239
|
+
@classmethod
|
|
240
|
+
def update_rbln_config_using_pipe(
|
|
241
|
+
cls, pipe: "RBLNDiffusionMixin", rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
|
|
242
|
+
) -> "RBLNDiffusionMixinConfig":
|
|
243
|
+
return rbln_config
|
|
244
|
+
|
|
245
|
+
@classmethod
|
|
246
|
+
def _update_rbln_config(
|
|
247
|
+
cls,
|
|
248
|
+
preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
|
|
249
|
+
model: Optional["PreTrainedModel"] = None,
|
|
250
|
+
model_config: "PixtralVisionConfig" = None,
|
|
251
|
+
rbln_config: Optional[RBLNPixtralVisionModelConfig] = None,
|
|
252
|
+
) -> RBLNPixtralVisionModelConfig:
|
|
253
|
+
if rbln_config.max_image_size is None:
|
|
254
|
+
rbln_config.max_image_size = (model_config.image_size, model_config.image_size)
|
|
255
|
+
|
|
256
|
+
if rbln_config.output_hidden_states is None:
|
|
257
|
+
rbln_config.output_hidden_states = getattr(model_config, "output_hidden_states", False)
|
|
258
|
+
|
|
259
|
+
num_total_patches = (rbln_config.max_image_size[0] // model_config.patch_size) * (
|
|
260
|
+
rbln_config.max_image_size[1] // model_config.patch_size
|
|
261
|
+
)
|
|
262
|
+
|
|
263
|
+
rbln_compile_config = RBLNCompileConfig(
|
|
264
|
+
input_info=[
|
|
265
|
+
(
|
|
266
|
+
"patch_embeds",
|
|
267
|
+
[1, num_total_patches, model_config.hidden_size],
|
|
268
|
+
"float32",
|
|
269
|
+
),
|
|
270
|
+
("attention_mask", [1, num_total_patches], "float32"),
|
|
271
|
+
(
|
|
272
|
+
"position_embeddings_1",
|
|
273
|
+
[
|
|
274
|
+
num_total_patches,
|
|
275
|
+
model_config.head_dim,
|
|
276
|
+
],
|
|
277
|
+
"float32",
|
|
278
|
+
),
|
|
279
|
+
(
|
|
280
|
+
"position_embeddings_2",
|
|
281
|
+
[
|
|
282
|
+
num_total_patches,
|
|
283
|
+
model_config.head_dim,
|
|
284
|
+
],
|
|
285
|
+
"float32",
|
|
286
|
+
),
|
|
287
|
+
]
|
|
288
|
+
)
|
|
289
|
+
|
|
290
|
+
rbln_config.set_compile_cfgs([rbln_compile_config])
|
|
291
|
+
return rbln_config
|
|
292
|
+
|
|
293
|
+
def forward(
|
|
294
|
+
self,
|
|
295
|
+
pixel_values: Optional[torch.FloatTensor] = None,
|
|
296
|
+
image_sizes: Optional[torch.FloatTensor] = None,
|
|
297
|
+
output_hidden_states: Optional[bool] = None,
|
|
298
|
+
return_dict: bool = True,
|
|
299
|
+
**kwargs,
|
|
300
|
+
) -> Union[Tuple, BaseModelOutput]:
|
|
301
|
+
output_hidden_states = (
|
|
302
|
+
output_hidden_states if output_hidden_states is not None else self.rbln_config.output_hidden_states
|
|
303
|
+
)
|
|
304
|
+
|
|
305
|
+
if output_hidden_states != self.rbln_config.output_hidden_states:
|
|
306
|
+
raise ValueError(
|
|
307
|
+
f"Variable output_hidden_states {output_hidden_states} is not equal to rbln_config.output_hidden_states {self.rbln_config.output_hidden_states} "
|
|
308
|
+
f"Please compile again with the correct argument."
|
|
309
|
+
)
|
|
310
|
+
|
|
311
|
+
output = self.model(
|
|
312
|
+
pixel_values, image_sizes, output_hidden_states=output_hidden_states, return_dict=return_dict
|
|
313
|
+
)
|
|
314
|
+
|
|
315
|
+
return output
|
|
@@ -0,0 +1,73 @@
|
|
|
1
|
+
# Copyright 2025 Rebellions Inc. All rights reserved.
|
|
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
|
+
from typing import Optional, Tuple
|
|
16
|
+
|
|
17
|
+
import torch
|
|
18
|
+
import torch.nn as nn
|
|
19
|
+
|
|
20
|
+
from ..decoderonly.decoderonly_architecture import apply_rotary_pos_emb
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class PixtralAttention(nn.Module):
|
|
24
|
+
def __init__(self, self_attention):
|
|
25
|
+
super().__init__()
|
|
26
|
+
self.original_model = self_attention
|
|
27
|
+
self.num_heads = getattr(self.original_model, "num_heads", None) or getattr(
|
|
28
|
+
self.original_model.config, "num_attention_heads"
|
|
29
|
+
)
|
|
30
|
+
self.head_dim = self.original_model.head_dim
|
|
31
|
+
self.scaling = self.head_dim**-0.5
|
|
32
|
+
|
|
33
|
+
self.__post_init__()
|
|
34
|
+
|
|
35
|
+
def __post_init__(self):
|
|
36
|
+
self.q_proj = self.original_model.q_proj
|
|
37
|
+
self.k_proj = self.original_model.k_proj
|
|
38
|
+
self.v_proj = self.original_model.v_proj
|
|
39
|
+
self.o_proj = self.original_model.o_proj
|
|
40
|
+
|
|
41
|
+
def forward(
|
|
42
|
+
self,
|
|
43
|
+
hidden_states: torch.Tensor,
|
|
44
|
+
attention_mask: Optional[torch.Tensor] = None,
|
|
45
|
+
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
|
|
46
|
+
output_attentions: Optional[bool] = False,
|
|
47
|
+
):
|
|
48
|
+
batch_size, patches, _ = hidden_states.size()
|
|
49
|
+
|
|
50
|
+
query_states = self.q_proj(hidden_states)
|
|
51
|
+
key_states = self.k_proj(hidden_states)
|
|
52
|
+
value_states = self.v_proj(hidden_states)
|
|
53
|
+
|
|
54
|
+
# TODO: return output attention
|
|
55
|
+
query_states = query_states.view(batch_size, patches, 1, self.num_heads, self.head_dim).transpose(1, 3)
|
|
56
|
+
key_states = key_states.view(batch_size, patches, 1, self.num_heads, self.head_dim).transpose(1, 3)
|
|
57
|
+
value_states = value_states.view(batch_size, patches, 1, self.num_heads, self.head_dim).transpose(1, 3)
|
|
58
|
+
|
|
59
|
+
cos, sin = position_embeddings
|
|
60
|
+
cos = cos[None, None, None, :, :]
|
|
61
|
+
sin = sin[None, None, None, :, :]
|
|
62
|
+
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
|
|
63
|
+
|
|
64
|
+
attn_weights = torch.matmul(query_states, key_states.transpose(3, 4)) * self.scaling
|
|
65
|
+
attn_weights = attn_weights + attention_mask
|
|
66
|
+
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32)
|
|
67
|
+
attn_output = torch.matmul(attn_weights, value_states)
|
|
68
|
+
attn_output = attn_output.transpose(1, 3)
|
|
69
|
+
|
|
70
|
+
attn_output = attn_output.reshape(batch_size, patches, -1)
|
|
71
|
+
attn_output = self.o_proj(attn_output)
|
|
72
|
+
|
|
73
|
+
return attn_output, _
|
|
@@ -12,5 +12,5 @@
|
|
|
12
12
|
# See the License for the specific language governing permissions and
|
|
13
13
|
# limitations under the License.
|
|
14
14
|
|
|
15
|
-
from .configuration_qwen2 import RBLNQwen2ForCausalLMConfig
|
|
16
|
-
from .modeling_qwen2 import RBLNQwen2ForCausalLM
|
|
15
|
+
from .configuration_qwen2 import RBLNQwen2ForCausalLMConfig, RBLNQwen2ModelConfig
|
|
16
|
+
from .modeling_qwen2 import RBLNQwen2ForCausalLM, RBLNQwen2Model
|
|
@@ -12,7 +12,7 @@
|
|
|
12
12
|
# See the License for the specific language governing permissions and
|
|
13
13
|
# limitations under the License.
|
|
14
14
|
|
|
15
|
-
from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
|
|
15
|
+
from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
|
|
16
16
|
|
|
17
17
|
|
|
18
18
|
class RBLNQwen2ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
|
|
@@ -40,3 +40,11 @@ class RBLNQwen2ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
|
|
|
40
40
|
)
|
|
41
41
|
```
|
|
42
42
|
"""
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
class RBLNQwen2ModelConfig(RBLNDecoderOnlyModelConfig):
|
|
46
|
+
"""
|
|
47
|
+
Configuration class for RBLN Qwen2 models.
|
|
48
|
+
|
|
49
|
+
This class is an alias of RBLNDecoderOnlyModelConfig.
|
|
50
|
+
"""
|
|
@@ -15,7 +15,11 @@
|
|
|
15
15
|
from transformers import PretrainedConfig
|
|
16
16
|
|
|
17
17
|
from ....utils import logging
|
|
18
|
-
from ...models.decoderonly import
|
|
18
|
+
from ...models.decoderonly import (
|
|
19
|
+
RBLNDecoderOnlyModel,
|
|
20
|
+
RBLNDecoderOnlyModelForCausalLM,
|
|
21
|
+
RBLNDecoderOnlyModelForCausalLMConfig,
|
|
22
|
+
)
|
|
19
23
|
from .qwen2_architecture import QWEN2Wrapper
|
|
20
24
|
|
|
21
25
|
|
|
@@ -95,3 +99,25 @@ class RBLNQwen2ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
|
|
|
95
99
|
rbln_config.sliding_window = model_config.sliding_window
|
|
96
100
|
rbln_config.sliding_window_layers = list(range(model_config.num_hidden_layers))
|
|
97
101
|
return rbln_config
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
class RBLNQwen2Model(RBLNDecoderOnlyModel):
|
|
105
|
+
"""
|
|
106
|
+
The Qwen2 Model transformer without a language modeling head.
|
|
107
|
+
This model inherits from [`RBLNDecoderOnlyModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
|
|
108
|
+
"""
|
|
109
|
+
|
|
110
|
+
_decoder_wrapper_cls = QWEN2Wrapper
|
|
111
|
+
|
|
112
|
+
@classmethod
|
|
113
|
+
def _update_sliding_window_config(
|
|
114
|
+
cls, model_config: PretrainedConfig, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig
|
|
115
|
+
):
|
|
116
|
+
# https://github.com/huggingface/transformers/issues/35896
|
|
117
|
+
# There seems to be a bug in transformers(v4.52.4). Therefore, similar to when attn_implementation is eager,
|
|
118
|
+
# we set all layers to use sliding window in this version. This should be updated once the bug is fixed.
|
|
119
|
+
|
|
120
|
+
rbln_config.cache_impl = "sliding_window"
|
|
121
|
+
rbln_config.sliding_window = model_config.sliding_window
|
|
122
|
+
rbln_config.sliding_window_layers = list(range(model_config.num_hidden_layers))
|
|
123
|
+
return rbln_config
|
|
@@ -12,7 +12,7 @@
|
|
|
12
12
|
# See the License for the specific language governing permissions and
|
|
13
13
|
# limitations under the License.
|
|
14
14
|
|
|
15
|
-
from typing import Any,
|
|
15
|
+
from typing import Any, List, Optional, Union
|
|
16
16
|
|
|
17
17
|
from ....configuration_utils import RBLNModelConfig
|
|
18
18
|
from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
|
|
@@ -33,7 +33,7 @@ class RBLNQwen2_5_VLForConditionalGenerationConfig(RBLNDecoderOnlyModelForCausal
|
|
|
33
33
|
self,
|
|
34
34
|
visual: Optional[RBLNModelConfig] = None,
|
|
35
35
|
use_inputs_embeds: bool = True,
|
|
36
|
-
**kwargs:
|
|
36
|
+
**kwargs: Any,
|
|
37
37
|
):
|
|
38
38
|
super().__init__(use_inputs_embeds=use_inputs_embeds, **kwargs)
|
|
39
39
|
if not self.use_inputs_embeds:
|
|
@@ -53,7 +53,7 @@ class RBLNQwen2_5_VisionTransformerPretrainedModelConfig(RBLNModelConfig):
|
|
|
53
53
|
mechanisms for processing images and videos.
|
|
54
54
|
"""
|
|
55
55
|
|
|
56
|
-
def __init__(self, max_seq_lens: Union[int, List[int]] = None, **kwargs:
|
|
56
|
+
def __init__(self, max_seq_lens: Union[int, List[int]] = None, **kwargs: Any):
|
|
57
57
|
"""
|
|
58
58
|
Args:
|
|
59
59
|
max_seq_lens (Optional[Union[int, List[int]]]): Maximum sequence lengths for Vision
|
|
@@ -17,12 +17,7 @@ from pathlib import Path
|
|
|
17
17
|
from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple, Union
|
|
18
18
|
|
|
19
19
|
import torch
|
|
20
|
-
from transformers import
|
|
21
|
-
AutoModelForVision2Seq,
|
|
22
|
-
PretrainedConfig,
|
|
23
|
-
PreTrainedModel,
|
|
24
|
-
Qwen2_5_VLForConditionalGeneration,
|
|
25
|
-
)
|
|
20
|
+
from transformers import AutoModelForVision2Seq, PretrainedConfig, PreTrainedModel, Qwen2_5_VLForConditionalGeneration
|
|
26
21
|
from transformers.modeling_utils import no_init_weights
|
|
27
22
|
from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
|
|
28
23
|
Qwen2_5_VisionPatchEmbed,
|
|
@@ -34,7 +29,7 @@ from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
|
|
|
34
29
|
from ....configuration_utils import RBLNCompileConfig
|
|
35
30
|
from ....modeling import RBLNModel
|
|
36
31
|
from ....utils.logging import get_logger
|
|
37
|
-
from ..decoderonly.modeling_decoderonly import
|
|
32
|
+
from ..decoderonly.modeling_decoderonly import RBLNDecoderOnlyForCausalLMOutput, RBLNDecoderOnlyModelForCausalLM
|
|
38
33
|
from .configuration_qwen2_5_vl import (
|
|
39
34
|
RBLNQwen2_5_VisionTransformerPretrainedModelConfig,
|
|
40
35
|
RBLNQwen2_5_VLForConditionalGenerationConfig,
|
|
@@ -45,12 +40,7 @@ from .qwen2_5_vl_architecture import Qwen2_5_VisionTransformerWrapper, Qwen2_5_V
|
|
|
45
40
|
logger = get_logger(__name__)
|
|
46
41
|
|
|
47
42
|
if TYPE_CHECKING:
|
|
48
|
-
from transformers import
|
|
49
|
-
AutoFeatureExtractor,
|
|
50
|
-
AutoProcessor,
|
|
51
|
-
AutoTokenizer,
|
|
52
|
-
PretrainedConfig,
|
|
53
|
-
)
|
|
43
|
+
from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig
|
|
54
44
|
|
|
55
45
|
|
|
56
46
|
class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
|
|
@@ -595,7 +585,7 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
|
|
|
595
585
|
generate_idx: Optional[torch.Tensor] = None,
|
|
596
586
|
return_dict: Optional[bool] = None,
|
|
597
587
|
**kwargs,
|
|
598
|
-
) ->
|
|
588
|
+
) -> RBLNDecoderOnlyForCausalLMOutput:
|
|
599
589
|
# Prefill
|
|
600
590
|
if cache_position is None:
|
|
601
591
|
inputs_embeds, position_embed, rope_deltas = self._preprocess_prefill(
|
|
@@ -637,7 +627,7 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
|
|
|
637
627
|
if not return_dict:
|
|
638
628
|
return logits, generate_idx
|
|
639
629
|
else:
|
|
640
|
-
return
|
|
630
|
+
return RBLNDecoderOnlyForCausalLMOutput(
|
|
641
631
|
logits=logits,
|
|
642
632
|
generate_idx=generate_idx,
|
|
643
633
|
)
|
|
@@ -4,10 +4,7 @@ from typing import Tuple
|
|
|
4
4
|
import torch
|
|
5
5
|
import torch.nn as nn
|
|
6
6
|
|
|
7
|
-
from ..decoderonly.decoderonly_architecture import
|
|
8
|
-
DecoderOnlyWrapper,
|
|
9
|
-
apply_rotary_pos_emb,
|
|
10
|
-
)
|
|
7
|
+
from ..decoderonly.decoderonly_architecture import DecoderOnlyWrapper, apply_rotary_pos_emb
|
|
11
8
|
|
|
12
9
|
|
|
13
10
|
class Qwen2_5_VisionTransformerWrapper(nn.Module):
|
|
@@ -0,0 +1,16 @@
|
|
|
1
|
+
# Copyright 2025 Rebellions Inc. All rights reserved.
|
|
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
|
+
from .configuration_qwen3 import RBLNQwen3ForCausalLMConfig, RBLNQwen3ModelConfig
|
|
16
|
+
from .modeling_qwen3 import RBLNQwen3ForCausalLM, RBLNQwen3Model
|
|
@@ -0,0 +1,71 @@
|
|
|
1
|
+
# Copyright 2025 Rebellions Inc. All rights reserved.
|
|
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
|
+
from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class RBLNQwen3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
|
|
19
|
+
"""
|
|
20
|
+
Configuration class for RBLN Qwen3 models.
|
|
21
|
+
|
|
22
|
+
This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
|
|
23
|
+
|
|
24
|
+
Example usage:
|
|
25
|
+
```python
|
|
26
|
+
from optimum.rbln import RBLNQwen3ForCausalLM, RBLNQwen3ForCausalLMConfig
|
|
27
|
+
|
|
28
|
+
# Create a configuration object
|
|
29
|
+
config = RBLNQwen3ForCausalLMConfig(
|
|
30
|
+
batch_size=1,
|
|
31
|
+
max_seq_len=40960,
|
|
32
|
+
tensor_parallel_size=4,
|
|
33
|
+
kvcache_partition_len=16384
|
|
34
|
+
)
|
|
35
|
+
|
|
36
|
+
# Use the configuration with from_pretrained
|
|
37
|
+
model = RBLNQwen3ForCausalLM.from_pretrained(
|
|
38
|
+
"Qwen/Qwen3-4B",
|
|
39
|
+
export=True,
|
|
40
|
+
rbln_config=config
|
|
41
|
+
)
|
|
42
|
+
```
|
|
43
|
+
"""
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
class RBLNQwen3ModelConfig(RBLNDecoderOnlyModelConfig):
|
|
47
|
+
"""
|
|
48
|
+
Configuration class for RBLN Qwen3 models.
|
|
49
|
+
|
|
50
|
+
This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
|
|
51
|
+
|
|
52
|
+
Example usage:
|
|
53
|
+
```python
|
|
54
|
+
from optimum.rbln import RBLNQwen3Model, RBLNQwen3ModelConfig
|
|
55
|
+
|
|
56
|
+
# Create a configuration object
|
|
57
|
+
config = RBLNQwen3ModelConfig(
|
|
58
|
+
batch_size=1,
|
|
59
|
+
max_seq_len=40960,
|
|
60
|
+
tensor_parallel_size=4,
|
|
61
|
+
kvcache_partition_len=16384
|
|
62
|
+
)
|
|
63
|
+
|
|
64
|
+
# Use the configuration with from_pretrained
|
|
65
|
+
model = RBLNQwen3Model.from_pretrained(
|
|
66
|
+
"Qwen/Qwen3-Embedding-4B",
|
|
67
|
+
export=True,
|
|
68
|
+
rbln_config=config
|
|
69
|
+
)
|
|
70
|
+
```
|
|
71
|
+
"""
|