manaladan6/codegeex4-all-9b
078
1""" PyTorch ChatGLM model. """2import json3import math4import copy5import warnings6import re7import sys8 9import torch10import torch.utils.checkpoint11import torch.nn.functional as F12from torch import nn13from torch.nn import CrossEntropyLoss, LayerNorm, MSELoss, BCEWithLogitsLoss14from torch.nn.utils import skip_init15from typing import Optional, Tuple, Union, List, Callable, Dict, Any16from copy import deepcopy17 18from transformers.modeling_outputs import (19 BaseModelOutputWithPast,20 CausalLMOutputWithPast,21 SequenceClassifierOutputWithPast,22)23from transformers.modeling_utils import PreTrainedModel24from transformers.utils import logging, is_torch_npu_available25from transformers.generation.logits_process import LogitsProcessor26from transformers.generation.utils import LogitsProcessorList, StoppingCriteriaList, GenerationConfig, ModelOutput27 28from .configuration_chatglm import ChatGLMConfig29 30try:31 from transformers.utils import is_flash_attn_greater_or_equal_2_10, is_flash_attn_2_available32 if is_flash_attn_2_available():33 from flash_attn import flash_attn_func, flash_attn_varlen_func34 from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa35except:36 pass37 38 39# flags required to enable jit fusion kernels40 41if sys.platform != 'darwin' and not is_torch_npu_available():42 torch._C._jit_set_profiling_mode(False)43 torch._C._jit_set_profiling_executor(False)44 torch._C._jit_override_can_fuse_on_cpu(True)45 torch._C._jit_override_can_fuse_on_gpu(True)46 47logger = logging.get_logger(__name__)48 49_CHECKPOINT_FOR_DOC = "THUDM/ChatGLM"50_CONFIG_FOR_DOC = "ChatGLMConfig"51 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[..., 198] = 5e462 return scores63 64 65def split_tensor_along_last_dim(66 tensor: torch.Tensor,67 num_partitions: int,68 contiguous_split_chunks: bool = False,69) -> List[torch.Tensor]:70 """Split a tensor along its last dimension.71 72 Arguments:73 tensor: input tensor.74 num_partitions: number of partitions to split the tensor75 contiguous_split_chunks: If True, make each chunk contiguous76 in memory.77 78 Returns:79 A list of Tensors80 """81 # Get the size and dimension.82 last_dim = tensor.dim() - 183 last_dim_size = tensor.size()[last_dim] // num_partitions84 # Split.85 tensor_list = torch.split(tensor, last_dim_size, dim=last_dim)86 # Note: torch.split does not create contiguous tensors by default.87 if contiguous_split_chunks:88 return tuple(chunk.contiguous() for chunk in tensor_list)89 90 return tensor_list91 92 93class RotaryEmbedding(nn.Module):94 def __init__(self, dim, rope_ratio=1, original_impl=False, device=None, dtype=None):95 super().__init__()96 inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, device=device).to(dtype=dtype) / dim))97 self.register_buffer("inv_freq", inv_freq)98 self.dim = dim99 self.original_impl = original_impl100 self.rope_ratio = rope_ratio101 102 def forward_impl(103 self, seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000104 ):105 """Enhanced Transformer with Rotary Position Embedding.106 107 Derived from: https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/108 transformers/rope/__init__.py. MIT License:109 https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license.110 """111 # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$112 base = base * self.rope_ratio113 theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, dtype=torch.float, device=device) / n_elem))114 115 # Create position indexes `[0, 1, ..., seq_len - 1]`116 seq_idx = torch.arange(seq_len, dtype=torch.float, device=device)117 118 # Calculate the product of position index and $\theta_i$119 idx_theta = torch.outer(seq_idx, theta).float()120 121 cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1)122 123 # this is to mimic the behaviour of complex32, else we will get different results124 if dtype in (torch.float16, torch.bfloat16, torch.int8):125 cache = cache.bfloat16() if dtype == torch.bfloat16 else cache.half()126 return cache127 128 def forward(self, max_seq_len, offset=0):129 return self.forward_impl(130 max_seq_len, self.dim, dtype=self.inv_freq.dtype, device=self.inv_freq.device131 )132 133 134@torch.jit.script135def apply_rotary_pos_emb(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor:136 # x: [b, np, sq, hn]137 b, np, sq, hn = x.size(0), x.size(1), x.size(2), x.size(3)138 rot_dim = rope_cache.shape[-2] * 2139 x, x_pass = x[..., :rot_dim], x[..., rot_dim:]140 # truncate to support variable sizes141 rope_cache = rope_cache[:, :sq]142 xshaped = x.reshape(b, np, sq, rot_dim // 2, 2)143 rope_cache = rope_cache.view(-1, 1, sq, xshaped.size(3), 2)144 x_out2 = torch.stack(145 [146 xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1],147 xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1],148 ],149 -1,150 )151 x_out2 = x_out2.flatten(3)152 return torch.cat((x_out2, x_pass), dim=-1)153 154 155class RMSNorm(torch.nn.Module):156 def __init__(self, normalized_shape, eps=1e-5, device=None, dtype=None, **kwargs):157 super().__init__()158 self.weight = torch.nn.Parameter(torch.empty(normalized_shape, device=device, dtype=dtype))159 self.eps = eps160 161 def forward(self, hidden_states: torch.Tensor):162 input_dtype = hidden_states.dtype163 variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)164 hidden_states = hidden_states * torch.rsqrt(variance + self.eps)165 166 return (self.weight * hidden_states).to(input_dtype)167 168 169class CoreAttention(torch.nn.Module):170 def __init__(self, config: ChatGLMConfig, layer_number):171 super(CoreAttention, self).__init__()172 self.config = config173 self.apply_query_key_layer_scaling = config.apply_query_key_layer_scaling174 self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32175 if self.apply_query_key_layer_scaling:176 self.attention_softmax_in_fp32 = True177 self.layer_number = max(1, layer_number)178 self.is_causal = True179 180 projection_size = config.kv_channels * config.num_attention_heads181 182 # Per attention head and per partition values.183 self.hidden_size_per_partition = projection_size184 self.hidden_size_per_attention_head = projection_size // config.num_attention_heads185 self.num_attention_heads_per_partition = config.num_attention_heads186 187 coeff = None188 self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)189 if self.apply_query_key_layer_scaling:190 coeff = self.layer_number191 self.norm_factor *= coeff192 self.coeff = coeff193 194 self.attention_dropout = torch.nn.Dropout(config.attention_dropout)195 196 def forward(self, query_layer, key_layer, value_layer, attention_mask):197 # [b, np, sq, sk]198 output_size = (query_layer.size(0), query_layer.size(1), query_layer.size(2), key_layer.size(2))199 200 # [b, np, sq, hn] -> [b * np, sq, hn]201 query_layer = query_layer.view(output_size[0] * output_size[1], output_size[2], -1)202 # [b, np, sk, hn] -> [b * np, sk, hn]203 key_layer = key_layer.view(output_size[0] * output_size[1], output_size[3], -1)204 205 # preallocting input tensor: [b * np, sq, sk]206 matmul_input_buffer = torch.empty(207 output_size[0] * output_size[1], output_size[2], output_size[3], dtype=query_layer.dtype,208 device=query_layer.device209 )210 211 # Raw attention scores. [b * np, sq, sk]212 matmul_result = torch.baddbmm(213 matmul_input_buffer,214 query_layer, # [b * np, sq, hn]215 key_layer.transpose(1, 2), # [b * np, hn, sk]216 beta=0.0,217 alpha=(1.0 / self.norm_factor),218 )219 220 # change view to [b, np, sq, sk]221 attention_scores = matmul_result.view(*output_size)222 223 # ===========================224 # Attention probs and dropout225 # ===========================226 227 # attention scores and attention mask [b, np, sq, sk]228 if self.attention_softmax_in_fp32:229 attention_scores = attention_scores.float()230 if self.coeff is not None:231 attention_scores = attention_scores * self.coeff232 if attention_mask is None and attention_scores.shape[2] == attention_scores.shape[3]:233 attention_mask = torch.ones(output_size[0], 1, output_size[2], output_size[3],234 device=attention_scores.device, dtype=torch.bool)235 attention_mask.tril_()236 attention_mask = ~attention_mask237 if attention_mask is not None:238 attention_scores = attention_scores.masked_fill(attention_mask, float("-inf"))239 attention_probs = F.softmax(attention_scores, dim=-1)240 attention_probs = attention_probs.type_as(value_layer)241 242 # This is actually dropping out entire tokens to attend to, which might243 # seem a bit unusual, but is taken from the original Transformer paper.244 attention_probs = self.attention_dropout(attention_probs)245 246 # query layer shape: [b * np, sq, hn]247 # value layer shape: [b, np, sk, hn]248 # attention shape: [b, np, sq, sk]249 # context layer shape: [b, np, sq, hn]250 output_size = (value_layer.size(0), value_layer.size(1), query_layer.size(1), value_layer.size(3))251 # change view [b * np, sk, hn]252 value_layer = value_layer.view(output_size[0] * output_size[1], value_layer.size(2), -1)253 # change view [b * np, sq, sk]254 attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)255 # matmul: [b * np, sq, hn]256 context_layer = torch.bmm(attention_probs, value_layer)257 # change view [b, np, sq, hn]258 context_layer = context_layer.view(*output_size)259 # [b, np, sq, hn] --> [b, sq, np, hn]260 context_layer = context_layer.transpose(1, 2).contiguous()261 # [b, sq, np, hn] --> [b, sq, hp]262 new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)263 context_layer = context_layer.reshape(*new_context_layer_shape)264 265 return context_layer266 267 268class SdpaAttention(CoreAttention):269 def forward(self, query_layer, key_layer, value_layer, attention_mask):270 if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]:271 context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,272 is_causal=True,273 dropout_p=self.config.attention_dropout if self.training else 0.0)274 else:275 if attention_mask is not None:276 attention_mask = ~attention_mask277 context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,278 attention_mask,279 dropout_p=self.config.attention_dropout if self.training else 0.0)280 context_layer = context_layer.transpose(1, 2).contiguous()281 new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)282 context_layer = context_layer.reshape(*new_context_layer_shape)283 return context_layer284 285 286def _get_unpad_data(attention_mask):287 seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)288 indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()289 max_seqlen_in_batch = seqlens_in_batch.max().item()290 cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))291 return (292 indices,293 cu_seqlens,294 max_seqlen_in_batch,295 )296 297 298# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2299class FlashAttention2(CoreAttention):300 def __init__(self, *args, **kwargs):301 super().__init__(*args, **kwargs)302 self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()303 304 def forward(self, query_states, key_states, value_states, attention_mask):305 query_states = query_states.transpose(1, 2)306 key_states = key_states.transpose(1, 2)307 value_states = value_states.transpose(1, 2)308 batch_size, query_length = query_states.shape[:2]309 if not self._flash_attn_uses_top_left_mask:310 causal = self.is_causal311 else:312 # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.313 causal = self.is_causal and query_length != 1314 dropout = self.config.attention_dropout if self.training else 0.0315 # Contains at least one padding token in the sequence316 if attention_mask is not None:317 query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(318 query_states, key_states, value_states, attention_mask, query_length319 )320 321 cu_seqlens_q, cu_seqlens_k = cu_seq_lens322 max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens323 324 attn_output_unpad = flash_attn_varlen_func(325 query_states,326 key_states,327 value_states,328 cu_seqlens_q=cu_seqlens_q,329 cu_seqlens_k=cu_seqlens_k,330 max_seqlen_q=max_seqlen_in_batch_q,331 max_seqlen_k=max_seqlen_in_batch_k,332 dropout_p=dropout,333 softmax_scale=None,334 causal=causal,335 )336 337 attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)338 else:339 attn_output = flash_attn_func(340 query_states, key_states, value_states, dropout, softmax_scale=None, causal=causal341 )342 attn_output = attn_output.reshape(batch_size, query_length, self.hidden_size_per_partition).contiguous()343 return attn_output344 345 def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):346 indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)347 batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape348 349 key_layer = index_first_axis(350 key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k351 )352 value_layer = index_first_axis(353 value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k354 )355 if query_length == kv_seq_len:356 query_layer = index_first_axis(357 query_layer.reshape(batch_size * kv_seq_len, self.num_attention_heads_per_partition, head_dim), indices_k358 )359 cu_seqlens_q = cu_seqlens_k360 max_seqlen_in_batch_q = max_seqlen_in_batch_k361 indices_q = indices_k362 elif query_length == 1:363 max_seqlen_in_batch_q = 1364 cu_seqlens_q = torch.arange(365 batch_size + 1, dtype=torch.int32, device=query_layer.device366 ) # There is a memcpy here, that is very bad.367 indices_q = cu_seqlens_q[:-1]368 query_layer = query_layer.squeeze(1)369 else:370 # The -q_len: slice assumes left padding.371 attention_mask = attention_mask[:, -query_length:]372 query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)373 374 return (375 query_layer,376 key_layer,377 value_layer,378 indices_q,379 (cu_seqlens_q, cu_seqlens_k),380 (max_seqlen_in_batch_q, max_seqlen_in_batch_k),381 )382 383 384CORE_ATTENTION_CLASSES = {385 "eager": CoreAttention,386 "sdpa": SdpaAttention,387 "flash_attention_2": FlashAttention2388}389 390 391class SelfAttention(torch.nn.Module):392 """Parallel self-attention layer abstract class.393 394 Self-attention layer takes input with size [s, b, h]395 and returns output of the same size.396 """397 398 def __init__(self, config: ChatGLMConfig, layer_number, device=None):399 super(SelfAttention, self).__init__()400 self.layer_number = max(1, layer_number)401 402 self.projection_size = config.kv_channels * config.num_attention_heads403 404 # Per attention head and per partition values.405 self.hidden_size_per_attention_head = self.projection_size // config.num_attention_heads406 self.num_attention_heads_per_partition = config.num_attention_heads407 408 self.multi_query_attention = config.multi_query_attention409 self.qkv_hidden_size = 3 * self.projection_size410 if self.multi_query_attention:411 self.num_multi_query_groups_per_partition = config.multi_query_group_num412 self.qkv_hidden_size = (413 self.projection_size + 2 * self.hidden_size_per_attention_head * config.multi_query_group_num414 )415 self.query_key_value = nn.Linear(config.hidden_size, self.qkv_hidden_size,416 bias=config.add_bias_linear or config.add_qkv_bias,417 device=device, **_config_to_kwargs(config)418 )419 420 self.core_attention = CORE_ATTENTION_CLASSES[config._attn_implementation](config, self.layer_number)421 422 # Output.423 self.dense = nn.Linear(self.projection_size, config.hidden_size, bias=config.add_bias_linear,424 device=device, **_config_to_kwargs(config)425 )426 427 def _allocate_memory(self, inference_max_sequence_len, batch_size, device=None, dtype=None):428 if self.multi_query_attention:429 num_attention_heads = self.num_multi_query_groups_per_partition430 else:431 num_attention_heads = self.num_attention_heads_per_partition432 return torch.empty(433 inference_max_sequence_len,434 batch_size,435 num_attention_heads,436 self.hidden_size_per_attention_head,437 dtype=dtype,438 device=device,439 )440 441 def forward(442 self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True443 ):444 # hidden_states: [b, sq, h]445 446 # =================================================447 # Pre-allocate memory for key-values for inference.448 # =================================================449 # =====================450 # Query, Key, and Value451 # =====================452 453 # Attention heads [b, sq, h] --> [b, sq, (np * 3 * hn)]454 mixed_x_layer = self.query_key_value(hidden_states)455 456 if self.multi_query_attention:457 (query_layer, key_layer, value_layer) = mixed_x_layer.split(458 [459 self.num_attention_heads_per_partition * self.hidden_size_per_attention_head,460 self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,461 self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,462 ],463 dim=-1,464 )465 query_layer = query_layer.view(466 query_layer.size()[:-1] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)467 )468 key_layer = key_layer.view(469 key_layer.size()[:-1] + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)470 )471 value_layer = value_layer.view(472 value_layer.size()[:-1]473 + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)474 )475 else:476 new_tensor_shape = mixed_x_layer.size()[:-1] + \477 (self.num_attention_heads_per_partition,478 3 * self.hidden_size_per_attention_head)479 mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)480 481 # [b, sq, np, 3 * hn] --> 3 [b, sq, np, hn]482 (query_layer, key_layer, value_layer) = split_tensor_along_last_dim(mixed_x_layer, 3)483 484 # [b, sq, np, hn] -> [b, np, sq, hn]485 query_layer, key_layer, value_layer = [k.transpose(1, 2) for k in [query_layer, key_layer, value_layer]]486 487 # apply relative positional encoding (rotary embedding)488 if rotary_pos_emb is not None:489 query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb)490 key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb)491 492 # adjust key and value for inference493 if kv_cache is not None:494 cache_k, cache_v = kv_cache495 key_layer = torch.cat((cache_k, key_layer), dim=2)496 value_layer = torch.cat((cache_v, value_layer), dim=2)497 if use_cache:498 if kv_cache is None:499 kv_cache = torch.cat((key_layer.unsqueeze(0).unsqueeze(0), value_layer.unsqueeze(0).unsqueeze(0)),500 dim=1)501 else:502 kv_cache = (key_layer, value_layer)503 else:504 kv_cache = None505 506 if self.multi_query_attention:507 key_layer = key_layer.unsqueeze(2)508 key_layer = key_layer.expand(509 -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1, -1510 )511 key_layer = key_layer.contiguous().view(512 key_layer.size()[:1] + (self.num_attention_heads_per_partition,) + key_layer.size()[3:]513 )514 value_layer = value_layer.unsqueeze(2)515 value_layer = value_layer.expand(516 -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1, -1517 )518 value_layer = value_layer.contiguous().view(519 value_layer.size()[:1] + (self.num_attention_heads_per_partition,) + value_layer.size()[3:]520 )521 522 # ==================================523 # core attention computation524 # ==================================525 526 context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask)527 528 # =================529 # Output. [sq, b, h]530 # =================531 532 output = self.dense(context_layer)533 534 return output, kv_cache535 536 537def _config_to_kwargs(args):538 common_kwargs = {539 "dtype": args.torch_dtype,540 }541 return common_kwargs542 543 544class MLP(torch.nn.Module):545 """MLP.546 547 MLP will take the input with h hidden state, project it to 4*h548 hidden dimension, perform nonlinear transformation, and project the549 state back into h hidden dimension.550 """551 552 def __init__(self, config: ChatGLMConfig, device=None):553 super(MLP, self).__init__()554 555 self.add_bias = config.add_bias_linear556 557 # Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf558 self.dense_h_to_4h = nn.Linear(559 config.hidden_size,560 config.ffn_hidden_size * 2,561 bias=self.add_bias,562 device=device,563 **_config_to_kwargs(config)564 )565 566 def swiglu(x):567 x = torch.chunk(x, 2, dim=-1)568 return F.silu(x[0]) * x[1]569 570 self.activation_func = swiglu571 572 # Project back to h.573 self.dense_4h_to_h = nn.Linear(574 config.ffn_hidden_size,575 config.hidden_size,576 bias=self.add_bias,577 device=device,578 **_config_to_kwargs(config)579 )580 581 def forward(self, hidden_states):582 # [s, b, 4hp]583 intermediate_parallel = self.dense_h_to_4h(hidden_states)584 intermediate_parallel = self.activation_func(intermediate_parallel)585 # [s, b, h]586 output = self.dense_4h_to_h(intermediate_parallel)587 return output588 589 590class GLMBlock(torch.nn.Module):591 """A single transformer layer.592 593 Transformer layer takes input with size [s, b, h] and returns an594 output of the same size.595 """596 597 def __init__(self, config: ChatGLMConfig, layer_number, device=None):598 super(GLMBlock, self).__init__()599 self.layer_number = layer_number600 601 self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm602 603 self.fp32_residual_connection = config.fp32_residual_connection604 605 LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm606 # Layernorm on the input data.607 self.input_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,608 dtype=config.torch_dtype)609 610 # Self attention.611 self.self_attention = SelfAttention(config, layer_number, device=device)612 self.hidden_dropout = config.hidden_dropout613 614 # Layernorm on the attention output615 self.post_attention_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,616 dtype=config.torch_dtype)617 618 # MLP619 self.mlp = MLP(config, device=device)620 621 def forward(622 self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True,623 ):624 # hidden_states: [s, b, h]625 626 # Layer norm at the beginning of the transformer layer.627 layernorm_output = self.input_layernorm(hidden_states)628 # Self attention.629 attention_output, kv_cache = self.self_attention(630 layernorm_output,631 attention_mask,632 rotary_pos_emb,633 kv_cache=kv_cache,634 use_cache=use_cache635 )636 637 # Residual connection.638 if self.apply_residual_connection_post_layernorm:639 residual = layernorm_output640 else:641 residual = hidden_states642 643 layernorm_input = torch.nn.functional.dropout(attention_output, p=self.hidden_dropout, training=self.training)644 layernorm_input = residual + layernorm_input645 646 # Layer norm post the self attention.647 layernorm_output = self.post_attention_layernorm(layernorm_input)648 649 # MLP.650 mlp_output = self.mlp(layernorm_output)651 652 # Second residual connection.653 if self.apply_residual_connection_post_layernorm:654 residual = layernorm_output655 else:656 residual = layernorm_input657 658 output = torch.nn.functional.dropout(mlp_output, p=self.hidden_dropout, training=self.training)659 output = residual + output660 661 return output, kv_cache662 663 664class GLMTransformer(torch.nn.Module):665 """Transformer class."""666 667 def __init__(self, config: ChatGLMConfig, device=None):668 super(GLMTransformer, self).__init__()669 670 self.fp32_residual_connection = config.fp32_residual_connection671 self.post_layer_norm = config.post_layer_norm672 673 # Number of layers.674 self.num_layers = config.num_layers675 676 # Transformer layers.677 def build_layer(layer_number):678 return GLMBlock(config, layer_number, device=device)679 680 self.layers = torch.nn.ModuleList([build_layer(i + 1) for i in range(self.num_layers)])681 682 if self.post_layer_norm:683 LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm684 # Final layer norm before output.685 self.final_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,686 dtype=config.torch_dtype)687 688 self.gradient_checkpointing = False689 690 def _get_layer(self, layer_number):691 return self.layers[layer_number]692 693 def forward(694 self, hidden_states, attention_mask, rotary_pos_emb, kv_caches=None,695 use_cache: Optional[bool] = True,696 output_hidden_states: Optional[bool] = False,697 ):698 if not kv_caches:699 kv_caches = [None for _ in range(self.num_layers)]700 presents = () if use_cache else None701 if self.gradient_checkpointing and self.training:702 if use_cache:703 logger.warning_once(704 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."705 )706 use_cache = False707 708 all_self_attentions = None709 all_hidden_states = () if output_hidden_states else None710 for index in range(self.num_layers):711 if output_hidden_states:712 all_hidden_states = all_hidden_states + (hidden_states,)713 714 layer = self._get_layer(index)715 if self.gradient_checkpointing and self.training:716 layer_ret = torch.utils.checkpoint.checkpoint(717 layer,718 hidden_states,719 attention_mask,720 rotary_pos_emb,721 kv_caches[index],722 use_cache,723 use_reentrant=False724 )725 else:726 layer_ret = layer(727 hidden_states,728 attention_mask,729 rotary_pos_emb,730 kv_cache=kv_caches[index],731 use_cache=use_cache732 )733 hidden_states, kv_cache = layer_ret734 if use_cache:735 # token by token decoding, use tuple format736 if kv_caches[0] is not None:737 presents = presents + (kv_cache,)738 # prefilling in decoding, use tensor format to save cuda memory739 else:740 if len(presents) == 0:741 presents = kv_cache742 else:743 presents = torch.cat((presents, kv_cache.to(presents.device)), dim=0)744 745 if output_hidden_states:746 all_hidden_states = all_hidden_states + (hidden_states,)747 748 # Final layer norm.749 if self.post_layer_norm:750 hidden_states = self.final_layernorm(hidden_states)751 752 return hidden_states, presents, all_hidden_states, all_self_attentions753 754 755class ChatGLMPreTrainedModel(PreTrainedModel):756 """757 An abstract class to handle weights initialization and758 a simple interface for downloading and loading pretrained models.759 """760 761 is_parallelizable = False762 supports_gradient_checkpointing = True763 config_class = ChatGLMConfig764 base_model_prefix = "transformer"765 _no_split_modules = ["GLMBlock"]766 _supports_flash_attn_2 = True767 _supports_sdpa = True768 769 def _init_weights(self, module: nn.Module):770 """Initialize the weights."""771 return772 773 def get_masks(self, input_ids, past_key_values, padding_mask=None):774 if self.config._attn_implementation == "flash_attention_2":775 if padding_mask is not None and not padding_mask.all():776 return padding_mask777 return None778 batch_size, seq_length = input_ids.shape779 full_attention_mask = torch.ones(batch_size, seq_length, seq_length, device=input_ids.device)780 full_attention_mask.tril_()781 past_length = 0782 if past_key_values:783 past_length = past_key_values[0][0].shape[2]784 if past_length:785 full_attention_mask = torch.cat((torch.ones(batch_size, seq_length, past_length,786 device=input_ids.device), full_attention_mask), dim=-1)787 if padding_mask is not None:788 full_attention_mask = full_attention_mask * padding_mask.unsqueeze(1)789 if not past_length and padding_mask is not None:790 full_attention_mask -= padding_mask.unsqueeze(-1) - 1791 full_attention_mask = (full_attention_mask < 0.5).bool()792 full_attention_mask.unsqueeze_(1)793 return full_attention_mask794 795 def get_position_ids(self, input_ids, device):796 batch_size, seq_length = input_ids.shape797 position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1)798 return position_ids799 800 def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):801 if not self.supports_gradient_checkpointing:802 raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")803 804 805class Embedding(torch.nn.Module):806 """Language model embeddings."""807 808 def __init__(self, config: ChatGLMConfig, device=None):809 super(Embedding, self).__init__()810 811 self.hidden_size = config.hidden_size812 # Word embeddings (parallel).813 self.word_embeddings = nn.Embedding(814 config.padded_vocab_size,815 self.hidden_size,816 dtype=config.torch_dtype,817 device=device818 )819 self.fp32_residual_connection = config.fp32_residual_connection820 821 def forward(self, input_ids):822 # Embeddings.823 words_embeddings = self.word_embeddings(input_ids)824 embeddings = words_embeddings825 # If the input flag for fp32 residual connection is set, convert for float.826 if self.fp32_residual_connection:827 embeddings = embeddings.float()828 return embeddings829 830 831class ChatGLMModel(ChatGLMPreTrainedModel):832 def __init__(self, config: ChatGLMConfig, device=None, empty_init=True):833 super().__init__(config)834 if empty_init:835 init_method = skip_init836 else:837 init_method = default_init838 init_kwargs = {}839 if device is not None:840 init_kwargs["device"] = device841 self.embedding = init_method(Embedding, config, **init_kwargs)842 self.num_layers = config.num_layers843 self.multi_query_group_num = config.multi_query_group_num844 self.kv_channels = config.kv_channels845 846 # Rotary positional embeddings847 self.seq_length = config.seq_length848 rotary_dim = (849 config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels850 )851 852 self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, rope_ratio=config.rope_ratio,853 original_impl=config.original_rope,854 device=device, dtype=config.torch_dtype)855 self.encoder = init_method(GLMTransformer, config, **init_kwargs)856 self.output_layer = init_method(nn.Linear, config.hidden_size, config.padded_vocab_size, bias=False,857 dtype=config.torch_dtype, **init_kwargs)858 859 def get_input_embeddings(self):860 return self.embedding.word_embeddings861 862 def set_input_embeddings(self, value):863 self.embedding.word_embeddings = value864 865 def forward(866 self,867 input_ids,868 position_ids: Optional[torch.Tensor] = None,869 attention_mask: Optional[torch.BoolTensor] = None,870 full_attention_mask: Optional[torch.BoolTensor] = None,871 past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,872 inputs_embeds: Optional[torch.Tensor] = None,873 use_cache: Optional[bool] = None,874 output_attentions: Optional[bool] = None,875 output_hidden_states: Optional[bool] = None,876 return_dict: Optional[bool] = None,877 ):878 output_hidden_states = (879 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states880 )881 use_cache = use_cache if use_cache is not None else self.config.use_cache882 return_dict = return_dict if return_dict is not None else self.config.use_return_dict883 884 batch_size, seq_length = input_ids.shape885 886 if inputs_embeds is None:887 inputs_embeds = self.embedding(input_ids)888 889 if full_attention_mask is None:890 if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1):891 full_attention_mask = self.get_masks(input_ids, past_key_values, padding_mask=attention_mask)892 893 # Rotary positional embeddings894 rotary_pos_emb = self.rotary_pos_emb(self.seq_length)895 if position_ids is not None:896 rotary_pos_emb = rotary_pos_emb[position_ids]897 else:898 rotary_pos_emb = rotary_pos_emb[None, :seq_length]899 900 # Run encoder.901 hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder(902 inputs_embeds, full_attention_mask, rotary_pos_emb=rotary_pos_emb,903 kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states904 )905 if presents is not None and type(presents) is torch.Tensor:906 presents = presents.split(1, dim=0)907 presents = list(presents)908 presents = [list(x.squeeze(0).split(1, dim=0)) for x in presents]909 presents = [tuple([x.squeeze(0) for x in y]) for y in presents]910 presents = tuple(presents)911 912 if not return_dict:913 return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)914 915 return BaseModelOutputWithPast(916 last_hidden_state=hidden_states,917 past_key_values=presents,918 hidden_states=all_hidden_states,919 attentions=all_self_attentions,920 )921 922 923class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):924 def __init__(self, config: ChatGLMConfig, empty_init=True, device=None):925 super().__init__(config)926 927 self.max_sequence_length = config.max_length928 self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device)929 self.config = config930 931 def _update_model_kwargs_for_generation(932 self,933 outputs: ModelOutput,934 model_kwargs: Dict[str, Any],935 is_encoder_decoder: bool = False,936 standardize_cache_format: bool = False,937 ) -> Dict[str, Any]:938 # update past_key_values939 model_kwargs["past_key_values"] = self._extract_past_from_model_output(940 outputs, standardize_cache_format=standardize_cache_format941 )942 943 # update attention mask944 if "attention_mask" in model_kwargs:945 attention_mask = model_kwargs["attention_mask"]946 model_kwargs["attention_mask"] = torch.cat(947 [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1948 )949 950 # update position ids951 if "position_ids" in model_kwargs:952 position_ids = model_kwargs["position_ids"]953 new_position_id = position_ids[..., -1:].clone()954 new_position_id += 1955 model_kwargs["position_ids"] = torch.cat(956 [position_ids, new_position_id], dim=-1957 )958 959 model_kwargs["is_first_forward"] = False960 return model_kwargs961 962 def prepare_inputs_for_generation(963 self,964 input_ids: torch.LongTensor,965 past_key_values: Optional[torch.Tensor] = None,966 attention_mask: Optional[torch.Tensor] = None,967 position_ids: Optional[torch.Tensor] = None,968 use_cache: Optional[bool] = None,969 is_first_forward: bool = True,970 **kwargs971 ) -> dict:972 # only last token for input_ids if past is not None973 if position_ids is None:974 position_ids = self.get_position_ids(input_ids, device=input_ids.device)975 if not is_first_forward:976 if past_key_values is not None:977 position_ids = position_ids[..., -1:]978 input_ids = input_ids[:, -1:]979 return {980 "input_ids": input_ids,981 "past_key_values": past_key_values,982 "position_ids": position_ids,983 "attention_mask": attention_mask,984 "return_last_logit": True,985 "use_cache": use_cache986 }987 988 def forward(989 self,990 input_ids: Optional[torch.Tensor] = None,991 position_ids: Optional[torch.Tensor] = None,992 attention_mask: Optional[torch.Tensor] = None,993 past_key_values: Optional[Tuple[torch.FloatTensor]] = None,994 inputs_embeds: Optional[torch.Tensor] = None,995 labels: Optional[torch.Tensor] = None,996 use_cache: Optional[bool] = None,997 output_attentions: Optional[bool] = None,998 output_hidden_states: Optional[bool] = None,999 return_dict: Optional[bool] = None,1000 return_last_logit: Optional[bool] = False,1001 ):1002 use_cache = use_cache if use_cache is not None else self.config.use_cache1003 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1004 1005 transformer_outputs = self.transformer(1006 input_ids=input_ids,1007 position_ids=position_ids,1008 attention_mask=attention_mask,1009 past_key_values=past_key_values,1010 inputs_embeds=inputs_embeds,1011 use_cache=use_cache,1012 output_hidden_states=output_hidden_states,1013 return_dict=return_dict,1014 )1015 1016 hidden_states = transformer_outputs[0]1017 if return_last_logit:1018 hidden_states = hidden_states[:, -1:]1019 lm_logits = self.transformer.output_layer(hidden_states)1020 1021 loss = None1022 if labels is not None:1023 lm_logits = lm_logits.to(torch.float32)1024 1025 # Shift so that tokens < n predict n1026 shift_logits = lm_logits[..., :-1, :].contiguous()1027 shift_labels = labels[..., 1:].contiguous()1028 # Flatten the tokens1029 loss_fct = CrossEntropyLoss(ignore_index=-100)1030 loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))1031 1032 lm_logits = lm_logits.to(hidden_states.dtype)1033 loss = loss.to(hidden_states.dtype)1034 1035 if not return_dict:1036 output = (lm_logits,) + transformer_outputs[1:]1037 return ((loss,) + output) if loss is not None else output1038 1039 return CausalLMOutputWithPast(1040 loss=loss,1041 logits=lm_logits,1042 past_key_values=transformer_outputs.past_key_values,1043 hidden_states=transformer_outputs.hidden_states,1044 attentions=transformer_outputs.attentions,1045 )1046 1047 @staticmethod1048 def _reorder_cache(1049 past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor1050 ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:1051 """1052 This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or1053 [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct1054 beam_idx at every generation step.1055 1056 Output shares the same memory storage as `past`.1057 """1058 return tuple(1059 (1060 layer_past[0].index_select(0, beam_idx.to(layer_past[0].device)),1061 layer_past[1].index_select(0, beam_idx.to(layer_past[1].device)),1062 )1063 for layer_past in past1064 )1065 1066 def process_response(self, output, history):1067 content = ""1068 history = deepcopy(history)1069 for response in output.split("<|assistant|>"):1070 if "\n" in response:1071 metadata, content = response.split("\n", maxsplit=1)1072 else:1073 metadata, content = "", response1074 if not metadata.strip():1075 content = content.strip()1076 history.append({"role": "assistant", "metadata": metadata, "content": content})1077 content = content.replace("[[训练时间]]", "2023年")1078 else:1079 history.append({"role": "assistant", "metadata": metadata, "content": content})1080 if history[0]["role"] == "system" and "tools" in history[0]:1081 parameters = json.loads(content)1082 content = {"name": metadata.strip(), "parameters": parameters}1083 else:1084 content = {"name": metadata.strip(), "content": content}1085 return content, history1086 1087 @torch.inference_mode()1088 def chat(self, tokenizer, query: str, history: List[Dict] = None, role: str = "user",1089 max_length: int = 8192, num_beams=1, do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None,1090 **kwargs):1091 if history is None:1092 history = []1093 if logits_processor is None:1094 logits_processor = LogitsProcessorList()1095 logits_processor.append(InvalidScoreLogitsProcessor())1096 gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "do_sample": do_sample, "top_p": top_p,1097 "temperature": temperature, "logits_processor": logits_processor, **kwargs}1098 history.append({"role": role, "content": query})1099 inputs = tokenizer.apply_chat_template(history, add_generation_prompt=True, tokenize=True,1100 return_tensors="pt", return_dict=True)1101 inputs = inputs.to(self.device)1102 eos_token_id = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|user|>"),1103 tokenizer.convert_tokens_to_ids("<|observation|>")]1104 outputs = self.generate(**inputs, **gen_kwargs, eos_token_id=eos_token_id)1105 outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):-1]1106 response = tokenizer.decode(outputs)1107 response, history = self.process_response(response, history)1108 return response, history1109 1110 @torch.inference_mode()1111 def stream_chat(self, tokenizer, query: str, history: List[Dict] = None, role: str = "user",1112 past_key_values=None, max_length: int = 8192, do_sample=True, top_p=0.8, temperature=0.8,1113 logits_processor=None, return_past_key_values=False, **kwargs):1114 if history is None:1115 history = []1116 if logits_processor is None:1117 logits_processor = LogitsProcessorList()1118 logits_processor.append(InvalidScoreLogitsProcessor())1119 eos_token_id = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|user|>"),1120 tokenizer.convert_tokens_to_ids("<|observation|>")]1121 gen_kwargs = {"max_length": max_length, "do_sample": do_sample, "top_p": top_p,1122 "temperature": temperature, "logits_processor": logits_processor, **kwargs}1123 if past_key_values is None:1124 inputs = tokenizer.apply_chat_template(history + [{"role": role, "content": query}],1125 add_generation_prompt=True, tokenize=True, return_tensors="pt",1126 return_dict=True)1127 else:1128 inputs = tokenizer.apply_chat_template([{"role": role, "content": query}], add_special_tokens=False,1129 add_generation_prompt=True, tokenize=True, return_tensors="pt",1130 return_dict=True)1131 inputs = inputs.to(self.device)1132 if past_key_values is not None:1133 past_length = past_key_values[0][0].shape[2]1134 inputs.position_ids += past_length1135 attention_mask = inputs.attention_mask1136 attention_mask = torch.cat((attention_mask.new_ones(1, past_length), attention_mask), dim=1)1137 inputs['attention_mask'] = attention_mask1138 history.append({"role": role, "content": query})1139 for outputs in self.stream_generate(**inputs, past_key_values=past_key_values,1140 eos_token_id=eos_token_id, return_past_key_values=return_past_key_values,1141 **gen_kwargs):1142 if return_past_key_values:1143 outputs, past_key_values = outputs1144 outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):-1]1145 response = tokenizer.decode(outputs)1146 if response and response[-1] != "�":1147 response, new_history = self.process_response(response, history)1148 if return_past_key_values:1149 yield response, new_history, past_key_values1150 else:1151 yield response, new_history1152 1153 @torch.inference_mode()1154 def stream_generate(1155 self,1156 input_ids,1157 generation_config: Optional[GenerationConfig] = None,1158 logits_processor: Optional[LogitsProcessorList] = None,1159 stopping_criteria: Optional[StoppingCriteriaList] = None,1160 prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,1161 return_past_key_values=False,1162 **kwargs,1163 ):1164 batch_size, input_ids_seq_length = input_ids.shape[0], input_ids.shape[-1]1165 1166 if generation_config is None:1167 generation_config = self.generation_config1168 generation_config = copy.deepcopy(generation_config)1169 model_kwargs = generation_config.update(**kwargs)1170 model_kwargs["use_cache"] = generation_config.use_cache1171 bos_token_id, eos_token_id = generation_config.bos_token_id, generation_config.eos_token_id1172 1173 if isinstance(eos_token_id, int):1174 eos_token_id = [eos_token_id]1175 eos_token_id_tensor = torch.tensor(eos_token_id).to(input_ids.device) if eos_token_id is not None else None1176 1177 has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None1178 if has_default_max_length and generation_config.max_new_tokens is None:1179 warnings.warn(1180 f"Using `max_length`'s default ({generation_config.max_length}) to control the generation length. "1181 "This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we"1182 " recommend using `max_new_tokens` to control the maximum length of the generation.",1183 UserWarning,1184 )1185 elif generation_config.max_new_tokens is not None:1186 generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length1187 if not has_default_max_length:1188 logger.warn(1189 f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="1190 f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "1191 "Please refer to the documentation for more information. "1192 "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)",1193 UserWarning,1194 )1195 1196 if input_ids_seq_length >= generation_config.max_length:1197 input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"1198 logger.warning(1199 f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to"1200 f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"