""" ChatGLM model configuration """
import torch
from collections import OrderedDict
from typing import List, Mapping, Optional, Any
from transformers.configuration_utils import PretrainedConfig
from transformers.utils import logging
from transformers.onnx import OnnxConfigWithPast, PatchingSpec
from transformers import PreTrainedTokenizer, TensorType, is_torch_available
logger = logging.get_logger(__name__)
class ChatGLMConfig(PretrainedConfig):
r"""
This is the configuration class to store the configuration of a [`~ChatGLMModel`].
It is used to instantiate an ChatGLM model according to the specified arguments, defining the model
architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of
the ChatGLM-6B [THUDM/ChatGLM-6B](https://huggingface.co/THUDM/chatglm-6b) architecture.
Configuration objects inherit from [`PretrainedConfig`] and can be used
to control the model outputs. Read the documentation from [`PretrainedConfig`]
for more information.
Args:
vocab_size (`int`, *optional*, defaults to 150528):
Vocabulary size of the ChatGLM-6B model. Defines the number of different tokens that can be represented by the
`inputs_ids` passed when calling [`~ChatGLMModel`] or
[`~TFChatGLMModel`].
hidden_size (`int`, *optional*, defaults to 4096):
Dimension of the encoder layers and the pooler layer.
num_hidden_layers (`int`, *optional*, defaults to 28):
Number of hidden layers in the Transformer encoder.
num_attention_heads (`int`, *optional*, defaults to 32):
Number of attention heads for each attention layer in the Transformer encoder.
inner_hidden_size (`int`, *optional*, defaults to 16384):
Dimension of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
max_sequence_length (`int`, *optional*, defaults to 512):
The maximum sequence length that this model might ever be used with.
Typically set this to something large just in case (e.g., 512 or 1024 or 2048).
layernorm_epsilon (`float`, *optional*, defaults to 1e-5):
The epsilon used by the layer normalization layers.
use_cache (`bool`, *optional*, defaults to `True`):
Whether the model should return the last key/values attentions (not used by all models).
Example:
```python
>>> from configuration_chatglm import ChatGLMConfig
>>> from modeling_chatglm import ChatGLMModel
>>> # Initializing a ChatGLM-6B THUDM/ChatGLM-6B style configuration
>>> configuration = ChatGLMConfig()
>>> # Initializing a model from the THUDM/ChatGLM-6B style configuration
>>> model = ChatGLMModel(configuration)
>>> # Accessing the model configuration
>>> configuration = model.config
```
"""
model_type = "chatglm"
def __init__(
self,
vocab_size=150528,
hidden_size=4096,
num_layers=28,
num_attention_heads=32,
layernorm_epsilon=1e-5,
use_cache=False,
bos_token_id=150004,
eos_token_id=150005,
mask_token_id=150000,
gmask_token_id=150001,
pad_token_id=0,
max_sequence_length=2048,
inner_hidden_size=16384,
position_encoding_2d=True,
quantization_bit=0,
pre_seq_len=None,
prefix_projection=False,
**kwargs
):
self.num_layers = num_layers
self.vocab_size = vocab_size
self.hidden_size = hidden_size
self.num_attention_heads = num_attention_heads
self.max_sequence_length = max_sequence_length
self.layernorm_epsilon = layernorm_epsilon
self.inner_hidden_size = inner_hidden_size
self.use_cache = use_cache
self.bos_token_id = bos_token_id
self.eos_token_id = eos_token_id
self.pad_token_id = pad_token_id
self.mask_token_id = mask_token_id
self.gmask_token_id = gmask_token_id
self.position_encoding_2d = position_encoding_2d
self.quantization_bit = quantization_bit
self.pre_seq_len = pre_seq_len
self.prefix_projection = prefix_projection
super().__init__(
pad_token_id=pad_token_id,
bos_token_id=bos_token_id,
eos_token_id=eos_token_id,
**kwargs
)
class ChatGLMOnnxConfig(OnnxConfigWithPast):
r"""
This class is the custom configuration of a ChatGLMModel needed in exporting model to ONNX.
Currently this need to pre-fix several model struct in modeling_chatglm.py
Also there is still a TODO list of current ChatGLMOnnxConfig:
1. add support for batch_size > 1
2. add support for use_past
in modeling_chatglm.py and its attention_fn function,we need to change several view into
torch tensor action since reshape param may get frozen into constant in onnx model.
here is the code:
```python
>>> def attention_fn(
>>> self,
>>> query_layer,
>>> key_layer,
>>> value_layer,
>>> attention_mask,
>>> hidden_size_per_partition,
>>> layer_id,
>>> layer_past=None,
>>> scaling_attention_score=True,
>>> use_cache=False,
>>> ):
>>> if layer_past is not None:
>>> past_key, past_value = layer_past[0], layer_past[1]
>>> key_layer = torch.cat((past_key, key_layer), dim=0)
>>> value_layer = torch.cat((past_value, value_layer), dim=0)
>>>
>>> # seqlen, batch, num_attention_heads, hidden_size_per_attention_head
>>> seq_len, b, nh, hidden_size = key_layer.shape
>>>
>>> if use_cache:
>>> present = (key_layer, value_layer)
>>> else:
>>> present = None
>>>
>>> query_key_layer_scaling_coeff = float(layer_id + 1)
>>> if scaling_attention_score:
>>> query_layer = query_layer / (math.sqrt(hidden_size) * query_key_layer_scaling_coeff)
>>>
>>> # ===================================
>>> # Raw attention scores. [b, np, s, s]
>>> # ===================================
>>>
>>> # [b, np, sq, sk]
>>> # # output_size = (query_layer.size(1), query_layer.size(2), query_layer.size(0), key_layer.size(0))
>>>
>>> # [sq, b, np, hn] -> [sq, b * np, hn]
>>> # query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1)
>>> query_layer = query_layer.flatten(start_dim=1, end_dim=2)
>>>
>>> # [sk, b, np, hn] -> [sk, b * np, hn]
>>> # key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1)
>>> key_layer = key_layer.flatten(start_dim=1, end_dim=2)
>>>
>>> matmul_result = torch.zeros(
>>> 1, 1, 1,
>>> dtype=query_layer.dtype,
>>> device=query_layer.device,
>>> )
>>>
>>> matmul_result = torch.baddbmm(
>>> matmul_result,
>>> query_layer.transpose(0, 1), # [b * np, sq, hn]
>>> key_layer.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk]
>>> beta=0.0,
>>> alpha=1.0,
>>> )
>>>
>>> # [b * np, sq, sk] -> [b, np, sq, sk]
>>> # attention_scores = matmul_result.view(*output_size)
>>> attention_scores = matmul_result.unsqueeze(0)
>>>
>>> if self.scale_mask_softmax:
>>> self.scale_mask_softmax.scale = query_key_layer_scaling_coeff
>>> attention_probs = self.scale_mask_softmax(attention_scores, attention_mask.contiguous())
>>> else:
>>> # if not (attention_mask == 0).all():
>>> # # if auto-regressive, skip
>>> attention_scores.masked_fill_(attention_mask, -10000.0)
>>> dtype = attention_scores.dtype
>>> attention_scores = attention_scores.float()
>>> attention_scores = attention_scores * query_key_layer_scaling_coeff
>>>
>>> attention_probs = F.softmax(attention_scores, dim=-1)
>>>
>>> attention_probs = attention_probs.type(dtype)
>>>
>>> # =========================
>>> # Context layer. [sq, b, hp]
>>> # =========================
>>>
>>> # value_layer -> context layer.
>>> # [sk, b, np, hn] --> [b, np, sq, hn]
>>>
>>> # context layer shape: [b, np, sq, hn]
>>> # output_size = (value_layer.size(1), value_layer.size(2), query_layer.size(0), value_layer.size(3))
>>>
>>> # change view [sk, b * np, hn]
>>> # value_layer = value_layer.view(value_layer.size(0), output_size[0] * output_size[1], -1)
>>> value_layer = value_layer.flatten(start_dim=1, end_dim=2)
>>>
>>> # change view [b * np, sq, sk]
>>> # attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
>>> attention_probs = attention_probs.flatten(start_dim=0, end_dim=1)
>>>
>>> # matmul: [b * np, sq, hn]
>>> context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1))
>>>
>>> # change view [b, np, sq, hn]
>>> # context_layer = context_layer.reshape(b, np, sq, hidden_size)
>>> context_layer = context_layer.unsqueeze(0)
>>>
>>> # [b, np, sq, hn] --> [sq, b, np, hn]
>>> context_layer = context_layer.permute(2, 0, 1, 3).contiguous()
>>>
>>> # [sq, b, np, hn] --> [sq, b, hp]
>>> # new_context_layer_shape = context_layer.size()[:-2] + (hidden_size_per_partition,)
>>> # context_layer = context_layer.view(*new_context_layer_shape)
>>> context_layer = context_layer.flatten(start_dim=2)
>>>
>>> outputs = (context_layer, present, attention_probs)
>>>
>>> return outputs
'''
mainly aviod using view with dynamic size
after change the modeling_chatglm.py, you can simply use following code to export and test the onnx model
```python
>>> from pathlib import Path
>>> from transformers import AutoTokenizer, AutoModel
>>> from transformers.onnx import export, validate_model_outputs
>>>
>>> # load model
>>> tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
>>> pt_model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
>>> pt_model = pt_model.float() # only tested in CPU for now
>>> pt_model.eval()
>>> # define path for saving onnx model
>>> onnx_path = Path(f"model/chatglm-6b.onnx")
>>> onnx_path.parent.mkdir(exist_ok=True)
>>> # convert model to onnx
>>> onnx_config_chatglm = ChatGLMOnnxConfig(pt_model.config, task="causal-lm")
>>> onnx_inputs, onnx_outputs = export(tokenizer, pt_model,
>>> onnx_config_chatglm, onnx_config_chatglm.default_onnx_opset,
>>> onnx_path)
>>> # test onnx model
>>> validate_model_outputs(onnx_config_chatglm, tokenizer, pt_model, onnx_path, onnx_outputs, atol=1e-4)
```
"""
# TODO support dynamic batch size
default_fixed_batch = 1
def __init__(
self,
config: PretrainedConfig,
task: str = "default",
patching_specs: List[PatchingSpec] = None,
use_past: bool = False,
):
super().__init__(config, task=task, patching_specs=patching_specs, use_past=use_past)
@property
def inputs(self) -> Mapping[str, Mapping[int, str]]:
common_inputs = OrderedDict({"input_ids": {0: "batch", 1: "sequence"}})
if self.use_past:
# TODO support use_past
# self.fill_with_past_key_values_(common_inputs, direction="inputs")
# common_inputs["attention_mask"] = \
# {0: "batch", 1: "past_sequence + sequence", 2: "past_sequence + sequence"}
raise NotImplementedError('position_ids do not support past_key_values yet.')
else:
# remind the order
common_inputs["position_ids"] = {0: "batch", 2: "sequence"}
common_inputs["attention_mask"] = {0: "batch", 2: "sequence", 3: "sequence"}
return common_inputs
@property
def num_layers(self) -> int:
return self._config.n_layer
@property
def num_attention_heads(self) -> int:
return self._config.n_head
def get_masks(self, input_ids, device=None):
"""
reference from modeling_chatglm.get_masks
"""
batch_size, seq_length = input_ids.shape
context_lengths = [seq.tolist().index(self._config.bos_token_id) for seq in input_ids]
if device:
attention_mask = torch.ones((batch_size, seq_length, seq_length), device=device)
else:
attention_mask = torch.ones((batch_size, seq_length, seq_length), device=input_ids.device)
attention_mask.tril_()
for i, context_length in enumerate(context_lengths):
attention_mask[i, :, :context_length] = 1
attention_mask.unsqueeze_(1)
attention_mask = (attention_mask < 0.5).bool()
# print("attention_mask", attention_mask.shape)
return attention_mask
def get_position_ids(self, input_ids, mask_positions, device=None, use_gmasks=None):
batch_size, seq_length = input_ids.shape
if device is None:
device = input_ids.device
if use_gmasks is None:
use_gmasks = [False] * batch_size
context_lengths = [seq.tolist().index(self._config.bos_token_id) for seq in input_ids]
if self._config.position_encoding_2d:
position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1)
for i, context_length in enumerate(context_lengths):
position_ids[i, context_length:] = mask_positions[i]
block_position_ids = [torch.cat((
torch.zeros(context_length, dtype=torch.long, device=device),
torch.arange(seq_length - context_length, dtype=torch.long, device=device) + 1
)) for context_length in context_lengths]
block_position_ids = torch.stack(block_position_ids, dim=0)
position_ids = torch.stack((position_ids, block_position_ids), dim=1)
else:
position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1)
for i, context_length in enumerate(context_lengths):
if not use_gmasks[i]:
position_ids[context_length:] = mask_positions[i]
# print("position_ids", position_ids.shape)
return position_ids
def generate_dummy_inputs(
self,
tokenizer: PreTrainedTokenizer,
batch_size: int = default_fixed_batch,
seq_length: int = -1,
is_pair: bool = False,
framework: Optional[TensorType] = None,
) -> Mapping[str, Any]:
common_inputs = super(OnnxConfigWithPast, self).generate_dummy_inputs(
tokenizer, batch_size=self.default_fixed_batch, seq_length=seq_length, is_pair=is_pair, framework=framework
)
# check if the mode is using fixed batch size
if batch_size != self.default_fixed_batch:
logger.warning('batch size is not fixed, force change into fixed batch size: %d.'
% self.default_fixed_batch)
# We need to order the input in the way they appears in the forward()
ordered_inputs = OrderedDict({"input_ids": common_inputs["input_ids"]})
# Need to add the past_keys
if self.use_past:
if not is_torch_available():
raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed.")
else:
# TODO support use_past
# import torch
#
# batch, seqlen = common_inputs["input_ids"].shape
# # Not using the same length for past_key_values
# past_key_values_length = seqlen + 2
# past_shape = (
# batch,
# self.num_attention_heads,
# past_key_values_length,
# self._config.hidden_size // self.num_attention_heads,
# )
# ordered_inputs["past_key_values"] = [
# (torch.zeros(past_shape), torch.zeros(past_shape)) for _ in range(self.num_layers)
# ]
raise NotImplementedError('position_ids do not support past_key_values yet.')
# Need to add the attention_mask manually
# 1. add attention_mask
ordered_inputs["attention_mask"] = self.get_masks(common_inputs["input_ids"])
# 2. add position_ids
MASK, gMASK = self._config.mask_token_id, self._config.gmask_token_id
seqs = common_inputs["input_ids"].tolist()
mask_positions, use_gmasks = [], []
for seq in seqs:
mask_token = gMASK if gMASK in seq else MASK
use_gmask = mask_token == gMASK
mask_positions.append(seq.index(mask_token))
use_gmasks.append(use_gmask)
ordered_inputs["position_ids"] = self.get_position_ids(common_inputs["input_ids"],
mask_positions, use_gmasks=use_gmasks)
if self.use_past:
# mask_dtype = ordered_inputs["attention_mask"].dtype
# ordered_inputs["attention_mask"] = torch.cat(
# [ordered_inputs["attention_mask"], torch.ones(batch, past_key_values_length, dtype=mask_dtype)], dim=1
# )
raise NotImplementedError('position_ids do not support past_key_values yet.')
return ordered_inputs
@property
def default_onnx_opset(self) -> int:
return 13