|
from functools import wraps |
|
import warnings |
|
import torch |
|
import openvino as ov |
|
from pathlib import Path |
|
from typing import Tuple, Optional |
|
import types |
|
from transformers.modeling_outputs import BaseModelOutputWithPast |
|
|
|
try: |
|
from optimum.exporters.openvino.stateful import make_stateful |
|
from optimum.exporters.openvino.stateful import fuse_cache_reorder |
|
except ImportError: |
|
warnings.warn("We recommend to update optimum-intel for getting optimal performance") |
|
make_stateful = None |
|
fuse_cache_reorder = None |
|
|
|
|
|
def register_configs(): |
|
from optimum.exporters.tasks import TasksManager |
|
TasksManager._SUPPORTED_MODEL_TYPE["minicpm"] = TasksManager._SUPPORTED_MODEL_TYPE["llama"] |
|
TasksManager._SUPPORTED_MODEL_TYPE["qwen2"] = TasksManager._SUPPORTED_MODEL_TYPE["llama"] |
|
|
|
def patch_stateful(ov_model, model_type): |
|
key_value_input_names = [ |
|
key.get_any_name() for key in ov_model.inputs if any("key_values" in key_name for key_name in key.get_names()) |
|
] |
|
key_value_output_names = [ |
|
key.get_any_name() for key in ov_model.outputs if any("present" in key_name for key_name in key.get_names()) |
|
] |
|
not_kv_inputs = [ |
|
input for input in ov_model.inputs if not any(name in key_value_input_names for name in input.get_names()) |
|
] |
|
if not key_value_input_names or not key_value_output_names: |
|
return |
|
batch_dim = 1 if model_type == "chatglm" else 0 |
|
num_attention_heads = 1 |
|
|
|
fuse_cache_reorder(ov_model, not_kv_inputs, key_value_input_names, batch_dim) |
|
make_stateful( |
|
ov_model, not_kv_inputs, key_value_input_names, key_value_output_names, batch_dim, num_attention_heads, None |
|
) |
|
|
|
|
|
|
|
def flattenize_inputs(inputs): |
|
""" |
|
Helper function for making nested inputs flattens |
|
""" |
|
flatten_inputs = [] |
|
for input_data in inputs: |
|
if input_data is None: |
|
continue |
|
if isinstance(input_data, (list, tuple)): |
|
flatten_inputs.extend(flattenize_inputs(input_data)) |
|
else: |
|
flatten_inputs.append(input_data) |
|
return flatten_inputs |
|
|
|
|
|
def cleanup_torchscript_cache(): |
|
""" |
|
Helper for removing cached model representation |
|
""" |
|
torch._C._jit_clear_class_registry() |
|
torch.jit._recursive.concrete_type_store = torch.jit._recursive.ConcreteTypeStore() |
|
torch.jit._state._clear_class_state() |
|
|
|
|
|
def convert_mpt(pt_model: torch.nn.Module, model_path: Path): |
|
""" |
|
MPT model conversion function |
|
|
|
Params: |
|
pt_model: PyTorch model |
|
model_path: path for saving model |
|
Returns: |
|
None |
|
""" |
|
ov_out_path = Path(model_path) / "openvino_model.xml" |
|
pt_model.config.save_pretrained(ov_out_path.parent) |
|
pt_model.config.use_cache = True |
|
outs = pt_model( |
|
input_ids=torch.ones((1, 10), dtype=torch.long), |
|
attention_mask=torch.ones((1, 10), dtype=torch.long), |
|
) |
|
inputs = ["input_ids"] |
|
outputs = ["logits"] |
|
|
|
dynamic_shapes = {"input_ids": {0: "batch_size", 1: "seq_len"}, "attention_mask": {0: "batch_size", 1: "seq_len"}} |
|
for idx in range(len(outs.past_key_values)): |
|
inputs.extend([f"past_key_values.{idx}.key", f"past_key_values.{idx}.value"]) |
|
dynamic_shapes[inputs[-1]] = {0: "batch_size", 2: "past_sequence + sequence"} |
|
dynamic_shapes[inputs[-2]] = {0: "batch_size", 3: "past_sequence + sequence"} |
|
outputs.extend([f"present.{idx}.key", f"present.{idx}.value"]) |
|
|
|
inputs.append("attention_mask") |
|
dummy_inputs = { |
|
"input_ids": torch.ones((1, 2), dtype=torch.long), |
|
"past_key_values": outs.past_key_values, |
|
"attention_mask": torch.ones((1, 12), dtype=torch.long), |
|
} |
|
pt_model.config.torchscript = True |
|
orig_forward = pt_model.forward |
|
|
|
@wraps(orig_forward) |
|
def ts_patched_forward( |
|
input_ids: torch.Tensor, |
|
past_key_values: Tuple[Tuple[torch.Tensor]], |
|
attention_mask: torch.Tensor, |
|
): |
|
pkv_list = list(past_key_values) |
|
outs = orig_forward( |
|
input_ids=input_ids, past_key_values=pkv_list, attention_mask=attention_mask |
|
) |
|
return (outs.logits, tuple(outs.past_key_values)) |
|
|
|
pt_model.forward = ts_patched_forward |
|
ov_model = ov.convert_model(pt_model, example_input=dummy_inputs) |
|
pt_model.forward = orig_forward |
|
for inp_name, m_input, input_data in zip( |
|
inputs, ov_model.inputs, flattenize_inputs(dummy_inputs.values()) |
|
): |
|
input_node = m_input.get_node() |
|
if input_node.element_type == ov.Type.dynamic: |
|
m_input.get_node().set_element_type(ov.Type.f32) |
|
shape = list(input_data.shape) |
|
if inp_name in dynamic_shapes: |
|
for k in dynamic_shapes[inp_name]: |
|
shape[k] = -1 |
|
input_node.set_partial_shape(ov.PartialShape(shape)) |
|
m_input.get_tensor().set_names({inp_name}) |
|
|
|
for out, out_name in zip(ov_model.outputs, outputs): |
|
out.get_tensor().set_names({out_name}) |
|
|
|
ov_model.validate_nodes_and_infer_types() |
|
if make_stateful is not None: |
|
patch_stateful(ov_model, "mpt") |
|
ov.save_model(ov_model, ov_out_path) |
|
del ov_model |
|
cleanup_torchscript_cache() |
|
del pt_model |
|
|
|
|
|
def convert_baichuan(pt_model: torch.nn.Module, model_path: Path): |
|
""" |
|
Baichuan model conversion function |
|
Params: |
|
pt_model: PyTorch model |
|
model_path: path for saving model |
|
Returns: |
|
None |
|
""" |
|
ov_out_path = Path(model_path) / "openvino_model.xml" |
|
pt_model.config.save_pretrained(ov_out_path.parent) |
|
pt_model.config.use_cache = True |
|
outs = pt_model( |
|
input_ids=torch.ones((1, 10), dtype=torch.long), |
|
attention_mask=torch.ones((1, 10), dtype=torch.long), |
|
) |
|
inputs = ["input_ids", "attention_mask"] |
|
outputs = ["logits"] |
|
|
|
dynamic_shapes = { |
|
"input_ids": {0: "batch_size", 1: "seq_len"}, |
|
"attention_mask": {0: "batch_size", 1: "seq_len"}, |
|
} |
|
for idx in range(len(outs.past_key_values)): |
|
inputs.extend([f"past_key_values.{idx}.key", f"past_key_values.{idx}.value"]) |
|
dynamic_shapes[inputs[-1]] = {0: "batch_size", 2: "past_sequence + sequence"} |
|
dynamic_shapes[inputs[-2]] = {0: "batch_size", 2: "past_sequence + sequence"} |
|
outputs.extend([f"present.{idx}.key", f"present.{idx}.value"]) |
|
|
|
dummy_inputs = { |
|
"input_ids": torch.ones((1, 2), dtype=torch.long), |
|
"attention_mask": torch.ones((1, 12), dtype=torch.long), |
|
"past_key_values": outs.past_key_values, |
|
} |
|
pt_model.config.torchscript = True |
|
ov_model = ov.convert_model(pt_model, example_input=dummy_inputs) |
|
for inp_name, m_input, input_data in zip( |
|
inputs, ov_model.inputs, flattenize_inputs(dummy_inputs.values()) |
|
): |
|
input_node = m_input.get_node() |
|
if input_node.element_type == ov.Type.dynamic: |
|
m_input.get_node().set_element_type(ov.Type.f32) |
|
shape = list(input_data.shape) |
|
if inp_name in dynamic_shapes: |
|
for k in dynamic_shapes[inp_name]: |
|
shape[k] = -1 |
|
input_node.set_partial_shape(ov.PartialShape(shape)) |
|
m_input.get_tensor().set_names({inp_name}) |
|
|
|
for out, out_name in zip(ov_model.outputs, outputs): |
|
out.get_tensor().set_names({out_name}) |
|
|
|
ov_model.validate_nodes_and_infer_types() |
|
if make_stateful is not None: |
|
patch_stateful(ov_model, "baichuan") |
|
ov.save_model(ov_model, ov_out_path) |
|
del ov_model |
|
cleanup_torchscript_cache() |
|
del pt_model |
|
|
|
|
|
@torch.jit.script_if_tracing |
|
def _chatglm2_get_context_layer(query_layer: torch.Tensor, key_layer: torch.Tensor, value_layer: torch.Tensor): |
|
mask = torch.zeros((query_layer.shape[-2], key_layer.shape[-2]), dtype=query_layer.dtype) |
|
if query_layer.shape[2] == key_layer.shape[2]: |
|
tmp_mask = torch.ones((query_layer.shape[-2], key_layer.shape[-2]), dtype=torch.bool).triu(diagonal=1) |
|
mask.masked_fill_(tmp_mask, float("-inf")) |
|
|
|
context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer, attn_mask=mask) |
|
return context_layer |
|
|
|
|
|
def _core_attention_forward(self, query_layer, key_layer, value_layer, attention_mask): |
|
query_layer, key_layer, value_layer = [k.permute(1, 2, 0, 3) for k in [query_layer, key_layer, value_layer]] |
|
if attention_mask is None: |
|
context_layer = _chatglm2_get_context_layer(query_layer, key_layer, value_layer) |
|
else: |
|
context_layer = torch.nn.functional.scaled_dot_product_attention( |
|
query_layer, key_layer, value_layer, attention_mask |
|
) |
|
context_layer = context_layer.permute(2, 0, 1, 3) |
|
new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,) |
|
context_layer = context_layer.reshape(*new_context_layer_shape) |
|
|
|
return context_layer |
|
|
|
|
|
@torch.jit.script_if_tracing |
|
def _get_chatglm_attention_mask(input_ids, past_key): |
|
mask = torch.zeros((input_ids.shape[1], past_key.shape[0] + input_ids.shape[1]), dtype=past_key.dtype) |
|
if past_key.shape[0] == 0: |
|
tmp_mask = torch.ones((input_ids.shape[1], past_key.shape[0] + input_ids.shape[1]), dtype=torch.bool).triu(diagonal=1) |
|
mask.masked_fill_(tmp_mask, float("-inf")) |
|
return mask |
|
|
|
|
|
def _chatglm_transformer_forward( |
|
self, |
|
input_ids, |
|
position_ids: Optional[torch.Tensor] = None, |
|
attention_mask: Optional[torch.BoolTensor] = None, |
|
full_attention_mask: Optional[torch.BoolTensor] = None, |
|
past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None, |
|
inputs_embeds: Optional[torch.Tensor] = None, |
|
use_cache: Optional[bool] = None, |
|
output_hidden_states: Optional[bool] = None, |
|
return_dict: Optional[bool] = None |
|
): |
|
output_hidden_states = ( |
|
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states |
|
) |
|
use_cache = use_cache if use_cache is not None else self.config.use_cache |
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
|
|
|
batch_size, seq_length = input_ids.shape |
|
|
|
if inputs_embeds is None: |
|
inputs_embeds = self.embedding(input_ids) |
|
|
|
if self.pre_seq_len is not None: |
|
if past_key_values is None: |
|
past_key_values = self.get_prompt(batch_size=batch_size, device=input_ids.device, |
|
dtype=inputs_embeds.dtype) |
|
if attention_mask is not None: |
|
attention_mask = torch.cat([attention_mask.new_ones((batch_size, self.pre_seq_len)), attention_mask], dim=-1) |
|
|
|
if full_attention_mask is None: |
|
if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1): |
|
full_attention_mask = self.get_masks(input_ids, past_key_values, padding_mask=attention_mask) |
|
elif past_key_values is not None: |
|
full_attention_mask = torch.ones(batch_size, seq_length, seq_length, |
|
device=input_ids.device, |
|
dtype=torch.float) * float("-inf") |
|
full_attention_mask.triu_(diagonal=1) |
|
past_length = 0 |
|
if past_key_values: |
|
past_length = past_key_values[0][0].shape[0] |
|
if past_length: |
|
full_attention_mask = torch.cat((torch.zeros(batch_size, seq_length, past_length, |
|
device=input_ids.device), full_attention_mask), dim=-1) |
|
full_attention_mask.unsqueeze_(1) |
|
|
|
|
|
rotary_pos_emb = self.rotary_pos_emb(self.seq_length) |
|
if position_ids is not None: |
|
rotary_pos_emb = rotary_pos_emb[position_ids] |
|
else: |
|
rotary_pos_emb = rotary_pos_emb[None, :seq_length] |
|
rotary_pos_emb = rotary_pos_emb.transpose(0, 1).contiguous() |
|
|
|
|
|
hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder( |
|
inputs_embeds, full_attention_mask, rotary_pos_emb=rotary_pos_emb, |
|
kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states |
|
) |
|
|
|
if not return_dict: |
|
return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None) |
|
|
|
return BaseModelOutputWithPast( |
|
last_hidden_state=hidden_states, |
|
past_key_values=presents, |
|
hidden_states=all_hidden_states, |
|
attentions=all_self_attentions, |
|
) |
|
|
|
|
|
def _patch_chatglm_forward(model: "PreTrainedModel"): |
|
model.transformer.forward = types.MethodType(_chatglm_transformer_forward, model.transformer) |
|
for block in model.transformer.encoder.layers: |
|
block.self_attention.core_attention.forward = types.MethodType( |
|
_core_attention_forward, block.self_attention.core_attention |
|
) |
|
|
|
|
|
def convert_chatglm(pt_model: torch.nn.Module, model_path: Path): |
|
""" |
|
ChatGLM model conversion function |
|
|
|
Params: |
|
pt_model: PyTorch model |
|
model_path: path for saving model |
|
Returns: |
|
None |
|
""" |
|
_patch_chatglm_forward(pt_model) |
|
ov_out_path = Path(model_path) / "openvino_model.xml" |
|
pt_model.config.save_pretrained(ov_out_path.parent) |
|
pt_model.config.use_cache = True |
|
outs = pt_model( |
|
input_ids=torch.ones((1, 10), dtype=torch.long), |
|
position_ids=torch.arange(0, 10, dtype=torch.long), |
|
) |
|
inputs = ["input_ids"] |
|
outputs = ["logits"] |
|
|
|
dynamic_shapes = { |
|
"input_ids": {0: "batch_size", 1: "seq_len"}, |
|
"position_ids": {0: "batch_size", 1: "seq_len"}, |
|
"attention_mask": {0: "batch_size", 1: "seq_len"}, |
|
} |
|
inputs += ["position_ids", "attention_mask"] |
|
for idx in range(len(outs.past_key_values)): |
|
inputs.extend([f"past_key_values.{idx}.key", f"past_key_values.{idx}.value"]) |
|
dynamic_shapes[inputs[-1]] = {0: "past_sequence + sequence", 1: "batch_size"} |
|
dynamic_shapes[inputs[-2]] = {0: "past_sequence + sequence", 1: "batch_size"} |
|
outputs.extend([f"present.{idx}.key", f"present.{idx}.value"]) |
|
|
|
dummy_inputs = { |
|
"input_ids": torch.ones((1, 1), dtype=torch.long), |
|
"position_ids": torch.tensor([[10]], dtype=torch.long), |
|
"attention_mask": torch.ones((1, 11), dtype=torch.long), |
|
"past_key_values": outs.past_key_values, |
|
} |
|
pt_model.config.torchscript = True |
|
ov_model = ov.convert_model(pt_model, example_input=dummy_inputs) |
|
for inp_name, m_input, input_data in zip( |
|
inputs, ov_model.inputs, flattenize_inputs(dummy_inputs.values()) |
|
): |
|
input_node = m_input.get_node() |
|
if input_node.element_type == ov.Type.dynamic: |
|
m_input.get_node().set_element_type(ov.Type.f32) |
|
shape = list(input_data.shape) |
|
if inp_name in dynamic_shapes: |
|
for k in dynamic_shapes[inp_name]: |
|
shape[k] = -1 |
|
input_node.set_partial_shape(ov.PartialShape(shape)) |
|
m_input.get_tensor().set_names({inp_name}) |
|
|
|
for out, out_name in zip(ov_model.outputs, outputs): |
|
out.get_tensor().set_names({out_name}) |
|
|
|
ov_model.validate_nodes_and_infer_types() |
|
if make_stateful is not None: |
|
patch_stateful(ov_model, "chatglm") |
|
ov.save_model(ov_model, ov_out_path) |
|
del ov_model |
|
cleanup_torchscript_cache() |
|
del pt_model |
|
|
|
def convert_gemma(pt_model: torch.nn.Module, model_path: Path): |
|
""" |
|
Gamma model conversion function |
|
|
|
Params: |
|
pt_model: PyTorch model |
|
model_path: path for saving model |
|
Returns: |
|
None |
|
""" |
|
ov_out_path = Path(model_path) / "openvino_model.xml" |
|
pt_model.config.save_pretrained(ov_out_path.parent) |
|
pt_model.config.use_cache = True |
|
outs = pt_model(input_ids=torch.ones((2, 10), dtype=torch.long)) |
|
inputs = ["input_ids"] |
|
outputs = ["logits"] |
|
|
|
dynamic_shapes = { |
|
"input_ids": {0: "batch_size", 1: "seq_len"}, |
|
"attention_mask": {0: "batch_size", 1: "seq_len"}, |
|
"position_ids": {0: "batch_size", 1: "seq_len"}, |
|
} |
|
inputs += ["attention_mask", "position_ids"] |
|
for idx in range(len(outs.past_key_values)): |
|
inputs.extend([f"past_key_values.{idx}.key", f"past_key_values.{idx}.value"]) |
|
dynamic_shapes[inputs[-1]] = {0: "batch_size", 2: "past_sequence + sequence"} |
|
dynamic_shapes[inputs[-2]] = {0: "batch_size", 2: "past_sequence + sequence"} |
|
outputs.extend([f"present.{idx}.key", f"present.{idx}.value"]) |
|
|
|
dummy_inputs = { |
|
"input_ids": torch.ones((2, 2), dtype=torch.long), |
|
"attention_mask": torch.ones((2, 12), dtype=torch.long), |
|
"position_ids": torch.tensor([[10, 11], [10, 11]], dtype=torch.long), |
|
"past_key_values": outs.past_key_values, |
|
} |
|
pt_model.config.torchscript = True |
|
ov_model = ov.convert_model(pt_model, example_input=dummy_inputs) |
|
for inp_name, m_input, input_data in zip( |
|
inputs, ov_model.inputs, flattenize_inputs(dummy_inputs.values()) |
|
): |
|
input_node = m_input.get_node() |
|
if input_node.element_type == ov.Type.dynamic: |
|
m_input.get_node().set_element_type(ov.Type.f32) |
|
shape = list(input_data.shape) |
|
if inp_name in dynamic_shapes: |
|
for k in dynamic_shapes[inp_name]: |
|
shape[k] = -1 |
|
input_node.set_partial_shape(ov.PartialShape(shape)) |
|
m_input.get_tensor().set_names({inp_name}) |
|
|
|
for out, out_name in zip(ov_model.outputs, outputs): |
|
out.get_tensor().set_names({out_name}) |
|
|
|
ov_model.validate_nodes_and_infer_types() |
|
if make_stateful is not None: |
|
patch_stateful(ov_model, "gemma") |
|
ov.save_model(ov_model, ov_out_path) |
|
del ov_model |
|
cleanup_torchscript_cache() |
|
del pt_model |
|
|
|
|
|
|
|
def convert_mpnet(pt_model: torch.nn.Module, model_path: Path): |
|
ov_out_path = Path(model_path) / "openvino_model.xml" |
|
dummy_inputs = {"input_ids": torch.ones((1, 10), dtype=torch.long), "attention_mask": torch.ones( |
|
(1, 10), dtype=torch.long)} |
|
ov_model = ov.convert_model(pt_model, example_input=dummy_inputs) |
|
ov.save_model(ov_model, ov_out_path) |
|
|
|
def convert_bert(pt_model: torch.nn.Module, model_path: Path): |
|
ov_out_path = Path(model_path) / "openvino_model.xml" |
|
dummy_inputs = {"input_ids": torch.ones((1, 10), dtype=torch.long), "attention_mask": torch.ones( |
|
(1, 10), dtype=torch.long), "token_type_ids": torch.zeros((1, 10), dtype=torch.long)} |
|
ov_model = ov.convert_model(pt_model, example_input=dummy_inputs) |
|
ov.save_model(ov_model, ov_out_path) |
|
|
|
|
|
converters = { |
|
|
|
"mpt": convert_mpt, |
|
"chatglm3": convert_chatglm, |
|
"baichuan2": convert_baichuan, |
|
"gemma": convert_gemma, |
|
|
|
"all-mpnet-base-v2": convert_mpnet, |
|
"text2vec-large-chinese": convert_bert, |
|
} |
|
|
