katuni4ka/tiny-random-chatglm2
0195
1""" PyTorch ChatGLM model. """2 3import math4import copy5import warnings6import re7import sys8 9import torch10import torch.utils.checkpoint11import torch.nn.functional as F12from torch import nn13from torch.nn import CrossEntropyLoss, LayerNorm14from torch.nn import CrossEntropyLoss, LayerNorm, MSELoss, BCEWithLogitsLoss15from torch.nn.utils import skip_init16from typing import Optional, Tuple, Union, List, Callable, Dict, Any17import transformers18 19from transformers.modeling_outputs import (20 BaseModelOutputWithPast,21 CausalLMOutputWithPast,22 SequenceClassifierOutputWithPast,23)24from transformers.modeling_utils import PreTrainedModel25from transformers.utils import logging26from transformers.generation.logits_process import LogitsProcessor27from transformers.generation.utils import LogitsProcessorList, StoppingCriteriaList, GenerationConfig, ModelOutput28 29from .configuration_chatglm import ChatGLMConfig30 31# flags required to enable jit fusion kernels32 33if sys.platform != 'darwin':34 torch._C._jit_set_profiling_mode(False)35 torch._C._jit_set_profiling_executor(False)36 torch._C._jit_override_can_fuse_on_cpu(True)37 torch._C._jit_override_can_fuse_on_gpu(True)38 39logger = logging.get_logger(__name__)40 41_CHECKPOINT_FOR_DOC = "THUDM/ChatGLM2-6B"42_CONFIG_FOR_DOC = "ChatGLM6BConfig"43 44CHATGLM_6B_PRETRAINED_MODEL_ARCHIVE_LIST = [45 "THUDM/chatglm2-6b",46 # See all ChatGLM models at https://huggingface.co/models?filter=chatglm47]48 49is_transformers_4_42_or_higher = int(transformers.__version__.split(".")[1]) >= 4250is_transformers_4_44_or_higher = int(transformers.__version__.split(".")[1]) >= 4451 52 53def default_init(cls, *args, **kwargs):54 return cls(*args, **kwargs)55 56 57class InvalidScoreLogitsProcessor(LogitsProcessor):58 def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:59 if torch.isnan(scores).any() or torch.isinf(scores).any():60 scores.zero_()61 scores[..., 5] = 5e462 return scores63 64 65class PrefixEncoder(torch.nn.Module):66 """67 The torch.nn model to encode the prefix68 Input shape: (batch-size, prefix-length)69 Output shape: (batch-size, prefix-length, 2*layers*hidden)70 """71 72 def __init__(self, config: ChatGLMConfig):73 super().__init__()74 self.prefix_projection = config.prefix_projection75 if self.prefix_projection:76 # Use a two-layer MLP to encode the prefix77 kv_size = config.num_layers * config.kv_channels * config.multi_query_group_num * 278 self.embedding = torch.nn.Embedding(config.pre_seq_len, kv_size)79 self.trans = torch.nn.Sequential(80 torch.nn.Linear(kv_size, config.hidden_size),81 torch.nn.Tanh(),82 torch.nn.Linear(config.hidden_size, kv_size)83 )84 else:85 self.embedding = torch.nn.Embedding(config.pre_seq_len,86 config.num_layers * config.kv_channels * config.multi_query_group_num * 2)87 88 def forward(self, prefix: torch.Tensor):89 if self.prefix_projection:90 prefix_tokens = self.embedding(prefix)91 past_key_values = self.trans(prefix_tokens)92 else:93 past_key_values = self.embedding(prefix)94 return past_key_values95 96 97def split_tensor_along_last_dim(98 tensor: torch.Tensor,99 num_partitions: int,100 contiguous_split_chunks: bool = False,101) -> List[torch.Tensor]:102 """Split a tensor along its last dimension.103 104 Arguments:105 tensor: input tensor.106 num_partitions: number of partitions to split the tensor107 contiguous_split_chunks: If True, make each chunk contiguous108 in memory.109 110 Returns:111 A list of Tensors112 """113 # Get the size and dimension.114 last_dim = tensor.dim() - 1115 last_dim_size = tensor.size()[last_dim] // num_partitions116 # Split.117 tensor_list = torch.split(tensor, last_dim_size, dim=last_dim)118 # Note: torch.split does not create contiguous tensors by default.119 if contiguous_split_chunks:120 return tuple(chunk.contiguous() for chunk in tensor_list)121 122 return tensor_list123 124 125class RotaryEmbedding(nn.Module):126 def __init__(self, dim, original_impl=False, device=None, dtype=None):127 super().__init__()128 inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, device=device).to(dtype=dtype) / dim))129 self.register_buffer("inv_freq", inv_freq)130 self.dim = dim131 self.original_impl = original_impl132 133 def forward_impl(134 self, seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000135 ):136 """Enhanced Transformer with Rotary Position Embedding.137 138 Derived from: https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/139 transformers/rope/__init__.py. MIT License:140 https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license.141 """142 # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$143 theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, dtype=dtype, device=device) / n_elem))144 145 # Create position indexes `[0, 1, ..., seq_len - 1]`146 seq_idx = torch.arange(seq_len, dtype=dtype, device=device)147 148 # Calculate the product of position index and $\theta_i$149 idx_theta = torch.outer(seq_idx, theta).float()150 151 cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1)152 153 # this is to mimic the behaviour of complex32, else we will get different results154 if dtype in (torch.float16, torch.bfloat16, torch.int8):155 cache = cache.bfloat16() if dtype == torch.bfloat16 else cache.half()156 return cache157 158 def forward(self, max_seq_len, offset=0):159 return self.forward_impl(160 max_seq_len, self.dim, dtype=self.inv_freq.dtype, device=self.inv_freq.device161 )162 163 164@torch.jit.script165def apply_rotary_pos_emb(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor:166 # x: [sq, b, np, hn]167 sq, b, np, hn = x.size(0), x.size(1), x.size(2), x.size(3)168 rot_dim = rope_cache.shape[-2] * 2169 x, x_pass = x[..., :rot_dim], x[..., rot_dim:]170 # truncate to support variable sizes171 rope_cache = rope_cache[:sq]172 xshaped = x.reshape(sq, -1, np, rot_dim // 2, 2)173 rope_cache = rope_cache.view(sq, -1, 1, xshaped.size(3), 2)174 x_out2 = torch.stack(175 [176 xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1],177 xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1],178 ],179 -1,180 )181 x_out2 = x_out2.flatten(3)182 return torch.cat((x_out2, x_pass), dim=-1)183 184 185class RMSNorm(torch.nn.Module):186 def __init__(self, normalized_shape, eps=1e-5, device=None, dtype=None, **kwargs):187 super().__init__()188 self.weight = torch.nn.Parameter(torch.empty(normalized_shape, device=device, dtype=dtype))189 self.eps = eps190 191 def forward(self, hidden_states: torch.Tensor):192 input_dtype = hidden_states.dtype193 variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)194 hidden_states = hidden_states * torch.rsqrt(variance + self.eps)195 196 return (self.weight * hidden_states).to(input_dtype)197 198 199class CoreAttention(torch.nn.Module):200 def __init__(self, config: ChatGLMConfig, layer_number):201 super(CoreAttention, self).__init__()202 203 self.apply_query_key_layer_scaling = config.apply_query_key_layer_scaling204 self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32205 if self.apply_query_key_layer_scaling:206 self.attention_softmax_in_fp32 = True207 self.layer_number = max(1, layer_number)208 209 projection_size = config.kv_channels * config.num_attention_heads210 211 # Per attention head and per partition values.212 self.hidden_size_per_partition = projection_size213 self.hidden_size_per_attention_head = projection_size // config.num_attention_heads214 self.num_attention_heads_per_partition = config.num_attention_heads215 216 coeff = None217 self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)218 if self.apply_query_key_layer_scaling:219 coeff = self.layer_number220 self.norm_factor *= coeff221 self.coeff = coeff222 223 self.attention_dropout = torch.nn.Dropout(config.attention_dropout)224 225 def forward(self, query_layer, key_layer, value_layer, attention_mask):226 pytorch_major_version = int(torch.__version__.split('.')[0])227 if pytorch_major_version >= 2:228 query_layer, key_layer, value_layer = [k.permute(1, 2, 0, 3) for k in [query_layer, key_layer, value_layer]]229 if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]:230 context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,231 is_causal=True)232 else:233 if attention_mask is not None:234 attention_mask = ~attention_mask235 context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,236 attention_mask)237 context_layer = context_layer.permute(2, 0, 1, 3)238 new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)239 context_layer = context_layer.reshape(*new_context_layer_shape)240 else:241 # Raw attention scores242 243 # [b, np, sq, sk]244 output_size = (query_layer.size(1), query_layer.size(2), query_layer.size(0), key_layer.size(0))245 246 # [sq, b, np, hn] -> [sq, b * np, hn]247 query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1)248 # [sk, b, np, hn] -> [sk, b * np, hn]249 key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1)250 251 # preallocting input tensor: [b * np, sq, sk]252 matmul_input_buffer = torch.empty(253 output_size[0] * output_size[1], output_size[2], output_size[3], dtype=query_layer.dtype,254 device=query_layer.device255 )256 257 # Raw attention scores. [b * np, sq, sk]258 matmul_result = torch.baddbmm(259 matmul_input_buffer,260 query_layer.transpose(0, 1), # [b * np, sq, hn]261 key_layer.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk]262 beta=0.0,263 alpha=(1.0 / self.norm_factor),264 )265 266 # change view to [b, np, sq, sk]267 attention_scores = matmul_result.view(*output_size)268 269 # ===========================270 # Attention probs and dropout271 # ===========================272 273 # attention scores and attention mask [b, np, sq, sk]274 if self.attention_softmax_in_fp32:275 attention_scores = attention_scores.float()276 if self.coeff is not None:277 attention_scores = attention_scores * self.coeff278 if attention_mask is None and attention_scores.shape[2] == attention_scores.shape[3]:279 attention_mask = torch.ones(output_size[0], 1, output_size[2], output_size[3],280 device=attention_scores.device, dtype=torch.bool)281 attention_mask.tril_()282 attention_mask = ~attention_mask283 if attention_mask is not None:284 attention_scores = attention_scores.masked_fill(attention_mask, float("-inf"))285 attention_probs = F.softmax(attention_scores, dim=-1)286 attention_probs = attention_probs.type_as(value_layer)287 288 # This is actually dropping out entire tokens to attend to, which might289 # seem a bit unusual, but is taken from the original Transformer paper.290 attention_probs = self.attention_dropout(attention_probs)291 # =========================292 # Context layer. [sq, b, hp]293 # =========================294 295 # value_layer -> context layer.296 # [sk, b, np, hn] --> [b, np, sq, hn]297 298 # context layer shape: [b, np, sq, hn]299 output_size = (value_layer.size(1), value_layer.size(2), query_layer.size(0), value_layer.size(3))300 # change view [sk, b * np, hn]301 value_layer = value_layer.view(value_layer.size(0), output_size[0] * output_size[1], -1)302 # change view [b * np, sq, sk]303 attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)304 # matmul: [b * np, sq, hn]305 context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1))306 # change view [b, np, sq, hn]307 context_layer = context_layer.view(*output_size)308 # [b, np, sq, hn] --> [sq, b, np, hn]309 context_layer = context_layer.permute(2, 0, 1, 3).contiguous()310 # [sq, b, np, hn] --> [sq, b, hp]311 new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)312 context_layer = context_layer.view(*new_context_layer_shape)313 314 return context_layer315 316 317class SelfAttention(torch.nn.Module):318 """Parallel self-attention layer abstract class.319 320 Self-attention layer takes input with size [s, b, h]321 and returns output of the same size.322 """323 324 def __init__(self, config: ChatGLMConfig, layer_number, device=None):325 super(SelfAttention, self).__init__()326 self.layer_number = max(1, layer_number)327 328 self.projection_size = config.kv_channels * config.num_attention_heads329 330 # Per attention head and per partition values.331 self.hidden_size_per_attention_head = self.projection_size // config.num_attention_heads332 self.num_attention_heads_per_partition = config.num_attention_heads333 334 self.multi_query_attention = config.multi_query_attention335 self.qkv_hidden_size = 3 * self.projection_size336 if self.multi_query_attention:337 self.num_multi_query_groups_per_partition = config.multi_query_group_num338 self.qkv_hidden_size = (339 self.projection_size + 2 * self.hidden_size_per_attention_head * config.multi_query_group_num340 )341 self.query_key_value = nn.Linear(config.hidden_size, self.qkv_hidden_size,342 bias=config.add_bias_linear or config.add_qkv_bias,343 device=device, **_config_to_kwargs(config)344 )345 346 self.core_attention = CoreAttention(config, self.layer_number)347 348 # Output.349 self.dense = nn.Linear(self.projection_size, config.hidden_size, bias=config.add_bias_linear,350 device=device, **_config_to_kwargs(config)351 )352 353 def _allocate_memory(self, inference_max_sequence_len, batch_size, device=None, dtype=None):354 if self.multi_query_attention:355 num_attention_heads = self.num_multi_query_groups_per_partition356 else:357 num_attention_heads = self.num_attention_heads_per_partition358 return torch.empty(359 inference_max_sequence_len,360 batch_size,361 num_attention_heads,362 self.hidden_size_per_attention_head,363 dtype=dtype,364 device=device,365 )366 367 def forward(368 self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True369 ):370 # hidden_states: [sq, b, h]371 372 # =================================================373 # Pre-allocate memory for key-values for inference.374 # =================================================375 # =====================376 # Query, Key, and Value377 # =====================378 379 # Attention heads [sq, b, h] --> [sq, b, (np * 3 * hn)]380 mixed_x_layer = self.query_key_value(hidden_states)381 382 if self.multi_query_attention:383 (query_layer, key_layer, value_layer) = mixed_x_layer.split(384 [385 self.num_attention_heads_per_partition * self.hidden_size_per_attention_head,386 self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,387 self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,388 ],389 dim=-1,390 )391 query_layer = query_layer.view(392 query_layer.size()[:-1] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)393 )394 key_layer = key_layer.view(395 key_layer.size()[:-1] + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)396 )397 value_layer = value_layer.view(398 value_layer.size()[:-1]399 + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)400 )401 else:402 new_tensor_shape = mixed_x_layer.size()[:-1] + \403 (self.num_attention_heads_per_partition,404 3 * self.hidden_size_per_attention_head)405 mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)406 407 # [sq, b, np, 3 * hn] --> 3 [sq, b, np, hn]408 (query_layer, key_layer, value_layer) = split_tensor_along_last_dim(mixed_x_layer, 3)409 410 # apply relative positional encoding (rotary embedding)411 if rotary_pos_emb is not None:412 query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb)413 key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb)414 415 # adjust key and value for inference416 if kv_cache is not None:417 cache_k, cache_v = kv_cache418 key_layer = torch.cat((cache_k, key_layer), dim=0)419 value_layer = torch.cat((cache_v, value_layer), dim=0)420 if use_cache:421 kv_cache = (key_layer, value_layer)422 else:423 kv_cache = None424 425 if self.multi_query_attention:426 key_layer = key_layer.unsqueeze(-2)427 key_layer = key_layer.expand(428 -1, -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1429 )430 key_layer = key_layer.contiguous().view(431 key_layer.size()[:2] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)432 )433 value_layer = value_layer.unsqueeze(-2)434 value_layer = value_layer.expand(435 -1, -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1436 )437 value_layer = value_layer.contiguous().view(438 value_layer.size()[:2] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)439 )440 441 # ==================================442 # core attention computation443 # ==================================444 445 context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask)446 447 # =================448 # Output. [sq, b, h]449 # =================450 451 output = self.dense(context_layer)452 453 return output, kv_cache454 455 456def _config_to_kwargs(args):457 common_kwargs = {458 "dtype": args.torch_dtype if not isinstance(args.torch_dtype, str) else getattr(torch, args.torch_dtype)459 }460 return common_kwargs461 462 463class MLP(torch.nn.Module):464 """MLP.465 466 MLP will take the input with h hidden state, project it to 4*h467 hidden dimension, perform nonlinear transformation, and project the468 state back into h hidden dimension.469 """470 471 def __init__(self, config: ChatGLMConfig, device=None):472 super(MLP, self).__init__()473 474 self.add_bias = config.add_bias_linear475 476 # Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf477 self.dense_h_to_4h = nn.Linear(478 config.hidden_size,479 config.ffn_hidden_size * 2,480 bias=self.add_bias,481 device=device,482 **_config_to_kwargs(config)483 )484 485 def swiglu(x):486 x = torch.chunk(x, 2, dim=-1)487 return F.silu(x[0]) * x[1]488 489 self.activation_func = swiglu490 491 # Project back to h.492 self.dense_4h_to_h = nn.Linear(493 config.ffn_hidden_size,494 config.hidden_size,495 bias=self.add_bias,496 device=device,497 **_config_to_kwargs(config)498 )499 500 def forward(self, hidden_states):501 # [s, b, 4hp]502 intermediate_parallel = self.dense_h_to_4h(hidden_states)503 intermediate_parallel = self.activation_func(intermediate_parallel)504 # [s, b, h]505 output = self.dense_4h_to_h(intermediate_parallel)506 return output507 508 509class GLMBlock(torch.nn.Module):510 """A single transformer layer.511 512 Transformer layer takes input with size [s, b, h] and returns an513 output of the same size.514 """515 516 def __init__(self, config: ChatGLMConfig, layer_number, device=None):517 super(GLMBlock, self).__init__()518 self.layer_number = layer_number519 dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype520 521 self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm522 523 self.fp32_residual_connection = config.fp32_residual_connection524 525 LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm526 # Layernorm on the input data.527 self.input_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,528 dtype=dtype)529 530 # Self attention.531 self.self_attention = SelfAttention(config, layer_number, device=device)532 self.hidden_dropout = config.hidden_dropout533 534 # Layernorm on the attention output535 self.post_attention_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,536 dtype=dtype)537 538 # MLP539 self.mlp = MLP(config, device=device)540 541 def forward(542 self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True,543 ):544 # hidden_states: [s, b, h]545 546 # Layer norm at the beginning of the transformer layer.547 layernorm_output = self.input_layernorm(hidden_states)548 # Self attention.549 attention_output, kv_cache = self.self_attention(550 layernorm_output,551 attention_mask,552 rotary_pos_emb,553 kv_cache=kv_cache,554 use_cache=use_cache555 )556 557 # Residual connection.558 if self.apply_residual_connection_post_layernorm:559 residual = layernorm_output560 else:561 residual = hidden_states562 563 layernorm_input = torch.nn.functional.dropout(attention_output, p=self.hidden_dropout, training=self.training)564 layernorm_input = residual + layernorm_input565 566 # Layer norm post the self attention.567 layernorm_output = self.post_attention_layernorm(layernorm_input)568 569 # MLP.570 mlp_output = self.mlp(layernorm_output)571 572 # Second residual connection.573 if self.apply_residual_connection_post_layernorm:574 residual = layernorm_output575 else:576 residual = layernorm_input577 578 output = torch.nn.functional.dropout(mlp_output, p=self.hidden_dropout, training=self.training)579 output = residual + output580 581 return output, kv_cache582 583 584class GLMTransformer(torch.nn.Module):585 """Transformer class."""586 587 def __init__(self, config: ChatGLMConfig, device=None):588 super(GLMTransformer, self).__init__()589 590 self.fp32_residual_connection = config.fp32_residual_connection591 self.post_layer_norm = config.post_layer_norm592 593 # Number of layers.594 self.num_layers = config.num_layers595 596 # Transformer layers.597 def build_layer(layer_number):598 return GLMBlock(config, layer_number, device=device)599 600 self.layers = torch.nn.ModuleList([build_layer(i + 1) for i in range(self.num_layers)])601 602 if self.post_layer_norm:603 LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm604 dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype605 # Final layer norm before output.606 self.final_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,607 dtype=dtype)608 609 self.gradient_checkpointing = False610 611 def _get_layer(self, layer_number):612 return self.layers[layer_number]613 614 def forward(615 self, hidden_states, attention_mask, rotary_pos_emb, kv_caches=None,616 use_cache: Optional[bool] = True,617 output_hidden_states: Optional[bool] = False,618 ):619 if not kv_caches:620 kv_caches = [None for _ in range(self.num_layers)]621 presents = () if use_cache else None622 if self.gradient_checkpointing and self.training:623 if use_cache:624 logger.warning_once(625 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."626 )627 use_cache = False628 629 all_self_attentions = None630 all_hidden_states = () if output_hidden_states else None631 for index in range(self.num_layers):632 if output_hidden_states:633 all_hidden_states = all_hidden_states + (hidden_states,)634 635 layer = self._get_layer(index)636 if self.gradient_checkpointing and self.training:637 layer_ret = torch.utils.checkpoint.checkpoint(638 layer,639 hidden_states,640 attention_mask,641 rotary_pos_emb,642 kv_caches[index],643 use_cache644 )645 else:646 layer_ret = layer(647 hidden_states,648 attention_mask,649 rotary_pos_emb,650 kv_cache=kv_caches[index],651 use_cache=use_cache652 )653 hidden_states, kv_cache = layer_ret654 if use_cache:655 presents = presents + (kv_cache,)656 657 if output_hidden_states:658 all_hidden_states = all_hidden_states + (hidden_states,)659 660 # Final layer norm.661 if self.post_layer_norm:662 hidden_states = self.final_layernorm(hidden_states)663 664 return hidden_states, presents, all_hidden_states, all_self_attentions665 666 667class ChatGLMPreTrainedModel(PreTrainedModel):668 """669 An abstract class to handle weights initialization and670 a simple interface for downloading and loading pretrained models.671 """672 673 is_parallelizable = False674 supports_gradient_checkpointing = True675 config_class = ChatGLMConfig676 base_model_prefix = "transformer"677 _no_split_modules = ["GLMBlock"]678 679 def _init_weights(self, module: nn.Module):680 """Initialize the weights."""681 return682 683 def get_masks(self, input_ids, past_key_values, padding_mask=None):684 batch_size, seq_length = input_ids.shape685 full_attention_mask = torch.ones(batch_size, seq_length, seq_length, device=input_ids.device)686 full_attention_mask.tril_()687 past_length = 0688 if past_key_values:689 past_length = past_key_values[0][0].shape[0]690 if past_length:691 full_attention_mask = torch.cat((torch.ones(batch_size, seq_length, past_length,692 device=input_ids.device), full_attention_mask), dim=-1)693 if padding_mask is not None:694 full_attention_mask = full_attention_mask * padding_mask.unsqueeze(1)695 if not past_length and padding_mask is not None:696 full_attention_mask -= padding_mask.unsqueeze(-1) - 1697 full_attention_mask = (full_attention_mask < 0.5).bool()698 full_attention_mask.unsqueeze_(1)699 return full_attention_mask700 701 def get_position_ids(self, input_ids, device):702 batch_size, seq_length = input_ids.shape703 position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1)704 return position_ids705 706 def _set_gradient_checkpointing(self, module, value=False):707 if isinstance(module, GLMTransformer):708 module.gradient_checkpointing = value709 710 711class Embedding(torch.nn.Module):712 """Language model embeddings."""713 714 def __init__(self, config: ChatGLMConfig, device=None):715 super(Embedding, self).__init__()716 dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype717 718 self.hidden_size = config.hidden_size719 # Word embeddings (parallel).720 self.word_embeddings = nn.Embedding(721 config.padded_vocab_size,722 self.hidden_size,723 dtype=dtype,724 device=device725 )726 self.fp32_residual_connection = config.fp32_residual_connection727 728 def forward(self, input_ids):729 # Embeddings.730 words_embeddings = self.word_embeddings(input_ids)731 embeddings = words_embeddings732 # Data format change to avoid explicit tranposes : [b s h] --> [s b h].733 embeddings = embeddings.transpose(0, 1).contiguous()734 # If the input flag for fp32 residual connection is set, convert for float.735 if self.fp32_residual_connection:736 embeddings = embeddings.float()737 return embeddings738 739 740class ChatGLMModel(ChatGLMPreTrainedModel):741 def __init__(self, config: ChatGLMConfig, device=None, empty_init=False):742 super().__init__(config)743 if empty_init:744 init_method = skip_init745 else:746 init_method = default_init747 init_kwargs = {}748 if device is not None:749 init_kwargs["device"] = device if not isinstance(device, str) else torch.device(device)750 self.embedding = init_method(Embedding, config, **init_kwargs)751 self.num_layers = config.num_layers752 self.multi_query_group_num = config.multi_query_group_num753 self.kv_channels = config.kv_channels754 dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype755 756 # Rotary positional embeddings757 self.seq_length = config.seq_length758 rotary_dim = (759 config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels760 )761 762 self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, original_impl=config.original_rope, device=device,763 dtype=dtype)764 self.encoder = init_method(GLMTransformer, config, **init_kwargs)765 self.output_layer = init_method(nn.Linear, config.hidden_size, config.padded_vocab_size, bias=False,766 dtype=dtype, **init_kwargs)767 self.pre_seq_len = config.pre_seq_len768 self.prefix_projection = config.prefix_projection769 if self.pre_seq_len is not None:770 for param in self.parameters():771 param.requires_grad = False772 self.prefix_tokens = torch.arange(self.pre_seq_len).long()773 self.prefix_encoder = PrefixEncoder(config)774 self.dropout = torch.nn.Dropout(0.1)775 776 def get_input_embeddings(self):777 return self.embedding.word_embeddings778 779 def get_prompt(self, batch_size, device, dtype=torch.half):780 prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(device)781 past_key_values = self.prefix_encoder(prefix_tokens).type(dtype)782 past_key_values = past_key_values.view(783 batch_size,784 self.pre_seq_len,785 self.num_layers * 2,786 self.multi_query_group_num,787 self.kv_channels788 )789 # seq_len, b, nh, hidden_size790 past_key_values = self.dropout(past_key_values)791 past_key_values = past_key_values.permute([2, 1, 0, 3, 4]).split(2)792 return past_key_values793 794 def forward(795 self,796 input_ids,797 position_ids: Optional[torch.Tensor] = None,798 attention_mask: Optional[torch.BoolTensor] = None,799 full_attention_mask: Optional[torch.BoolTensor] = None,800 past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,801 inputs_embeds: Optional[torch.Tensor] = None,802 use_cache: Optional[bool] = None,803 output_hidden_states: Optional[bool] = None,804 return_dict: Optional[bool] = None,805 ):806 output_hidden_states = (807 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states808 )809 use_cache = use_cache if use_cache is not None else self.config.use_cache810 return_dict = return_dict if return_dict is not None else self.config.use_return_dict811 812 batch_size, seq_length = input_ids.shape813 814 if inputs_embeds is None:815 inputs_embeds = self.embedding(input_ids)816 817 if self.pre_seq_len is not None:818 if past_key_values is None:819 past_key_values = self.get_prompt(batch_size=batch_size, device=input_ids.device,820 dtype=inputs_embeds.dtype)821 if attention_mask is not None:822 attention_mask = torch.cat([attention_mask.new_ones((batch_size, self.pre_seq_len)),823 attention_mask], dim=-1)824 825 if full_attention_mask is None:826 if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1):827 full_attention_mask = self.get_masks(input_ids, past_key_values, padding_mask=attention_mask)828 829 # Rotary positional embeddings830 rotary_pos_emb = self.rotary_pos_emb(self.seq_length)831 if position_ids is not None:832 rotary_pos_emb = rotary_pos_emb[position_ids]833 else:834 rotary_pos_emb = rotary_pos_emb[None, :seq_length]835 rotary_pos_emb = rotary_pos_emb.transpose(0, 1).contiguous()836 837 # Run encoder.838 hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder(839 inputs_embeds, full_attention_mask, rotary_pos_emb=rotary_pos_emb,840 kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states841 )842 843 if not return_dict:844 return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)845 846 return BaseModelOutputWithPast(847 last_hidden_state=hidden_states,848 past_key_values=presents,849 hidden_states=all_hidden_states,850 attentions=all_self_attentions,851 )852 853 def quantize(self, weight_bit_width: int):854 from .quantization import quantize855 quantize(self.encoder, weight_bit_width)856 return self857 858 859class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):860 def __init__(self, config: ChatGLMConfig, empty_init=True, device=None):861 super().__init__(config)862 863 self.max_sequence_length = config.max_length864 self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device)865 self.config = config866 self.quantized = False867 868 if self.config.quantization_bit:869 self.quantize(self.config.quantization_bit, empty_init=True)870 871 872 @staticmethod873 def _extract_past_from_model_output(outputs: ModelOutput, *args, **kwargs):874 past_key_values = None875 if "past_key_values" in outputs:876 past_key_values = outputs.past_key_values877 if is_transformers_4_42_or_higher:878 return None, past_key_values879 return past_key_values880 881 def _update_model_kwargs_for_generation(882 self,883 outputs: ModelOutput,884 model_kwargs: Dict[str, Any],885 is_encoder_decoder: bool = False,886 standardize_cache_format: bool = False,887 ) -> Dict[str, Any]:888 if is_transformers_4_44_or_higher:889 model_kwargs["past_key_values"] = self._extract_past_from_model_output(890 outputs891 )[1]892 elif is_transformers_4_42_or_higher:893 # update past_key_values894 model_kwargs["past_key_values"] = self._extract_past_from_model_output(895 outputs, standardize_cache_format=standardize_cache_format896 )[1]897 else:898 model_kwargs["past_key_values"] = self._extract_past_from_model_output(899 outputs, standardize_cache_format=standardize_cache_format900 )901 902 903 # update attention mask904 if "attention_mask" in model_kwargs:905 attention_mask = model_kwargs["attention_mask"]906 model_kwargs["attention_mask"] = torch.cat(907 [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1908 )909 910 # update position ids911 if "position_ids" in model_kwargs:912 position_ids = model_kwargs["position_ids"]913 new_position_id = position_ids[..., -1:].clone()914 new_position_id += 1915 model_kwargs["position_ids"] = torch.cat(916 [position_ids, new_position_id], dim=-1917 )918 919 model_kwargs["is_first_forward"] = False920 return model_kwargs921 922 def prepare_inputs_for_generation(923 self,924 input_ids: torch.LongTensor,925 past_key_values: Optional[torch.Tensor] = None,926 attention_mask: Optional[torch.Tensor] = None,927 position_ids: Optional[torch.Tensor] = None,928 use_cache: Optional[bool] = None,929 is_first_forward: bool = True,930 **kwargs931 ) -> dict:932 # only last token for input_ids if past is not None933 if position_ids is None:934 position_ids = self.get_position_ids(input_ids, device=input_ids.device)935 if not is_first_forward:936 if past_key_values is not None:937 position_ids = position_ids[..., -1:]938 input_ids = input_ids[:, -1:]939 return {940 "input_ids": input_ids,941 "past_key_values": past_key_values,942 "position_ids": position_ids,943 "attention_mask": attention_mask,944 "return_last_logit": True,945 "use_cache": use_cache946 }947 948 def forward(949 self,950 input_ids: Optional[torch.Tensor] = None,951 position_ids: Optional[torch.Tensor] = None,952 attention_mask: Optional[torch.Tensor] = None,953 past_key_values: Optional[Tuple[torch.FloatTensor]] = None,954 inputs_embeds: Optional[torch.Tensor] = None,955 labels: Optional[torch.Tensor] = None,956 use_cache: Optional[bool] = None,957 output_attentions: Optional[bool] = None,958 output_hidden_states: Optional[bool] = None,959 return_dict: Optional[bool] = None,960 return_last_logit: Optional[bool] = False,961 ):962 use_cache = use_cache if use_cache is not None else self.config.use_cache963 return_dict = return_dict if return_dict is not None else self.config.use_return_dict964 965 transformer_outputs = self.transformer(966 input_ids=input_ids,967 position_ids=position_ids,968 attention_mask=attention_mask,969 past_key_values=past_key_values,970 inputs_embeds=inputs_embeds,971 use_cache=use_cache,972 output_hidden_states=output_hidden_states,973 return_dict=return_dict,974 )975 976 hidden_states = transformer_outputs[0]977 if return_last_logit:978 hidden_states = hidden_states[-1:]979 lm_logits = self.transformer.output_layer(hidden_states)980 lm_logits = lm_logits.transpose(0, 1).contiguous()981 982 loss = None983 if labels is not None:984 lm_logits = lm_logits.to(torch.float32)985 986 # Shift so that tokens < n predict n987 shift_logits = lm_logits[..., :-1, :].contiguous()988 shift_labels = labels[..., 1:].contiguous()989 # Flatten the tokens990 loss_fct = CrossEntropyLoss(ignore_index=-100)991 loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))992 993 lm_logits = lm_logits.to(hidden_states.dtype)994 loss = loss.to(hidden_states.dtype)995 996 if not return_dict:997 output = (lm_logits,) + transformer_outputs[1:]998 return ((loss,) + output) if loss is not None else output999 1000 return CausalLMOutputWithPast(1001 loss=loss,1002 logits=lm_logits,1003 past_key_values=transformer_outputs.past_key_values,1004 hidden_states=transformer_outputs.hidden_states,1005 attentions=transformer_outputs.attentions,1006 )1007 1008 @staticmethod1009 def _reorder_cache(1010 past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor1011 ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:1012 """1013 This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or1014 [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct1015 beam_idx at every generation step.1016 1017 Output shares the same memory storage as `past`.1018 """1019 return tuple(1020 (1021 layer_past[0].index_select(1, beam_idx.to(layer_past[0].device)),1022 layer_past[1].index_select(1, beam_idx.to(layer_past[1].device)),1023 )1024 for layer_past in past1025 )1026 1027 def process_response(self, response):1028 response = response.strip()1029 response = response.replace("[[训练时间]]", "2023年")1030 return response1031 1032 def build_inputs(self, tokenizer, query: str, history: List[Tuple[str, str]] = None):1033 prompt = tokenizer.build_prompt(query, history=history)1034 inputs = tokenizer([prompt], return_tensors="pt")1035 inputs = inputs.to(self.device)1036 return inputs1037 1038 def build_stream_inputs(self, tokenizer, query: str, history: List[Tuple[str, str]] = None):1039 if history:1040 prompt = "\n\n[Round {}]\n\n问：{}\n\n答：".format(len(history) + 1, query)1041 input_ids = tokenizer.encode(prompt, add_special_tokens=False)1042 input_ids = input_ids[1:]1043 inputs = tokenizer.batch_encode_plus([(input_ids, None)], return_tensors="pt", add_special_tokens=False)1044 else:1045 prompt = "[Round {}]\n\n问：{}\n\n答：".format(len(history) + 1, query)1046 inputs = tokenizer([prompt], return_tensors="pt")1047 inputs = inputs.to(self.device)1048 return inputs1049 1050 @torch.inference_mode()1051 def chat(self, tokenizer, query: str, history: List[Tuple[str, str]] = None, max_length: int = 8192, num_beams=1,1052 do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None, **kwargs):1053 if history is None:1054 history = []1055 if logits_processor is None:1056 logits_processor = LogitsProcessorList()1057 logits_processor.append(InvalidScoreLogitsProcessor())1058 gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "do_sample": do_sample, "top_p": top_p,1059 "temperature": temperature, "logits_processor": logits_processor, **kwargs}1060 inputs = self.build_inputs(tokenizer, query, history=history)1061 outputs = self.generate(**inputs, **gen_kwargs)1062 outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):]1063 response = tokenizer.decode(outputs)1064 response = self.process_response(response)1065 history = history + [(query, response)]1066 return response, history1067 1068 @torch.inference_mode()1069 def stream_chat(self, tokenizer, query: str, history: List[Tuple[str, str]] = None, past_key_values=None,1070 max_length: int = 8192, do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None,1071 return_past_key_values=False, **kwargs):1072 if history is None:1073 history = []1074 if logits_processor is None:1075 logits_processor = LogitsProcessorList()1076 logits_processor.append(InvalidScoreLogitsProcessor())1077 gen_kwargs = {"max_length": max_length, "do_sample": do_sample, "top_p": top_p,1078 "temperature": temperature, "logits_processor": logits_processor, **kwargs}1079 if past_key_values is None and not return_past_key_values:1080 inputs = self.build_inputs(tokenizer, query, history=history)1081 else:1082 inputs = self.build_stream_inputs(tokenizer, query, history=history)1083 if past_key_values is not None:1084 past_length = past_key_values[0][0].shape[0]1085 if self.transformer.pre_seq_len is not None:1086 past_length -= self.transformer.pre_seq_len1087 inputs.position_ids += past_length1088 attention_mask = inputs.attention_mask1089 attention_mask = torch.cat((attention_mask.new_ones(1, past_length), attention_mask), dim=1)1090 inputs['attention_mask'] = attention_mask1091 for outputs in self.stream_generate(**inputs, past_key_values=past_key_values,1092 return_past_key_values=return_past_key_values, **gen_kwargs):1093 if return_past_key_values:1094 outputs, past_key_values = outputs1095 outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):]1096 response = tokenizer.decode(outputs)1097 if response and response[-1] != "�":1098 response = self.process_response(response)1099 new_history = history + [(query, response)]1100 if return_past_key_values:1101 yield response, new_history, past_key_values1102 else:1103 yield response, new_history1104 1105 @torch.inference_mode()1106 def stream_generate(1107 self,1108 input_ids,1109 generation_config: Optional[GenerationConfig] = None,1110 logits_processor: Optional[LogitsProcessorList] = None,1111 stopping_criteria: Optional[StoppingCriteriaList] = None,1112 prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,1113 return_past_key_values=False,1114 **kwargs,1115 ):1116 batch_size, input_ids_seq_length = input_ids.shape[0], input_ids.shape[-1]1117 1118 if generation_config is None:1119 generation_config = self.generation_config1120 generation_config = copy.deepcopy(generation_config)1121 model_kwargs = generation_config.update(**kwargs)1122 model_kwargs["use_cache"] = generation_config.use_cache1123 bos_token_id, eos_token_id = generation_config.bos_token_id, generation_config.eos_token_id1124 1125 if isinstance(eos_token_id, int):1126 eos_token_id = [eos_token_id]1127 1128 has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None1129 if has_default_max_length and generation_config.max_new_tokens is None:1130 warnings.warn(1131 f"Using `max_length`'s default ({generation_config.max_length}) to control the generation length. "1132 "This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we"1133 " recommend using `max_new_tokens` to control the maximum length of the generation.",1134 UserWarning,1135 )1136 elif generation_config.max_new_tokens is not None:1137 generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length1138 if not has_default_max_length:1139 logger.warn(1140 f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="1141 f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "1142 "Please refer to the documentation for more information. "1143 "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)",1144 UserWarning,1145 )1146 1147 if input_ids_seq_length >= generation_config.max_length:1148 input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"1149 logger.warning(1150 f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to"1151 f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"1152 " increasing `max_new_tokens`."1153 )1154 1155 # 2. Set generation parameters if not already defined1156 logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()1157 stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()1158 1159 logits_processor = self._get_logits_processor(1160 generation_config=generation_config,1161 input_ids_seq_length=input_ids_seq_length,1162 encoder_input_ids=input_ids,1163 prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,1164 logits_processor=logits_processor,1165 )1166 1167 stopping_criteria = self._get_stopping_criteria(1168 generation_config=generation_config, stopping_criteria=stopping_criteria1169 )1170 logits_warper = self._get_logits_warper(generation_config)1171 1172 unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1)1173 scores = None1174 while True:1175 model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)1176 # forward pass to get next token1177 outputs = self(1178 **model_inputs,1179 return_dict=True,1180 output_attentions=False,1181 output_hidden_states=False,1182 )1183 1184 next_token_logits = outputs.logits[:, -1, :]1185 1186 # pre-process distribution1187 next_token_scores = logits_processor(input_ids, next_token_logits)1188 next_token_scores = logits_warper(input_ids, next_token_scores)1189 1190 # sample1191 probs = nn.functional.softmax(next_token_scores, dim=-1)1192 if generation_config.do_sample:1193 next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)1194 else:1195 next_tokens = torch.argmax(probs, dim=-1)1196 1197 # update generated ids, model inputs, and length for next step1198 input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)1199 model_kwargs = self._update_model_kwargs_for_generation(1200 outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder