zjzju commited on
Commit
bfbbecc
·
1 Parent(s): 2d445d4

Upload 76 files

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. checkpoint-1000/config.json +45 -0
  2. checkpoint-1000/configuration_chatglm.py +59 -0
  3. checkpoint-1000/finetuning_args.json +13 -0
  4. checkpoint-1000/generation_config.json +6 -0
  5. checkpoint-1000/modeling_chatglm.py +1193 -0
  6. checkpoint-1000/optimizer.pt +3 -0
  7. checkpoint-1000/pytorch_model.bin +3 -0
  8. checkpoint-1000/quantization.py +188 -0
  9. checkpoint-1000/rng_state_0.pth +3 -0
  10. checkpoint-1000/rng_state_1.pth +3 -0
  11. checkpoint-1000/rng_state_2.pth +3 -0
  12. checkpoint-1000/rng_state_3.pth +3 -0
  13. checkpoint-1000/scheduler.pt +3 -0
  14. checkpoint-1000/special_tokens_map.json +1 -0
  15. checkpoint-1000/tokenization_chatglm.py +253 -0
  16. checkpoint-1000/tokenizer.model +3 -0
  17. checkpoint-1000/tokenizer_config.json +14 -0
  18. checkpoint-1000/trainer_state.json +616 -0
  19. checkpoint-1000/training_args.bin +3 -0
  20. checkpoint-2000/config.json +45 -0
  21. checkpoint-2000/configuration_chatglm.py +59 -0
  22. checkpoint-2000/finetuning_args.json +13 -0
  23. checkpoint-2000/generation_config.json +6 -0
  24. checkpoint-2000/modeling_chatglm.py +1193 -0
  25. checkpoint-2000/optimizer.pt +3 -0
  26. checkpoint-2000/pytorch_model.bin +3 -0
  27. checkpoint-2000/quantization.py +188 -0
  28. checkpoint-2000/rng_state_0.pth +3 -0
  29. checkpoint-2000/rng_state_1.pth +3 -0
  30. checkpoint-2000/rng_state_2.pth +3 -0
  31. checkpoint-2000/rng_state_3.pth +3 -0
  32. checkpoint-2000/scheduler.pt +3 -0
  33. checkpoint-2000/special_tokens_map.json +1 -0
  34. checkpoint-2000/tokenization_chatglm.py +253 -0
  35. checkpoint-2000/tokenizer.model +3 -0
  36. checkpoint-2000/tokenizer_config.json +14 -0
  37. checkpoint-2000/trainer_state.json +1216 -0
  38. checkpoint-2000/training_args.bin +3 -0
  39. checkpoint-3000/config.json +45 -0
  40. checkpoint-3000/configuration_chatglm.py +59 -0
  41. checkpoint-3000/finetuning_args.json +13 -0
  42. checkpoint-3000/generation_config.json +6 -0
  43. checkpoint-3000/modeling_chatglm.py +1193 -0
  44. checkpoint-3000/optimizer.pt +3 -0
  45. checkpoint-3000/pytorch_model.bin +3 -0
  46. checkpoint-3000/quantization.py +188 -0
  47. checkpoint-3000/rng_state_0.pth +3 -0
  48. checkpoint-3000/rng_state_1.pth +3 -0
  49. checkpoint-3000/rng_state_2.pth +3 -0
  50. checkpoint-3000/rng_state_3.pth +3 -0
checkpoint-1000/config.json ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "/home/ubuntu/chatglm2-6b",
3
+ "add_bias_linear": false,
4
+ "add_qkv_bias": true,
5
+ "apply_query_key_layer_scaling": true,
6
+ "apply_residual_connection_post_layernorm": false,
7
+ "architectures": [
8
+ "ChatGLMForConditionalGeneration"
9
+ ],
10
+ "attention_dropout": 0.0,
11
+ "attention_softmax_in_fp32": true,
12
+ "auto_map": {
13
+ "AutoConfig": "configuration_chatglm.ChatGLMConfig",
14
+ "AutoModel": "modeling_chatglm.ChatGLMForConditionalGeneration",
15
+ "AutoModelForCausalLM": "modeling_chatglm.ChatGLMForConditionalGeneration",
16
+ "AutoModelForSeq2SeqLM": "modeling_chatglm.ChatGLMForConditionalGeneration"
17
+ },
18
+ "bias_dropout_fusion": true,
19
+ "eos_token_id": 2,
20
+ "ffn_hidden_size": 13696,
21
+ "fp32_residual_connection": false,
22
+ "hidden_dropout": 0.0,
23
+ "hidden_size": 4096,
24
+ "kv_channels": 128,
25
+ "layernorm_epsilon": 1e-05,
26
+ "model_type": "chatglm",
27
+ "multi_query_attention": true,
28
+ "multi_query_group_num": 2,
29
+ "num_attention_heads": 32,
30
+ "num_layers": 28,
31
+ "original_rope": true,
32
+ "pad_token_id": 0,
33
+ "padded_vocab_size": 65024,
34
+ "post_layer_norm": true,
35
+ "pre_seq_len": null,
36
+ "prefix_projection": false,
37
+ "quantization_bit": 0,
38
+ "rmsnorm": true,
39
+ "seq_length": 32768,
40
+ "tie_word_embeddings": false,
41
+ "torch_dtype": "float16",
42
+ "transformers_version": "4.30.2",
43
+ "use_cache": true,
44
+ "vocab_size": 65024
45
+ }
checkpoint-1000/configuration_chatglm.py ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import PretrainedConfig
2
+
3
+
4
+ class ChatGLMConfig(PretrainedConfig):
5
+ model_type = "chatglm"
6
+ def __init__(
7
+ self,
8
+ num_layers=28,
9
+ padded_vocab_size=65024,
10
+ hidden_size=4096,
11
+ ffn_hidden_size=13696,
12
+ kv_channels=128,
13
+ num_attention_heads=32,
14
+ seq_length=2048,
15
+ hidden_dropout=0.0,
16
+ attention_dropout=0.0,
17
+ layernorm_epsilon=1e-5,
18
+ rmsnorm=True,
19
+ apply_residual_connection_post_layernorm=False,
20
+ post_layer_norm=True,
21
+ add_bias_linear=False,
22
+ add_qkv_bias=False,
23
+ bias_dropout_fusion=True,
24
+ multi_query_attention=False,
25
+ multi_query_group_num=1,
26
+ apply_query_key_layer_scaling=True,
27
+ attention_softmax_in_fp32=True,
28
+ fp32_residual_connection=False,
29
+ quantization_bit=0,
30
+ pre_seq_len=None,
31
+ prefix_projection=False,
32
+ **kwargs
33
+ ):
34
+ self.num_layers = num_layers
35
+ self.vocab_size = padded_vocab_size
36
+ self.padded_vocab_size = padded_vocab_size
37
+ self.hidden_size = hidden_size
38
+ self.ffn_hidden_size = ffn_hidden_size
39
+ self.kv_channels = kv_channels
40
+ self.num_attention_heads = num_attention_heads
41
+ self.seq_length = seq_length
42
+ self.hidden_dropout = hidden_dropout
43
+ self.attention_dropout = attention_dropout
44
+ self.layernorm_epsilon = layernorm_epsilon
45
+ self.rmsnorm = rmsnorm
46
+ self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm
47
+ self.post_layer_norm = post_layer_norm
48
+ self.add_bias_linear = add_bias_linear
49
+ self.add_qkv_bias = add_qkv_bias
50
+ self.bias_dropout_fusion = bias_dropout_fusion
51
+ self.multi_query_attention = multi_query_attention
52
+ self.multi_query_group_num = multi_query_group_num
53
+ self.apply_query_key_layer_scaling = apply_query_key_layer_scaling
54
+ self.attention_softmax_in_fp32 = attention_softmax_in_fp32
55
+ self.fp32_residual_connection = fp32_residual_connection
56
+ self.quantization_bit = quantization_bit
57
+ self.pre_seq_len = pre_seq_len
58
+ self.prefix_projection = prefix_projection
59
+ super().__init__(**kwargs)
checkpoint-1000/finetuning_args.json ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "finetuning_type": "freeze",
3
+ "lora_alpha": 9.0,
4
+ "lora_dropout": 0.05,
5
+ "lora_rank": 100,
6
+ "lora_target": [
7
+ "query_key_value"
8
+ ],
9
+ "name_module_trainable": "mlp",
10
+ "num_layer_trainable": 1,
11
+ "pre_seq_len": 64,
12
+ "prefix_projection": false
13
+ }
checkpoint-1000/generation_config.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "eos_token_id": 2,
4
+ "pad_token_id": 0,
5
+ "transformers_version": "4.30.2"
6
+ }
checkpoint-1000/modeling_chatglm.py ADDED
@@ -0,0 +1,1193 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """ PyTorch ChatGLM model. """
2
+
3
+ import math
4
+ import copy
5
+ import warnings
6
+ import re
7
+ import sys
8
+
9
+ import torch
10
+ import torch.utils.checkpoint
11
+ import torch.nn.functional as F
12
+ from torch import nn
13
+ from torch.nn import CrossEntropyLoss, LayerNorm
14
+ from torch.nn.utils import skip_init
15
+ from typing import Optional, Tuple, Union, List, Callable, Dict, Any
16
+
17
+ from transformers.modeling_outputs import (
18
+ BaseModelOutputWithPast,
19
+ CausalLMOutputWithPast,
20
+ )
21
+ from transformers.modeling_utils import PreTrainedModel
22
+ from transformers.utils import logging
23
+ from transformers.generation.logits_process import LogitsProcessor
24
+ from transformers.generation.utils import LogitsProcessorList, StoppingCriteriaList, GenerationConfig, ModelOutput
25
+
26
+ from .configuration_chatglm import ChatGLMConfig
27
+
28
+ # flags required to enable jit fusion kernels
29
+
30
+ if sys.platform != 'darwin':
31
+ torch._C._jit_set_profiling_mode(False)
32
+ torch._C._jit_set_profiling_executor(False)
33
+ torch._C._jit_override_can_fuse_on_cpu(True)
34
+ torch._C._jit_override_can_fuse_on_gpu(True)
35
+
36
+ logger = logging.get_logger(__name__)
37
+
38
+ _CHECKPOINT_FOR_DOC = "THUDM/ChatGLM2-6B"
39
+ _CONFIG_FOR_DOC = "ChatGLM6BConfig"
40
+
41
+ CHATGLM_6B_PRETRAINED_MODEL_ARCHIVE_LIST = [
42
+ "THUDM/chatglm2-6b",
43
+ # See all ChatGLM models at https://huggingface.co/models?filter=chatglm
44
+ ]
45
+
46
+
47
+ def default_init(cls, *args, **kwargs):
48
+ return cls(*args, **kwargs)
49
+
50
+
51
+ class InvalidScoreLogitsProcessor(LogitsProcessor):
52
+ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
53
+ if torch.isnan(scores).any() or torch.isinf(scores).any():
54
+ scores.zero_()
55
+ scores[..., 5] = 5e4
56
+ return scores
57
+
58
+
59
+ class PrefixEncoder(torch.nn.Module):
60
+ """
61
+ The torch.nn model to encode the prefix
62
+ Input shape: (batch-size, prefix-length)
63
+ Output shape: (batch-size, prefix-length, 2*layers*hidden)
64
+ """
65
+
66
+ def __init__(self, config: ChatGLMConfig):
67
+ super().__init__()
68
+ self.prefix_projection = config.prefix_projection
69
+ if self.prefix_projection:
70
+ # Use a two-layer MLP to encode the prefix
71
+ kv_size = config.num_layers * config.kv_channels * config.multi_query_group_num * 2
72
+ self.embedding = torch.nn.Embedding(config.pre_seq_len, kv_size)
73
+ self.trans = torch.nn.Sequential(
74
+ torch.nn.Linear(kv_size, config.hidden_size),
75
+ torch.nn.Tanh(),
76
+ torch.nn.Linear(config.hidden_size, kv_size)
77
+ )
78
+ else:
79
+ self.embedding = torch.nn.Embedding(config.pre_seq_len,
80
+ config.num_layers * config.kv_channels * config.multi_query_group_num * 2)
81
+
82
+ def forward(self, prefix: torch.Tensor):
83
+ if self.prefix_projection:
84
+ prefix_tokens = self.embedding(prefix)
85
+ past_key_values = self.trans(prefix_tokens)
86
+ else:
87
+ past_key_values = self.embedding(prefix)
88
+ return past_key_values
89
+
90
+
91
+ def split_tensor_along_last_dim(
92
+ tensor: torch.Tensor,
93
+ num_partitions: int,
94
+ contiguous_split_chunks: bool = False,
95
+ ) -> List[torch.Tensor]:
96
+ """Split a tensor along its last dimension.
97
+
98
+ Arguments:
99
+ tensor: input tensor.
100
+ num_partitions: number of partitions to split the tensor
101
+ contiguous_split_chunks: If True, make each chunk contiguous
102
+ in memory.
103
+
104
+ Returns:
105
+ A list of Tensors
106
+ """
107
+ # Get the size and dimension.
108
+ last_dim = tensor.dim() - 1
109
+ last_dim_size = tensor.size()[last_dim] // num_partitions
110
+ # Split.
111
+ tensor_list = torch.split(tensor, last_dim_size, dim=last_dim)
112
+ # Note: torch.split does not create contiguous tensors by default.
113
+ if contiguous_split_chunks:
114
+ return tuple(chunk.contiguous() for chunk in tensor_list)
115
+
116
+ return tensor_list
117
+
118
+
119
+ class RotaryEmbedding(nn.Module):
120
+ def __init__(self, dim, original_impl=False, device=None, dtype=None):
121
+ super().__init__()
122
+ inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, device=device).to(dtype=dtype) / dim))
123
+ self.register_buffer("inv_freq", inv_freq)
124
+ self.dim = dim
125
+ self.original_impl = original_impl
126
+
127
+ def forward_impl(
128
+ self, seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000
129
+ ):
130
+ """Enhanced Transformer with Rotary Position Embedding.
131
+
132
+ Derived from: https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/
133
+ transformers/rope/__init__.py. MIT License:
134
+ https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license.
135
+ """
136
+ # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
137
+ theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, dtype=dtype, device=device) / n_elem))
138
+
139
+ # Create position indexes `[0, 1, ..., seq_len - 1]`
140
+ seq_idx = torch.arange(seq_len, dtype=dtype, device=device)
141
+
142
+ # Calculate the product of position index and $\theta_i$
143
+ idx_theta = torch.outer(seq_idx, theta).float()
144
+
145
+ cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1)
146
+
147
+ # this is to mimic the behaviour of complex32, else we will get different results
148
+ if dtype in (torch.float16, torch.bfloat16, torch.int8):
149
+ cache = cache.bfloat16() if dtype == torch.bfloat16 else cache.half()
150
+ return cache
151
+
152
+ def forward(self, max_seq_len, offset=0):
153
+ return self.forward_impl(
154
+ max_seq_len, self.dim, dtype=self.inv_freq.dtype, device=self.inv_freq.device
155
+ )
156
+
157
+
158
+ @torch.jit.script
159
+ def apply_rotary_pos_emb(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor:
160
+ # x: [sq, b, np, hn]
161
+ sq, b, np, hn = x.size(0), x.size(1), x.size(2), x.size(3)
162
+ rot_dim = rope_cache.shape[-2] * 2
163
+ x, x_pass = x[..., :rot_dim], x[..., rot_dim:]
164
+ # truncate to support variable sizes
165
+ rope_cache = rope_cache[:sq]
166
+ xshaped = x.reshape(sq, -1, np, rot_dim // 2, 2)
167
+ rope_cache = rope_cache.view(sq, -1, 1, xshaped.size(3), 2)
168
+ x_out2 = torch.stack(
169
+ [
170
+ xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1],
171
+ xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1],
172
+ ],
173
+ -1,
174
+ )
175
+ x_out2 = x_out2.flatten(3)
176
+ return torch.cat((x_out2, x_pass), dim=-1)
177
+
178
+
179
+ class RMSNorm(torch.nn.Module):
180
+ def __init__(self, normalized_shape, eps=1e-5, device=None, dtype=None, **kwargs):
181
+ super().__init__()
182
+ self.weight = torch.nn.Parameter(torch.empty(normalized_shape, device=device, dtype=dtype))
183
+ self.eps = eps
184
+
185
+ def forward(self, hidden_states: torch.Tensor):
186
+ input_dtype = hidden_states.dtype
187
+ variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
188
+ hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
189
+
190
+ return (self.weight * hidden_states).to(input_dtype)
191
+
192
+
193
+ class CoreAttention(torch.nn.Module):
194
+ def __init__(self, config: ChatGLMConfig, layer_number):
195
+ super(CoreAttention, self).__init__()
196
+
197
+ self.apply_query_key_layer_scaling = config.apply_query_key_layer_scaling
198
+ self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32
199
+ if self.apply_query_key_layer_scaling:
200
+ self.attention_softmax_in_fp32 = True
201
+ self.layer_number = max(1, layer_number)
202
+
203
+ projection_size = config.kv_channels * config.num_attention_heads
204
+
205
+ # Per attention head and per partition values.
206
+ self.hidden_size_per_partition = projection_size
207
+ self.hidden_size_per_attention_head = projection_size // config.num_attention_heads
208
+ self.num_attention_heads_per_partition = config.num_attention_heads
209
+
210
+ coeff = None
211
+ self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)
212
+ if self.apply_query_key_layer_scaling:
213
+ coeff = self.layer_number
214
+ self.norm_factor *= coeff
215
+ self.coeff = coeff
216
+
217
+ self.attention_dropout = torch.nn.Dropout(config.attention_dropout)
218
+
219
+ def forward(self, query_layer, key_layer, value_layer, attention_mask):
220
+ pytorch_major_version = int(torch.__version__.split('.')[0])
221
+ if pytorch_major_version >= 2:
222
+ query_layer, key_layer, value_layer = [k.permute(1, 2, 0, 3) for k in [query_layer, key_layer, value_layer]]
223
+ if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]:
224
+ context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
225
+ is_causal=True)
226
+ else:
227
+ if attention_mask is not None:
228
+ attention_mask = ~attention_mask
229
+ context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
230
+ attention_mask)
231
+ context_layer = context_layer.permute(2, 0, 1, 3)
232
+ new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
233
+ context_layer = context_layer.reshape(*new_context_layer_shape)
234
+ else:
235
+ # Raw attention scores
236
+
237
+ # [b, np, sq, sk]
238
+ output_size = (query_layer.size(1), query_layer.size(2), query_layer.size(0), key_layer.size(0))
239
+
240
+ # [sq, b, np, hn] -> [sq, b * np, hn]
241
+ query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1)
242
+ # [sk, b, np, hn] -> [sk, b * np, hn]
243
+ key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1)
244
+
245
+ # preallocting input tensor: [b * np, sq, sk]
246
+ matmul_input_buffer = torch.empty(
247
+ output_size[0] * output_size[1], output_size[2], output_size[3], dtype=query_layer.dtype,
248
+ device=query_layer.device
249
+ )
250
+
251
+ # Raw attention scores. [b * np, sq, sk]
252
+ matmul_result = torch.baddbmm(
253
+ matmul_input_buffer,
254
+ query_layer.transpose(0, 1), # [b * np, sq, hn]
255
+ key_layer.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk]
256
+ beta=0.0,
257
+ alpha=(1.0 / self.norm_factor),
258
+ )
259
+
260
+ # change view to [b, np, sq, sk]
261
+ attention_scores = matmul_result.view(*output_size)
262
+
263
+ # ===========================
264
+ # Attention probs and dropout
265
+ # ===========================
266
+
267
+ # attention scores and attention mask [b, np, sq, sk]
268
+ if self.attention_softmax_in_fp32:
269
+ attention_scores = attention_scores.float()
270
+ if self.coeff is not None:
271
+ attention_scores = attention_scores * self.coeff
272
+ if attention_mask is None and attention_scores.shape[2] == attention_scores.shape[3]:
273
+ attention_mask = torch.ones(output_size[0], 1, output_size[2], output_size[3],
274
+ device=attention_scores.device, dtype=torch.bool)
275
+ attention_mask.tril_()
276
+ attention_mask = ~attention_mask
277
+ if attention_mask is not None:
278
+ attention_scores = attention_scores.masked_fill(attention_mask, float("-inf"))
279
+ attention_probs = F.softmax(attention_scores, dim=-1)
280
+ attention_probs = attention_probs.type_as(value_layer)
281
+
282
+ # This is actually dropping out entire tokens to attend to, which might
283
+ # seem a bit unusual, but is taken from the original Transformer paper.
284
+ attention_probs = self.attention_dropout(attention_probs)
285
+ # =========================
286
+ # Context layer. [sq, b, hp]
287
+ # =========================
288
+
289
+ # value_layer -> context layer.
290
+ # [sk, b, np, hn] --> [b, np, sq, hn]
291
+
292
+ # context layer shape: [b, np, sq, hn]
293
+ output_size = (value_layer.size(1), value_layer.size(2), query_layer.size(0), value_layer.size(3))
294
+ # change view [sk, b * np, hn]
295
+ value_layer = value_layer.view(value_layer.size(0), output_size[0] * output_size[1], -1)
296
+ # change view [b * np, sq, sk]
297
+ attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
298
+ # matmul: [b * np, sq, hn]
299
+ context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1))
300
+ # change view [b, np, sq, hn]
301
+ context_layer = context_layer.view(*output_size)
302
+ # [b, np, sq, hn] --> [sq, b, np, hn]
303
+ context_layer = context_layer.permute(2, 0, 1, 3).contiguous()
304
+ # [sq, b, np, hn] --> [sq, b, hp]
305
+ new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
306
+ context_layer = context_layer.view(*new_context_layer_shape)
307
+
308
+ return context_layer
309
+
310
+
311
+ class SelfAttention(torch.nn.Module):
312
+ """Parallel self-attention layer abstract class.
313
+
314
+ Self-attention layer takes input with size [s, b, h]
315
+ and returns output of the same size.
316
+ """
317
+
318
+ def __init__(self, config: ChatGLMConfig, layer_number, device=None):
319
+ super(SelfAttention, self).__init__()
320
+ self.layer_number = max(1, layer_number)
321
+
322
+ self.projection_size = config.kv_channels * config.num_attention_heads
323
+
324
+ # Per attention head and per partition values.
325
+ self.hidden_size_per_attention_head = self.projection_size // config.num_attention_heads
326
+ self.num_attention_heads_per_partition = config.num_attention_heads
327
+
328
+ self.multi_query_attention = config.multi_query_attention
329
+ self.qkv_hidden_size = 3 * self.projection_size
330
+ if self.multi_query_attention:
331
+ self.num_multi_query_groups_per_partition = config.multi_query_group_num
332
+ self.qkv_hidden_size = (
333
+ self.projection_size + 2 * self.hidden_size_per_attention_head * config.multi_query_group_num
334
+ )
335
+ self.query_key_value = nn.Linear(config.hidden_size, self.qkv_hidden_size,
336
+ bias=config.add_bias_linear or config.add_qkv_bias,
337
+ device=device, **_config_to_kwargs(config)
338
+ )
339
+
340
+ self.core_attention = CoreAttention(config, self.layer_number)
341
+
342
+ # Output.
343
+ self.dense = nn.Linear(self.projection_size, config.hidden_size, bias=config.add_bias_linear,
344
+ device=device, **_config_to_kwargs(config)
345
+ )
346
+
347
+ def _allocate_memory(self, inference_max_sequence_len, batch_size, device=None, dtype=None):
348
+ if self.multi_query_attention:
349
+ num_attention_heads = self.num_multi_query_groups_per_partition
350
+ else:
351
+ num_attention_heads = self.num_attention_heads_per_partition
352
+ return torch.empty(
353
+ inference_max_sequence_len,
354
+ batch_size,
355
+ num_attention_heads,
356
+ self.hidden_size_per_attention_head,
357
+ dtype=dtype,
358
+ device=device,
359
+ )
360
+
361
+ def forward(
362
+ self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True
363
+ ):
364
+ # hidden_states: [sq, b, h]
365
+
366
+ # =================================================
367
+ # Pre-allocate memory for key-values for inference.
368
+ # =================================================
369
+ # =====================
370
+ # Query, Key, and Value
371
+ # =====================
372
+
373
+ # Attention heads [sq, b, h] --> [sq, b, (np * 3 * hn)]
374
+ mixed_x_layer = self.query_key_value(hidden_states)
375
+
376
+ if self.multi_query_attention:
377
+ (query_layer, key_layer, value_layer) = mixed_x_layer.split(
378
+ [
379
+ self.num_attention_heads_per_partition * self.hidden_size_per_attention_head,
380
+ self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
381
+ self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
382
+ ],
383
+ dim=-1,
384
+ )
385
+ query_layer = query_layer.view(
386
+ query_layer.size()[:-1] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)
387
+ )
388
+ key_layer = key_layer.view(
389
+ key_layer.size()[:-1] + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)
390
+ )
391
+ value_layer = value_layer.view(
392
+ value_layer.size()[:-1]
393
+ + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)
394
+ )
395
+ else:
396
+ new_tensor_shape = mixed_x_layer.size()[:-1] + \
397
+ (self.num_attention_heads_per_partition,
398
+ 3 * self.hidden_size_per_attention_head)
399
+ mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)
400
+
401
+ # [sq, b, np, 3 * hn] --> 3 [sq, b, np, hn]
402
+ (query_layer, key_layer, value_layer) = split_tensor_along_last_dim(mixed_x_layer, 3)
403
+
404
+ # apply relative positional encoding (rotary embedding)
405
+ if rotary_pos_emb is not None:
406
+ query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb)
407
+ key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb)
408
+
409
+ # adjust key and value for inference
410
+ if kv_cache is not None:
411
+ cache_k, cache_v = kv_cache
412
+ key_layer = torch.cat((cache_k, key_layer), dim=0)
413
+ value_layer = torch.cat((cache_v, value_layer), dim=0)
414
+ if use_cache:
415
+ kv_cache = (key_layer, value_layer)
416
+ else:
417
+ kv_cache = None
418
+
419
+ if self.multi_query_attention:
420
+ key_layer = key_layer.unsqueeze(-2)
421
+ key_layer = key_layer.expand(
422
+ -1, -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1
423
+ )
424
+ key_layer = key_layer.contiguous().view(
425
+ key_layer.size()[:2] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)
426
+ )
427
+ value_layer = value_layer.unsqueeze(-2)
428
+ value_layer = value_layer.expand(
429
+ -1, -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1
430
+ )
431
+ value_layer = value_layer.contiguous().view(
432
+ value_layer.size()[:2] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)
433
+ )
434
+
435
+ # ==================================
436
+ # core attention computation
437
+ # ==================================
438
+
439
+ context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask)
440
+
441
+ # =================
442
+ # Output. [sq, b, h]
443
+ # =================
444
+
445
+ output = self.dense(context_layer)
446
+
447
+ return output, kv_cache
448
+
449
+
450
+ def _config_to_kwargs(args):
451
+ common_kwargs = {
452
+ "dtype": args.torch_dtype,
453
+ }
454
+ return common_kwargs
455
+
456
+
457
+ class MLP(torch.nn.Module):
458
+ """MLP.
459
+
460
+ MLP will take the input with h hidden state, project it to 4*h
461
+ hidden dimension, perform nonlinear transformation, and project the
462
+ state back into h hidden dimension.
463
+ """
464
+
465
+ def __init__(self, config: ChatGLMConfig, device=None):
466
+ super(MLP, self).__init__()
467
+
468
+ self.add_bias = config.add_bias_linear
469
+
470
+ # Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf
471
+ self.dense_h_to_4h = nn.Linear(
472
+ config.hidden_size,
473
+ config.ffn_hidden_size * 2,
474
+ bias=self.add_bias,
475
+ device=device,
476
+ **_config_to_kwargs(config)
477
+ )
478
+
479
+ def swiglu(x):
480
+ x = torch.chunk(x, 2, dim=-1)
481
+ return F.silu(x[0]) * x[1]
482
+
483
+ self.activation_func = swiglu
484
+
485
+ # Project back to h.
486
+ self.dense_4h_to_h = nn.Linear(
487
+ config.ffn_hidden_size,
488
+ config.hidden_size,
489
+ bias=self.add_bias,
490
+ device=device,
491
+ **_config_to_kwargs(config)
492
+ )
493
+
494
+ def forward(self, hidden_states):
495
+ # [s, b, 4hp]
496
+ intermediate_parallel = self.dense_h_to_4h(hidden_states)
497
+ intermediate_parallel = self.activation_func(intermediate_parallel)
498
+ # [s, b, h]
499
+ output = self.dense_4h_to_h(intermediate_parallel)
500
+ return output
501
+
502
+
503
+ class GLMBlock(torch.nn.Module):
504
+ """A single transformer layer.
505
+
506
+ Transformer layer takes input with size [s, b, h] and returns an
507
+ output of the same size.
508
+ """
509
+
510
+ def __init__(self, config: ChatGLMConfig, layer_number, device=None):
511
+ super(GLMBlock, self).__init__()
512
+ self.layer_number = layer_number
513
+
514
+ self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm
515
+
516
+ self.fp32_residual_connection = config.fp32_residual_connection
517
+
518
+ LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm
519
+ # Layernorm on the input data.
520
+ self.input_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
521
+ dtype=config.torch_dtype)
522
+
523
+ # Self attention.
524
+ self.self_attention = SelfAttention(config, layer_number, device=device)
525
+ self.hidden_dropout = config.hidden_dropout
526
+
527
+ # Layernorm on the attention output
528
+ self.post_attention_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
529
+ dtype=config.torch_dtype)
530
+
531
+ # MLP
532
+ self.mlp = MLP(config, device=device)
533
+
534
+ def forward(
535
+ self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True,
536
+ ):
537
+ # hidden_states: [s, b, h]
538
+
539
+ # Layer norm at the beginning of the transformer layer.
540
+ layernorm_output = self.input_layernorm(hidden_states)
541
+ # Self attention.
542
+ attention_output, kv_cache = self.self_attention(
543
+ layernorm_output,
544
+ attention_mask,
545
+ rotary_pos_emb,
546
+ kv_cache=kv_cache,
547
+ use_cache=use_cache
548
+ )
549
+
550
+ # Residual connection.
551
+ if self.apply_residual_connection_post_layernorm:
552
+ residual = layernorm_output
553
+ else:
554
+ residual = hidden_states
555
+
556
+ layernorm_input = torch.nn.functional.dropout(attention_output, p=self.hidden_dropout, training=self.training)
557
+ layernorm_input = residual + layernorm_input
558
+
559
+ # Layer norm post the self attention.
560
+ layernorm_output = self.post_attention_layernorm(layernorm_input)
561
+
562
+ # MLP.
563
+ mlp_output = self.mlp(layernorm_output)
564
+
565
+ # Second residual connection.
566
+ if self.apply_residual_connection_post_layernorm:
567
+ residual = layernorm_output
568
+ else:
569
+ residual = layernorm_input
570
+
571
+ output = torch.nn.functional.dropout(mlp_output, p=self.hidden_dropout, training=self.training)
572
+ output = residual + output
573
+
574
+ return output, kv_cache
575
+
576
+
577
+ class GLMTransformer(torch.nn.Module):
578
+ """Transformer class."""
579
+
580
+ def __init__(self, config: ChatGLMConfig, device=None):
581
+ super(GLMTransformer, self).__init__()
582
+
583
+ self.fp32_residual_connection = config.fp32_residual_connection
584
+ self.post_layer_norm = config.post_layer_norm
585
+
586
+ # Number of layers.
587
+ self.num_layers = config.num_layers
588
+
589
+ # Transformer layers.
590
+ def build_layer(layer_number):
591
+ return GLMBlock(config, layer_number, device=device)
592
+
593
+ self.layers = torch.nn.ModuleList([build_layer(i + 1) for i in range(self.num_layers)])
594
+
595
+ if self.post_layer_norm:
596
+ LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm
597
+ # Final layer norm before output.
598
+ self.final_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
599
+ dtype=config.torch_dtype)
600
+
601
+ self.gradient_checkpointing = False
602
+
603
+ def _get_layer(self, layer_number):
604
+ return self.layers[layer_number]
605
+
606
+ def forward(
607
+ self, hidden_states, attention_mask, rotary_pos_emb, kv_caches=None,
608
+ use_cache: Optional[bool] = True,
609
+ output_hidden_states: Optional[bool] = False,
610
+ ):
611
+ if not kv_caches:
612
+ kv_caches = [None for _ in range(self.num_layers)]
613
+ presents = () if use_cache else None
614
+ if self.gradient_checkpointing and self.training:
615
+ if use_cache:
616
+ logger.warning_once(
617
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
618
+ )
619
+ use_cache = False
620
+
621
+ all_self_attentions = None
622
+ all_hidden_states = () if output_hidden_states else None
623
+ for index in range(self.num_layers):
624
+ if output_hidden_states:
625
+ all_hidden_states = all_hidden_states + (hidden_states,)
626
+
627
+ layer = self._get_layer(index)
628
+ if self.gradient_checkpointing and self.training:
629
+ layer_ret = torch.utils.checkpoint.checkpoint(
630
+ layer,
631
+ hidden_states,
632
+ attention_mask,
633
+ rotary_pos_emb,
634
+ kv_caches[index],
635
+ use_cache
636
+ )
637
+ else:
638
+ layer_ret = layer(
639
+ hidden_states,
640
+ attention_mask,
641
+ rotary_pos_emb,
642
+ kv_cache=kv_caches[index],
643
+ use_cache=use_cache
644
+ )
645
+ hidden_states, kv_cache = layer_ret
646
+ if use_cache:
647
+ presents = presents + (kv_cache,)
648
+
649
+ if output_hidden_states:
650
+ all_hidden_states = all_hidden_states + (hidden_states,)
651
+
652
+ # Final layer norm.
653
+ if self.post_layer_norm:
654
+ hidden_states = self.final_layernorm(hidden_states)
655
+
656
+ return hidden_states, presents, all_hidden_states, all_self_attentions
657
+
658
+
659
+ class ChatGLMPreTrainedModel(PreTrainedModel):
660
+ """
661
+ An abstract class to handle weights initialization and
662
+ a simple interface for downloading and loading pretrained models.
663
+ """
664
+
665
+ is_parallelizable = False
666
+ supports_gradient_checkpointing = True
667
+ config_class = ChatGLMConfig
668
+ base_model_prefix = "transformer"
669
+ _no_split_modules = ["GLMBlock"]
670
+
671
+ def _init_weights(self, module: nn.Module):
672
+ """Initialize the weights."""
673
+ return
674
+
675
+ def get_masks(self, input_ids, past_key_values, padding_mask=None):
676
+ batch_size, seq_length = input_ids.shape
677
+ full_attention_mask = torch.ones(batch_size, seq_length, seq_length, device=input_ids.device)
678
+ full_attention_mask.tril_()
679
+ past_length = 0
680
+ if past_key_values:
681
+ past_length = past_key_values[0][0].shape[0]
682
+ if past_length:
683
+ full_attention_mask = torch.cat((torch.ones(batch_size, seq_length, past_length,
684
+ device=input_ids.device), full_attention_mask), dim=-1)
685
+ if padding_mask is not None:
686
+ full_attention_mask = full_attention_mask * padding_mask.unsqueeze(1)
687
+ if not past_length and padding_mask is not None:
688
+ full_attention_mask -= padding_mask.unsqueeze(-1) - 1
689
+ full_attention_mask = (full_attention_mask < 0.5).bool()
690
+ full_attention_mask.unsqueeze_(1)
691
+ return full_attention_mask
692
+
693
+ def get_position_ids(self, input_ids, device):
694
+ batch_size, seq_length = input_ids.shape
695
+ position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1)
696
+ return position_ids
697
+
698
+ def _set_gradient_checkpointing(self, module, value=False):
699
+ if isinstance(module, GLMTransformer):
700
+ module.gradient_checkpointing = value
701
+
702
+
703
+ class Embedding(torch.nn.Module):
704
+ """Language model embeddings."""
705
+
706
+ def __init__(self, config: ChatGLMConfig, device=None):
707
+ super(Embedding, self).__init__()
708
+
709
+ self.hidden_size = config.hidden_size
710
+ # Word embeddings (parallel).
711
+ self.word_embeddings = nn.Embedding(
712
+ config.padded_vocab_size,
713
+ self.hidden_size,
714
+ dtype=config.torch_dtype,
715
+ device=device
716
+ )
717
+ self.fp32_residual_connection = config.fp32_residual_connection
718
+
719
+ def forward(self, input_ids):
720
+ # Embeddings.
721
+ words_embeddings = self.word_embeddings(input_ids)
722
+ embeddings = words_embeddings
723
+ # Data format change to avoid explicit tranposes : [b s h] --> [s b h].
724
+ embeddings = embeddings.transpose(0, 1).contiguous()
725
+ # If the input flag for fp32 residual connection is set, convert for float.
726
+ if self.fp32_residual_connection:
727
+ embeddings = embeddings.float()
728
+ return embeddings
729
+
730
+
731
+ class ChatGLMModel(ChatGLMPreTrainedModel):
732
+ def __init__(self, config: ChatGLMConfig, device=None, empty_init=True):
733
+ super().__init__(config)
734
+ if empty_init:
735
+ init_method = skip_init
736
+ else:
737
+ init_method = default_init
738
+ init_kwargs = {}
739
+ if device is not None:
740
+ init_kwargs["device"] = device
741
+ self.embedding = init_method(Embedding, config, **init_kwargs)
742
+ self.num_layers = config.num_layers
743
+ self.multi_query_group_num = config.multi_query_group_num
744
+ self.kv_channels = config.kv_channels
745
+
746
+ # Rotary positional embeddings
747
+ self.seq_length = config.seq_length
748
+ rotary_dim = (
749
+ config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels
750
+ )
751
+
752
+ self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, original_impl=config.original_rope, device=device,
753
+ dtype=config.torch_dtype)
754
+ self.encoder = init_method(GLMTransformer, config, **init_kwargs)
755
+ self.output_layer = init_method(nn.Linear, config.hidden_size, config.padded_vocab_size, bias=False,
756
+ dtype=config.torch_dtype, **init_kwargs)
757
+ self.pre_seq_len = config.pre_seq_len
758
+ self.prefix_projection = config.prefix_projection
759
+ if self.pre_seq_len is not None:
760
+ for param in self.parameters():
761
+ param.requires_grad = False
762
+ self.prefix_tokens = torch.arange(self.pre_seq_len).long()
763
+ self.prefix_encoder = PrefixEncoder(config)
764
+ self.dropout = torch.nn.Dropout(0.1)
765
+
766
+ def get_input_embeddings(self):
767
+ return self.embedding.word_embeddings
768
+
769
+ def get_prompt(self, batch_size, device, dtype=torch.half):
770
+ prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(device)
771
+ past_key_values = self.prefix_encoder(prefix_tokens).type(dtype)
772
+ past_key_values = past_key_values.view(
773
+ batch_size,
774
+ self.pre_seq_len,
775
+ self.num_layers * 2,
776
+ self.multi_query_group_num,
777
+ self.kv_channels
778
+ )
779
+ # seq_len, b, nh, hidden_size
780
+ past_key_values = self.dropout(past_key_values)
781
+ past_key_values = past_key_values.permute([2, 1, 0, 3, 4]).split(2)
782
+ return past_key_values
783
+
784
+ def forward(
785
+ self,
786
+ input_ids,
787
+ position_ids: Optional[torch.Tensor] = None,
788
+ attention_mask: Optional[torch.BoolTensor] = None,
789
+ full_attention_mask: Optional[torch.BoolTensor] = None,
790
+ past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
791
+ inputs_embeds: Optional[torch.Tensor] = None,
792
+ use_cache: Optional[bool] = None,
793
+ output_hidden_states: Optional[bool] = None,
794
+ return_dict: Optional[bool] = None,
795
+ ):
796
+ output_hidden_states = (
797
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
798
+ )
799
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
800
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
801
+
802
+ batch_size, seq_length = input_ids.shape
803
+
804
+ if inputs_embeds is None:
805
+ inputs_embeds = self.embedding(input_ids)
806
+
807
+ if self.pre_seq_len is not None:
808
+ if past_key_values is None:
809
+ past_key_values = self.get_prompt(batch_size=batch_size, device=input_ids.device,
810
+ dtype=inputs_embeds.dtype)
811
+ if attention_mask is not None:
812
+ attention_mask = torch.cat([attention_mask.new_ones((batch_size, self.pre_seq_len)),
813
+ attention_mask], dim=-1)
814
+
815
+ if full_attention_mask is None:
816
+ if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1):
817
+ full_attention_mask = self.get_masks(input_ids, past_key_values, padding_mask=attention_mask)
818
+
819
+ # Rotary positional embeddings
820
+ rotary_pos_emb = self.rotary_pos_emb(self.seq_length)
821
+ if position_ids is not None:
822
+ rotary_pos_emb = rotary_pos_emb[position_ids]
823
+ else:
824
+ rotary_pos_emb = rotary_pos_emb[None, :seq_length]
825
+ rotary_pos_emb = rotary_pos_emb.transpose(0, 1).contiguous()
826
+
827
+ # Run encoder.
828
+ hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder(
829
+ inputs_embeds, full_attention_mask, rotary_pos_emb=rotary_pos_emb,
830
+ kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states
831
+ )
832
+
833
+ if not return_dict:
834
+ return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)
835
+
836
+ return BaseModelOutputWithPast(
837
+ last_hidden_state=hidden_states,
838
+ past_key_values=presents,
839
+ hidden_states=all_hidden_states,
840
+ attentions=all_self_attentions,
841
+ )
842
+
843
+ def quantize(self, weight_bit_width: int):
844
+ from .quantization import quantize
845
+ quantize(self.encoder, weight_bit_width)
846
+ return self
847
+
848
+
849
+ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
850
+ def __init__(self, config: ChatGLMConfig, empty_init=True, device=None):
851
+ super().__init__(config)
852
+
853
+ self.max_sequence_length = config.max_length
854
+ self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device)
855
+ self.config = config
856
+ self.quantized = False
857
+
858
+ if self.config.quantization_bit:
859
+ self.quantize(self.config.quantization_bit, empty_init=True)
860
+
861
+ def _update_model_kwargs_for_generation(
862
+ self,
863
+ outputs: ModelOutput,
864
+ model_kwargs: Dict[str, Any],
865
+ is_encoder_decoder: bool = False,
866
+ standardize_cache_format: bool = False,
867
+ ) -> Dict[str, Any]:
868
+ # update past_key_values
869
+ model_kwargs["past_key_values"] = self._extract_past_from_model_output(
870
+ outputs, standardize_cache_format=standardize_cache_format
871
+ )
872
+
873
+ # update attention mask
874
+ if "attention_mask" in model_kwargs:
875
+ attention_mask = model_kwargs["attention_mask"]
876
+ model_kwargs["attention_mask"] = torch.cat(
877
+ [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
878
+ )
879
+
880
+ # update position ids
881
+ if "position_ids" in model_kwargs:
882
+ position_ids = model_kwargs["position_ids"]
883
+ new_position_id = position_ids[..., -1:].clone()
884
+ new_position_id += 1
885
+ model_kwargs["position_ids"] = torch.cat(
886
+ [position_ids, new_position_id], dim=-1
887
+ )
888
+
889
+ model_kwargs["is_first_forward"] = False
890
+ return model_kwargs
891
+
892
+ def prepare_inputs_for_generation(
893
+ self,
894
+ input_ids: torch.LongTensor,
895
+ past_key_values: Optional[torch.Tensor] = None,
896
+ attention_mask: Optional[torch.Tensor] = None,
897
+ position_ids: Optional[torch.Tensor] = None,
898
+ is_first_forward: bool = True,
899
+ **kwargs
900
+ ) -> dict:
901
+ # only last token for input_ids if past is not None
902
+ if position_ids is None:
903
+ position_ids = self.get_position_ids(input_ids, device=input_ids.device)
904
+ if not is_first_forward:
905
+ position_ids = position_ids[..., -1:]
906
+ input_ids = input_ids[:, -1:]
907
+ return {
908
+ "input_ids": input_ids,
909
+ "past_key_values": past_key_values,
910
+ "position_ids": position_ids,
911
+ "attention_mask": attention_mask,
912
+ "return_last_logit": True
913
+ }
914
+
915
+ def forward(
916
+ self,
917
+ input_ids: Optional[torch.Tensor] = None,
918
+ position_ids: Optional[torch.Tensor] = None,
919
+ attention_mask: Optional[torch.Tensor] = None,
920
+ past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
921
+ inputs_embeds: Optional[torch.Tensor] = None,
922
+ labels: Optional[torch.Tensor] = None,
923
+ use_cache: Optional[bool] = None,
924
+ output_attentions: Optional[bool] = None,
925
+ output_hidden_states: Optional[bool] = None,
926
+ return_dict: Optional[bool] = None,
927
+ return_last_logit: Optional[bool] = False,
928
+ ):
929
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
930
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
931
+
932
+ transformer_outputs = self.transformer(
933
+ input_ids=input_ids,
934
+ position_ids=position_ids,
935
+ attention_mask=attention_mask,
936
+ past_key_values=past_key_values,
937
+ inputs_embeds=inputs_embeds,
938
+ use_cache=use_cache,
939
+ output_hidden_states=output_hidden_states,
940
+ return_dict=return_dict,
941
+ )
942
+
943
+ hidden_states = transformer_outputs[0]
944
+ if return_last_logit:
945
+ hidden_states = hidden_states[-1:]
946
+ lm_logits = self.transformer.output_layer(hidden_states)
947
+ lm_logits = lm_logits.transpose(0, 1).contiguous()
948
+
949
+ loss = None
950
+ if labels is not None:
951
+ lm_logits = lm_logits.to(torch.float32)
952
+
953
+ # Shift so that tokens < n predict n
954
+ shift_logits = lm_logits[..., :-1, :].contiguous()
955
+ shift_labels = labels[..., 1:].contiguous()
956
+ # Flatten the tokens
957
+ loss_fct = CrossEntropyLoss(ignore_index=-100)
958
+ loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
959
+
960
+ lm_logits = lm_logits.to(hidden_states.dtype)
961
+ loss = loss.to(hidden_states.dtype)
962
+
963
+ if not return_dict:
964
+ output = (lm_logits,) + transformer_outputs[1:]
965
+ return ((loss,) + output) if loss is not None else output
966
+
967
+ return CausalLMOutputWithPast(
968
+ loss=loss,
969
+ logits=lm_logits,
970
+ past_key_values=transformer_outputs.past_key_values,
971
+ hidden_states=transformer_outputs.hidden_states,
972
+ attentions=transformer_outputs.attentions,
973
+ )
974
+
975
+ @staticmethod
976
+ def _reorder_cache(
977
+ past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor
978
+ ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:
979
+ """
980
+ This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
981
+ [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
982
+ beam_idx at every generation step.
983
+
984
+ Output shares the same memory storage as `past`.
985
+ """
986
+ return tuple(
987
+ (
988
+ layer_past[0].index_select(1, beam_idx.to(layer_past[0].device)),
989
+ layer_past[1].index_select(1, beam_idx.to(layer_past[1].device)),
990
+ )
991
+ for layer_past in past
992
+ )
993
+
994
+ def process_response(self, response):
995
+ response = response.strip()
996
+ response = response.replace("[[训练时间]]", "2023年")
997
+ return response
998
+
999
+ def build_inputs(self, tokenizer, query: str, history: List[Tuple[str, str]] = None):
1000
+ prompt = tokenizer.build_prompt(query, history=history)
1001
+ inputs = tokenizer([prompt], return_tensors="pt")
1002
+ inputs = inputs.to(self.device)
1003
+ return inputs
1004
+
1005
+ def build_stream_inputs(self, tokenizer, query: str, history: List[Tuple[str, str]] = None):
1006
+ if history:
1007
+ prompt = "\n\n[Round {}]\n\n问:{}\n\n答:".format(len(history) + 1, query)
1008
+ input_ids = tokenizer.encode(prompt, add_special_tokens=False)
1009
+ input_ids = input_ids[1:]
1010
+ inputs = tokenizer.batch_encode_plus([(input_ids, None)], return_tensors="pt", add_special_tokens=False)
1011
+ else:
1012
+ prompt = "[Round {}]\n\n问:{}\n\n答:".format(len(history) + 1, query)
1013
+ inputs = tokenizer([prompt], return_tensors="pt")
1014
+ inputs = inputs.to(self.device)
1015
+ return inputs
1016
+
1017
+ @torch.no_grad()
1018
+ def chat(self, tokenizer, query: str, history: List[Tuple[str, str]] = None, max_length: int = 8192, num_beams=1,
1019
+ do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None, **kwargs):
1020
+ if history is None:
1021
+ history = []
1022
+ if logits_processor is None:
1023
+ logits_processor = LogitsProcessorList()
1024
+ logits_processor.append(InvalidScoreLogitsProcessor())
1025
+ gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "do_sample": do_sample, "top_p": top_p,
1026
+ "temperature": temperature, "logits_processor": logits_processor, **kwargs}
1027
+ inputs = self.build_inputs(tokenizer, query, history=history)
1028
+ outputs = self.generate(**inputs, **gen_kwargs)
1029
+ outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):]
1030
+ response = tokenizer.decode(outputs)
1031
+ response = self.process_response(response)
1032
+ history = history + [(query, response)]
1033
+ return response, history
1034
+
1035
+ @torch.no_grad()
1036
+ def stream_chat(self, tokenizer, query: str, history: List[Tuple[str, str]] = None, past_key_values=None,
1037
+ max_length: int = 8192, do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None,
1038
+ return_past_key_values=False, **kwargs):
1039
+ if history is None:
1040
+ history = []
1041
+ if logits_processor is None:
1042
+ logits_processor = LogitsProcessorList()
1043
+ logits_processor.append(InvalidScoreLogitsProcessor())
1044
+ gen_kwargs = {"max_length": max_length, "do_sample": do_sample, "top_p": top_p,
1045
+ "temperature": temperature, "logits_processor": logits_processor, **kwargs}
1046
+ if past_key_values is None and not return_past_key_values:
1047
+ inputs = self.build_inputs(tokenizer, query, history=history)
1048
+ else:
1049
+ inputs = self.build_stream_inputs(tokenizer, query, history=history)
1050
+ if past_key_values is not None:
1051
+ past_length = past_key_values[0][0].shape[0]
1052
+ if self.transformer.pre_seq_len is not None:
1053
+ past_length -= self.transformer.pre_seq_len
1054
+ inputs.position_ids += past_length
1055
+ attention_mask = inputs.attention_mask
1056
+ attention_mask = torch.cat((attention_mask.new_ones(1, past_length), attention_mask), dim=1)
1057
+ inputs['attention_mask'] = attention_mask
1058
+ for outputs in self.stream_generate(**inputs, past_key_values=past_key_values,
1059
+ return_past_key_values=return_past_key_values, **gen_kwargs):
1060
+ if return_past_key_values:
1061
+ outputs, past_key_values = outputs
1062
+ outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):]
1063
+ response = tokenizer.decode(outputs)
1064
+ if response and response[-1] != "�":
1065
+ response = self.process_response(response)
1066
+ new_history = history + [(query, response)]
1067
+ if return_past_key_values:
1068
+ yield response, new_history, past_key_values
1069
+ else:
1070
+ yield response, new_history
1071
+
1072
+ @torch.no_grad()
1073
+ def stream_generate(
1074
+ self,
1075
+ input_ids,
1076
+ generation_config: Optional[GenerationConfig] = None,
1077
+ logits_processor: Optional[LogitsProcessorList] = None,
1078
+ stopping_criteria: Optional[StoppingCriteriaList] = None,
1079
+ prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
1080
+ return_past_key_values=False,
1081
+ **kwargs,
1082
+ ):
1083
+ batch_size, input_ids_seq_length = input_ids.shape[0], input_ids.shape[-1]
1084
+
1085
+ if generation_config is None:
1086
+ generation_config = self.generation_config
1087
+ generation_config = copy.deepcopy(generation_config)
1088
+ model_kwargs = generation_config.update(**kwargs)
1089
+ bos_token_id, eos_token_id = generation_config.bos_token_id, generation_config.eos_token_id
1090
+
1091
+ if isinstance(eos_token_id, int):
1092
+ eos_token_id = [eos_token_id]
1093
+
1094
+ has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
1095
+ if has_default_max_length and generation_config.max_new_tokens is None:
1096
+ warnings.warn(
1097
+ f"Using `max_length`'s default ({generation_config.max_length}) to control the generation length. "
1098
+ "This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we"
1099
+ " recommend using `max_new_tokens` to control the maximum length of the generation.",
1100
+ UserWarning,
1101
+ )
1102
+ elif generation_config.max_new_tokens is not None:
1103
+ generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length
1104
+ if not has_default_max_length:
1105
+ logger.warn(
1106
+ f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="
1107
+ f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "
1108
+ "Please refer to the documentation for more information. "
1109
+ "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)",
1110
+ UserWarning,
1111
+ )
1112
+
1113
+ if input_ids_seq_length >= generation_config.max_length:
1114
+ input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"
1115
+ logger.warning(
1116
+ f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to"
1117
+ f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
1118
+ " increasing `max_new_tokens`."
1119
+ )
1120
+
1121
+ # 2. Set generation parameters if not already defined
1122
+ logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
1123
+ stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
1124
+
1125
+ logits_processor = self._get_logits_processor(
1126
+ generation_config=generation_config,
1127
+ input_ids_seq_length=input_ids_seq_length,
1128
+ encoder_input_ids=input_ids,
1129
+ prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
1130
+ logits_processor=logits_processor,
1131
+ )
1132
+
1133
+ stopping_criteria = self._get_stopping_criteria(
1134
+ generation_config=generation_config, stopping_criteria=stopping_criteria
1135
+ )
1136
+ logits_warper = self._get_logits_warper(generation_config)
1137
+
1138
+ unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1)
1139
+ scores = None
1140
+ while True:
1141
+ model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
1142
+ # forward pass to get next token
1143
+ outputs = self(
1144
+ **model_inputs,
1145
+ return_dict=True,
1146
+ output_attentions=False,
1147
+ output_hidden_states=False,
1148
+ )
1149
+
1150
+ next_token_logits = outputs.logits[:, -1, :]
1151
+
1152
+ # pre-process distribution
1153
+ next_token_scores = logits_processor(input_ids, next_token_logits)
1154
+ next_token_scores = logits_warper(input_ids, next_token_scores)
1155
+
1156
+ # sample
1157
+ probs = nn.functional.softmax(next_token_scores, dim=-1)
1158
+ if generation_config.do_sample:
1159
+ next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
1160
+ else:
1161
+ next_tokens = torch.argmax(probs, dim=-1)
1162
+
1163
+ # update generated ids, model inputs, and length for next step
1164
+ input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
1165
+ model_kwargs = self._update_model_kwargs_for_generation(
1166
+ outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder
1167
+ )
1168
+ unfinished_sequences = unfinished_sequences.mul((sum(next_tokens != i for i in eos_token_id)).long())
1169
+ if return_past_key_values:
1170
+ yield input_ids, outputs.past_key_values
1171
+ else:
1172
+ yield input_ids
1173
+ # stop when each sentence is finished, or if we exceed the maximum length
1174
+ if unfinished_sequences.max() == 0 or stopping_criteria(input_ids, scores):
1175
+ break
1176
+
1177
+ def quantize(self, bits: int, empty_init=False, device=None, **kwargs):
1178
+ if bits == 0:
1179
+ return
1180
+
1181
+ from .quantization import quantize
1182
+
1183
+ if self.quantized:
1184
+ logger.info("Already quantized.")
1185
+ return self
1186
+
1187
+ self.quantized = True
1188
+
1189
+ self.config.quantization_bit = bits
1190
+
1191
+ self.transformer.encoder = quantize(self.transformer.encoder, bits, empty_init=empty_init, device=device,
1192
+ **kwargs)
1193
+ return self
checkpoint-1000/optimizer.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:a951e4053f67c9a8feede35b809b11fcae48b03032bd565ab1df2b64d3badb7e
3
+ size 1346374055
checkpoint-1000/pytorch_model.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c788dd8fa4a4721a8a86994609e4d08f4965703b89a9aeb216fc8f46a73b13c2
3
+ size 673186943
checkpoint-1000/quantization.py ADDED
@@ -0,0 +1,188 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from torch.nn import Linear
2
+ from torch.nn.parameter import Parameter
3
+
4
+ import bz2
5
+ import torch
6
+ import base64
7
+ import ctypes
8
+ from transformers.utils import logging
9
+
10
+ from typing import List
11
+ from functools import partial
12
+
13
+ logger = logging.get_logger(__name__)
14
+
15
+ try:
16
+ from cpm_kernels.kernels.base import LazyKernelCModule, KernelFunction, round_up
17
+
18
+ class Kernel:
19
+ def __init__(self, code: bytes, function_names: List[str]):
20
+ self.code = code
21
+ self._function_names = function_names
22
+ self._cmodule = LazyKernelCModule(self.code)
23
+
24
+ for name in self._function_names:
25
+ setattr(self, name, KernelFunction(self._cmodule, name))
26
+
27
+ quantization_code = "$QlpoOTFBWSZTWU9yuJUAQHN//////////f/n/8/n///n//bt4dTidcVx8X3V9FV/92/v4B7/AD5FBQFAAAChSgKpFCFAFVSigUAAAEKhSgUUqgFBKigqVREQAABQBQIANDTTIGI00BkZBkNGE0A0BkBkGQGRkaNAaAGQNBoGgDIAAYIGTI0DQAQAaGmmQMRpoDIyDIaMJoBoDIDIMgMjI0aA0AMgaDQNAGQAAwQMmRoGgAgA0NNMgYjTQGRkGQ0YTQDQGQGQZAZGRo0BoAZA0GgaAMgABggZMjQNABABoaaZAxGmgMjIMhowmgGgMgMgyAyMjRoDQAyBoNA0AZAADBAyZGgaAAmqU1NEgJqnptU/Sn4jRR6J6epk2pqb1Q/SgAPUGgyNNGjQ2SBpoAZAAGg0NB6mgDIAAAAA2oaApSREBNAARhGiYEaEwU8pvImlP0k2aam1GaGqbFNM1MHpTwmkepmyU9R6nqPKekHqNNPUxNGhp6n6p6QaZ6o9TG1GMqcoV9ly6nRanHlq6zPNbnGZNi6HSug+2nPiZ13XcnFYZW+45W11CumhzYhchOJ2GLLV1OBjBjGf4TptOddTSOcVxhqYZMYwZXZZY00zI1paX5X9J+b+f4e+x43RXSxXPOdquiGpduatGyXneN696M9t4HU2eR5XX/kPhP261NTx3JO1Ow7LyuDmeo9a7d351T1ZxnvnrvYnrXv/hXxPCeuYx2XsNmO003eg9J3Z6U7b23meJ4ri01OdzTk9BNO96brz+qT5nuvvH3ds/G+m/JcG/F2XYuhXlvO+jP7U3XgrzPN/lr8Sf1n6j4j7jZs+s/T0tNaNNYzTs12rxjwztHlnire3Nzc3N1wuBwOBwXBvZfoHpD7rFmR99V5vj3aXza3xdBbXMalubTg/jIv5dfAi54Pdc75j4z412n3Npj3Ld/ENm7a3b/Cod6h/ret1/5vn/C+l+gdslMvgPSLJ8d8q+U66fevYn/tW1chleEtNTGlcHCbLRlq0tHzF5tsbbZZfHjjLgZu42XCuC3NrdjTasZGNzgxPIrGqp7r3p7L2p5XjnpPSmTd5XtzqnB6U87zzg1Ol0zd0zsLszxR6lkxp35u6/teL0L0W922cR7Lu1lpL9CsHirzuM2T+BgsyViT6LHcm0/Vr6U/7LGGyJeqTEjt0PHWhF5mCT7R9mtlDwriYv0Tyr/OxYt6qp5r0mPVT0608TqnqMZaarU2nFwrTzzlrs1ed7z1ux60wyr4ydCaTi3enW8x68x0zU7tXSlcmPSW1mGpWJMg4zmPC2lK96tp0OE80y4MfEvnZj8zGluR6b22ki1Ou9V2nCd9xovcPvcYMZYy0lvN60ScZ45vN6yeCeeXFb1lVjnnCar5fwXwE2bzJ4HI1XVPXfXZMm44GUsMpYsmLB65TuVdm0cl0b+i/wGNN66XjeV7zuPpHcnK/juhhjdfId5jMdE5nN0dGmmm2zZs2cexD5n9p/dY352XsvXHaZNWWsmmS1atjR452nYudzvqv2HMRyvNNnlMcDl3R2+yx2uVrBubTW9icHDVtbNXlZm7jma1rM4VurZZd2y6nUau7ZXZ7bVU+mnoOVxZGMrVmvX60605JwmzGZhhhjTWtaaaMaaGTGmNMZasY0iX8VMUl8eepaIrzGSpemWOQyZORk2bNpjUybMmxqYmknCGCFynutfksaZpjTNMaaatM0xsxcGR0sociNqxNSmhhR1ZJPbsn8qyF0t2qH6iYBclclalbtTTcHTDsPaX6rlnElph2Jyumumtynv2Kk8GI7rsvXbIcJgHJOSaSXnnGaI3m87RtVXJOZ/YtgdTE6Wpha6ZlE8ayXkef1fh602r2WwvfMXtMdLlkfnLFdYYwYso+bWqm7yJqHXZGw2nrS5ZanSYnWlxBxMF1V940K2wdrI7R6OYf7DGGamMmTSbRhlS45xmVOumF1EyPCmHrrN8wwZOOrdNtLeMtzFzDlWnfTBxMk2NaXIZHBYxYLD4w8yju0ao65Vz1OIXoS9dLanwCe1PWrYuWMqf1if1z2k2yYfKJ741PDgno1ZQ8DRqvUny3mNoWTzGO6m1DkrJI8JiR5cSd+vZdGOO8nrMoc5+NDUFsMSXaZJeNlMmGLtJsovOsUp7I9S5VojKxF6bTVEelXqlfJobQr3LozSh2Jk7VcrVMfhXqszGWMzNqGhqZY0OadxkyyMssKugZR0KNFXBHlqwmJgTE/BNVMk6ItJXZMR0H47GpXv/DMOvNkmVuaV1PRfEdxuqc7Hcd+ZV/zTLaRxWk0nl9CdCeM6mn5rstHIBcpiuwmUZXeq81DacHI2rmrZ5SuE5mOZd6LQrZg9mx32TprA8BMo5jKN6yLTCi3WzQaZSuhzTtM1fUTGVpG8Tw+KXI0tjEpiWxtLYynOlktSbVlaI5kxP8TDH8kx50xoxi5KcA4pcja8KWLRlO/Ks6q06ergnvm1ca3Tq8Uw7LTUsmWyctXPWmpitl/uvGcWTGXGuAXDfhqazGmjkxcJW5hMMMMpYsXl2TZYtVOddG3XCarUt6Ptq9CZXSNzyuRzqRZOjsxdBbFVz6OA5HI43r1jityVlVpVkxmOsyaYWE1NTGq1sOVh36mHMcxtSvcy70edG0ZGR3I1Go1GRlV7mWWo1G0ZGRqlvH40l7o4m5xMWLLLYyNjnqc8556mdPqLJ31n/1nWOncxzG1tizrHs/Z+d2vP/B/l8wdJ6rHUn2nbbDq4p6htFtYzMMMTaZis1K5GKzGNmxhmUx2DDlZ/qNnIx41xnaMfCZWYaZWtNLTNW8ND4Fw1MyZOCdM428suKG1ehW8TesOydg7J+YYcD4cYR+8dFK6M4E3HM9ZfRNNL+Sn6rsl4DsrDl2HpPCnfxjGXtbZtYys1ttlyJ4T+BvexjGWRjMszK4Jpc77D3GyuVD7q0+G8m9G+2+rGm7cOR2y7FdtY2XUYx/oNlfRYxhMYyYZkyyg55enna9Kt/FFi6GMMwYwdwxWgxGMLKYmUyGExTKMZkMFhkymKuh0NOBNnBu+23LdwDoZYYzGGMxtORaTU1pjTGWTTGGtMrNWUsyyTTLLG1qy2ZjbK2DBllWqxMtBMaYZQmcE7zvvRcTkclUwdkxTaSdyySt/7fpL+T1v516Ji97fwr5JbLu305zMn5+GMTTZ9F+y7ExwmGVfG44yxn3dLv6l5i+Wth1jCrDq21nW9LqvvDzz3Vf3LLH/O/32TJ/erx3bXftO4eF+G956D952K/An4NfvOpjFjExjevP/UmE0fIoZXx6/w6lX/no3D0bLt+ixjieBM6ksRd0yB4Lt2SwYNE+gd1detlZWUnpiZfGfFaK+4PyCa/v18V8X75pe9fLXzp7l3VjF76vWZmHwGz1IZNWT7b8yddJ4q5kyrVdfru6atWc7bVYztL9Jf4GXvT+Y8m9/YsXP6H018a8D4XVOqvfzqeR+6yZOD8dPv0+U7/q5Pl+2dNb0MjzGVH5p6MNQ7cOWvw62U9aHE8DprDek+McLyvDz+te+9Zhq5+YTruufMcWMabqysTmZVWjKPfnK0wyVcrsuhjZRdLkHNvD72b9abriOSGIxiLixMOoalNPXzy+wT/tf+U6HHONfsz+xe8ufHBdQWWGWLA9if0rsnmrxK5LvRZQeWsTCsrmOYy8VteVfuRfcVTtDLItLIsMYxZLdU/DbtSemxF6Z6Zo5WBXE4tFdCyVMMXMTEMZXVlS6Xec2T4e0tHsRcEuWshcJ2YsNF5rUx1E8ifCq6Z+ZP7qdCeu/aTwFd53l16/o0NOw6O3dLavP4Hbi4RdmuDk6DoYaninC0+o4uZjbJ7Rxeu0/FbuFg+q7DVS6fQe0rZ6NDGUNNU6DEqOaLTicKnYZMnBWruljQxoaS3dZhocDge0bSTyOvdAbG5hxe2xji7E/L55xX13wWNDi6HCekcFxfCPGxY0MXC+s7afWaMdDyjyr+o8Rudm/NabOZvdl274zH4f5XK9z6On1Pe/K5TdPAslg77BjuO6Y3eO7GqvOPG/stknp1leyvLL0Z7bl9I4noMvLkzytLhWYzrOZzLXCORe028rORzOg4N/L0HlMOQ3Pgmnbb6KczlabORpu980q37TBqRu0/p3PO6234Bl03Ynuz+9W7gnsEcmvYaYY3aMYY0wx3pYd+ujsXauWdaY5Xkbtl23fPzFHiDB/QMo0yFjBllYxTQYYyxkrwn7JufwJ/PfgJ+C83X69ni6zvXcnyXabv0ncbLwsceS+RNlyN2mnneJtX0ngYO0+e+0+UnA+Wch3ji8hj5an4h+i6XBySU4n+R0roVcbw5yvHrmr4Yw8Y7x6c+9POPYHI5HI5HI5HI5HGXGww4nE4nrVyOR8XeqPEO7PLOiukYa3Novk5hV4cdtYZLI93e+uxff2jRo0aNGjRo0aNG1bVtW1dy3m83m8+tQ5ZzHw3nObwOu8La9Rc1dtkdS8A3eTk823tnktXWlxN6Oixe06zrN70Isd9jiOgZFq9yfkPqP/SLhN2Myl8jDM43bl1nbcb4cO57jlh8Jow6pzXZdL4dyODTuuhu77FyO27DdwdRxmvO+O+3N2+BdqyTwLHVczDVY4UPE4O66/ZO2cx1LFzVdSXtF7G4HMbrauOHRw6c8FdZ5m9fHZHYZXfTlZquyynSyTTKke6vcffSD9pzPA/G7n7jxPmuhc1DHMynPMrGL6AdewYmwu5ko+UUyTwrMv27rPH1v1nGqd87+p6N6LU8k3NEng53xXyHS97+44OSg/sy/hn+Se6yfYNjW0/uTgP+PvWYzLMmjhcLB/gGpri6H83/84eUXWT6T9Hsv7785z/7z4icpW+zfXypuR7rx/gMdZb1/wC678pcs8/2a3mDitGHxl9mfPlll5MafWWqxk/eYuTDgcNMzDGWLWvsuglNxs53GtN6uWpktlW1tZZYcuinMMWmnNnJydze3b2Y1McBxrBkXw799izLMZZYyy0TkbsGM4p03S2uVu5s/XXUdSdec6smVxZYYGpVmT8A+8ajuEyV5FatkvVru2x6uxGXXbH4A+jvgP4GMYy3iPLXzq/6z65+E005ey+cwMZD3fZcqc6xpjTFjQ0P3U+e++cPYmTIwj0nrK5NPTfl3WvpfLtXDcb2HQMudYOxFXQBor4L4T6vrOauFctYXJQ++NUWmJe5bmx1jDiZS1dTqWxo4GR8jm3fttpmPHppk9PEyv4/y8/sO07XacOmcqc0x2Vi9BvNJvN5oW8x4mOsydpidRxMYJPx06m1bqPzq9KtK8sxXNXFodD/+MYYaJTLwOhc9brCsV18oOR1i4tXChyTkq4lf4y1Ke+9axjDHqs1mfBbMXuP4Hzi+X7t8vzv7bHerrUPgPCxhjre4fXdfLNtNM+Jd+Zdh8xd8wP87uNPoPgv4W7/5P2BuxfsMabNnMnza+54Pdi5U671GPZY8CehX8Voeoo7FHpkeEc6715FwHZrIrUrHaviPUbPZHND+IhczrP6FcYvhOZ0Di/ETt0OI+YwNWR9r7tpf6WDeZKZDB1+z2IthOl1mPyb5FluvEx9h9d0NnM0Y1XPFkWIsk1WotJ0PBMmkvjvQTd0e71tfeV+8r8lQ/tpzpsmxJ+InrI/dj2UajUajVTUajatRqNRtGo1Go1Go4wjeMpZFMVV9CHbofPraLsJ3JpWV2XOoanCuFky4y3PPNxucK2uKC1Lbdb1eo+m5XomN6HfeZsabHLHRX/K+offtNGGmHWctcVcG44MdSqsOLY9VzX+Zxfxn2HPdWTpzWvkrtJ8M5zorrKcquRytJ5N5DZmcaW02l76nWO+BqPXm1A2Ry/0q71dH/mqrqeFjkYxjEXtsX8qubTk67rGycyqsdm4tZx5D6D5hhi0waaWmiaMP81Yjii5qxPlPuU/GfTL1Y5E6Jyfiq63qTa39A4J0sOGDgO9WF9bOXl0XfPRbsY2bPNKPy1YrFYrFYmRhhlTIyMjJWJYZHXuCXI8OoXsvfljGLFicNifpp2XunoPiG1wtx3p1Tah+/DD66OnVtVXP9rKbVxOnL0tR/rHtqB5UDErUVcl11D4qqvjpOcxX7armUNJB3LpW6bxVvD08e8h3odKKvyCFZBdSh2FVcST9xV3n3T8t1j7Kr9qgrqXg+13Pt5U7JCvFXVIV1YG5lRhkVYZJYYDDD4KOIMoHCp26WS8GB7uBh2zIdgq/PKyInjV2STShuoapUdCpX1yTwqq/z1VvET7Kh5nVPkO8YyxjLt2MaaMmWTLQvx3qnzltnXW0p2jxgbEtSny/Osv8Y9pLMXYoHVPAhkVdWVeODhR6q9/Sxe2liwwZWMVvFXfRkeIDxAePUPIrdJ4ey6yquzH+PD/bUOWAu05qVHtFd8rrKHSoeNIOUqrYr3FXyToqfYJgwmJdKpXXOwYYegNNGMzfZPp/t3t/DVs4zjNTN61rRqaWaa4NYbRjTa0tWwy2Y2tGN8ZO8ofNKq4j9SL7I+cSm4/6ovLV5HNXLI0jJidwrtk6ynCaP6Z++GjRlWS3tLeW129Mi9evxU9mtz6s5J3Z7M2ngTgnKvmpomxpaLCzPfmx0JWE+m3NLDDGOX47RctdYYNK5jakdqLkRlI39n590T5zctGSwwZZDJj6kW8XSi6ot2MmWWJ0DUT3nuvebBudScjZ79g8cWJ8av0k+/bE5WKd5MdbFpbDVMxu1DVMmtNZGJvq1mtRbn6M+g/kP0FwDwr7quZs7xosNGpbscyxhhd9TyJyFwbLcxlTasg75vW7TsV5K7ji44XPMMrdoj+Y3rT0Hie62nlYV/pwczzOmdLqLhYkzGMzCZWGMQzGMSsZYY6Di1t4nlJ+Em63mJxrVLxPbYxNEdgc1dU2iOKyoYYWjNrEeHTYybVk0atSa7ehuwsWMWTqn1TrnS6hYsi71d1+s+k+ic70e20fzE/VaTdxT9ZtU4GIXdeNx3X77guYYfpHeTQjaMX6brOu4OY4K7Y2d9mbHarI5ox3p4GpJ2Vd/Tst60f7j999pppjR+Q/Qf8J/VaORs3cji7FfFuN61+ui9s8hix1OCh5KGVV23BPXvZfz3CLyHpix+exi8z/KnCnosY2eunor+cxyPO/xJ0vKey9OvE9VjqaYu0x3Z3jd6o2b1T12D+F8l232lwaaacD5LE8LBxu7WTlbWraWpew8Xexjel3E+wWD4APITdNqR8F3R3T0lunCQ4GaE9R37DxeCYfcHi4xci5ovKfxVs55y2hf+65E/Xdp6jR5nrebTmi5incpkyOjs50JvrZwstbbW6kfuuQw+2mykf/EXNFzxfKTrxew929TR6bWnGL//F3JFOFCQT3K4lQ"
28
+
29
+ kernels = Kernel(
30
+ bz2.decompress(base64.b64decode(quantization_code)),
31
+ [
32
+ "int4WeightCompression",
33
+ "int4WeightExtractionFloat",
34
+ "int4WeightExtractionHalf",
35
+ "int8WeightExtractionFloat",
36
+ "int8WeightExtractionHalf",
37
+ ],
38
+ )
39
+ except Exception as exception:
40
+ kernels = None
41
+ logger.warning("Failed to load cpm_kernels:" + str(exception))
42
+
43
+
44
+ class W8A16Linear(torch.autograd.Function):
45
+ @staticmethod
46
+ def forward(ctx, inp: torch.Tensor, quant_w: torch.Tensor, scale_w: torch.Tensor, weight_bit_width):
47
+ ctx.inp_shape = inp.size()
48
+ ctx.weight_bit_width = weight_bit_width
49
+ out_features = quant_w.size(0)
50
+ inp = inp.contiguous().view(-1, inp.size(-1))
51
+ weight = extract_weight_to_half(quant_w, scale_w, weight_bit_width)
52
+ ctx.weight_shape = weight.size()
53
+ output = inp.mm(weight.t())
54
+ ctx.save_for_backward(inp, quant_w, scale_w)
55
+ return output.view(*(ctx.inp_shape[:-1] + (out_features,)))
56
+
57
+ @staticmethod
58
+ def backward(ctx, grad_output: torch.Tensor):
59
+ inp, quant_w, scale_w = ctx.saved_tensors
60
+ weight = extract_weight_to_half(quant_w, scale_w, ctx.weight_bit_width)
61
+ grad_output = grad_output.contiguous().view(-1, weight.size(0))
62
+ grad_input = grad_output.mm(weight)
63
+ grad_weight = grad_output.t().mm(inp)
64
+ return grad_input.view(ctx.inp_shape), grad_weight.view(ctx.weight_shape), None, None
65
+
66
+
67
+ def compress_int4_weight(weight: torch.Tensor): # (n, m)
68
+ with torch.cuda.device(weight.device):
69
+ n, m = weight.size(0), weight.size(1)
70
+ assert m % 2 == 0
71
+ m = m // 2
72
+ out = torch.empty(n, m, dtype=torch.int8, device="cuda")
73
+ stream = torch.cuda.current_stream()
74
+
75
+ gridDim = (n, 1, 1)
76
+ blockDim = (min(round_up(m, 32), 1024), 1, 1)
77
+
78
+ kernels.int4WeightCompression(
79
+ gridDim,
80
+ blockDim,
81
+ 0,
82
+ stream,
83
+ [ctypes.c_void_p(weight.data_ptr()), ctypes.c_void_p(out.data_ptr()), ctypes.c_int32(n), ctypes.c_int32(m)],
84
+ )
85
+ return out
86
+
87
+
88
+ def extract_weight_to_half(weight: torch.Tensor, scale_list: torch.Tensor, source_bit_width: int):
89
+ assert scale_list.dtype in [torch.half, torch.bfloat16]
90
+ assert weight.dtype in [torch.int8]
91
+ if source_bit_width == 8:
92
+ return weight.to(scale_list.dtype) * scale_list[:, None]
93
+ elif source_bit_width == 4:
94
+ func = (
95
+ kernels.int4WeightExtractionHalf if scale_list.dtype == torch.half else kernels.int4WeightExtractionBFloat16
96
+ )
97
+ else:
98
+ assert False, "Unsupported bit-width"
99
+
100
+ with torch.cuda.device(weight.device):
101
+ n, m = weight.size(0), weight.size(1)
102
+ out = torch.empty(n, m * (8 // source_bit_width), dtype=scale_list.dtype, device="cuda")
103
+ stream = torch.cuda.current_stream()
104
+
105
+ gridDim = (n, 1, 1)
106
+ blockDim = (min(round_up(m, 32), 1024), 1, 1)
107
+
108
+ func(
109
+ gridDim,
110
+ blockDim,
111
+ 0,
112
+ stream,
113
+ [
114
+ ctypes.c_void_p(weight.data_ptr()),
115
+ ctypes.c_void_p(scale_list.data_ptr()),
116
+ ctypes.c_void_p(out.data_ptr()),
117
+ ctypes.c_int32(n),
118
+ ctypes.c_int32(m),
119
+ ],
120
+ )
121
+ return out
122
+
123
+
124
+ class QuantizedLinear(torch.nn.Module):
125
+ def __init__(self, weight_bit_width: int, weight, bias=None, device="cpu", dtype=None, empty_init=False, *args,
126
+ **kwargs):
127
+ super().__init__()
128
+ self.weight_bit_width = weight_bit_width
129
+
130
+ shape = weight.shape
131
+
132
+ if weight is None or empty_init:
133
+ self.weight = torch.empty(shape[0], shape[1] * weight_bit_width // 8, dtype=torch.int8, device=device)
134
+ self.weight_scale = torch.empty(shape[0], dtype=dtype, device=device)
135
+ else:
136
+ self.weight_scale = weight.abs().max(dim=-1).values / ((2 ** (weight_bit_width - 1)) - 1)
137
+ self.weight = torch.round(weight / self.weight_scale[:, None]).to(torch.int8)
138
+ if weight_bit_width == 4:
139
+ self.weight = compress_int4_weight(self.weight)
140
+
141
+ self.weight = Parameter(self.weight.to(device), requires_grad=False)
142
+ self.weight_scale = Parameter(self.weight_scale.to(device), requires_grad=False)
143
+ self.bias = Parameter(bias.to(device), requires_grad=False) if bias is not None else None
144
+
145
+ def forward(self, input):
146
+ output = W8A16Linear.apply(input, self.weight, self.weight_scale, self.weight_bit_width)
147
+ if self.bias is not None:
148
+ output = output + self.bias
149
+ return output
150
+
151
+
152
+ def quantize(model, weight_bit_width, empty_init=False, device=None):
153
+ """Replace fp16 linear with quantized linear"""
154
+ for layer in model.layers:
155
+ layer.self_attention.query_key_value = QuantizedLinear(
156
+ weight_bit_width=weight_bit_width,
157
+ weight=layer.self_attention.query_key_value.weight.to(torch.cuda.current_device()),
158
+ bias=layer.self_attention.query_key_value.bias,
159
+ dtype=layer.self_attention.query_key_value.weight.dtype,
160
+ device=layer.self_attention.query_key_value.weight.device if device is None else device,
161
+ empty_init=empty_init
162
+ )
163
+ layer.self_attention.dense = QuantizedLinear(
164
+ weight_bit_width=weight_bit_width,
165
+ weight=layer.self_attention.dense.weight.to(torch.cuda.current_device()),
166
+ bias=layer.self_attention.dense.bias,
167
+ dtype=layer.self_attention.dense.weight.dtype,
168
+ device=layer.self_attention.dense.weight.device if device is None else device,
169
+ empty_init=empty_init
170
+ )
171
+ layer.mlp.dense_h_to_4h = QuantizedLinear(
172
+ weight_bit_width=weight_bit_width,
173
+ weight=layer.mlp.dense_h_to_4h.weight.to(torch.cuda.current_device()),
174
+ bias=layer.mlp.dense_h_to_4h.bias,
175
+ dtype=layer.mlp.dense_h_to_4h.weight.dtype,
176
+ device=layer.mlp.dense_h_to_4h.weight.device if device is None else device,
177
+ empty_init=empty_init
178
+ )
179
+ layer.mlp.dense_4h_to_h = QuantizedLinear(
180
+ weight_bit_width=weight_bit_width,
181
+ weight=layer.mlp.dense_4h_to_h.weight.to(torch.cuda.current_device()),
182
+ bias=layer.mlp.dense_4h_to_h.bias,
183
+ dtype=layer.mlp.dense_4h_to_h.weight.dtype,
184
+ device=layer.mlp.dense_4h_to_h.weight.device if device is None else device,
185
+ empty_init=empty_init
186
+ )
187
+
188
+ return model
checkpoint-1000/rng_state_0.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6bd8bdb59e9192d49197cb061374fd4c50d6d8537f5a113b12bf4b45127d57bc
3
+ size 17655
checkpoint-1000/rng_state_1.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c4ee77112b679d1d95da2ae8da1e21a46c05310a444fab01332534849afbf5fe
3
+ size 17655
checkpoint-1000/rng_state_2.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:298ba88fe1fab9a0075aa69e702e2cd646f3213bb45d4c7361f2177c6b6f34d1
3
+ size 17655
checkpoint-1000/rng_state_3.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:480b2689eb2712c09ad603537b3b042e208e7574e1e7b0d6bd721ac2ba83f095
3
+ size 17655
checkpoint-1000/scheduler.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:a7de62a4100c28c74ddc336f94c6734017a18fcaf4a289d3a4f1ef9794bb24a1
3
+ size 627
checkpoint-1000/special_tokens_map.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {}
checkpoint-1000/tokenization_chatglm.py ADDED
@@ -0,0 +1,253 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import torch
3
+ from typing import List, Optional, Union, Dict
4
+ from sentencepiece import SentencePieceProcessor
5
+ from transformers import PreTrainedTokenizer
6
+ from transformers.utils import logging, PaddingStrategy
7
+ from transformers.tokenization_utils_base import EncodedInput, BatchEncoding
8
+
9
+
10
+ class SPTokenizer:
11
+ def __init__(self, model_path: str):
12
+ # reload tokenizer
13
+ assert os.path.isfile(model_path), model_path
14
+ self.sp_model = SentencePieceProcessor(model_file=model_path)
15
+
16
+ # BOS / EOS token IDs
17
+ self.n_words: int = self.sp_model.vocab_size()
18
+ self.bos_id: int = self.sp_model.bos_id()
19
+ self.eos_id: int = self.sp_model.eos_id()
20
+ self.pad_id: int = self.sp_model.unk_id()
21
+ assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()
22
+
23
+ special_tokens = ["[MASK]", "[gMASK]", "[sMASK]", "sop", "eop"]
24
+ self.special_tokens = {}
25
+ self.index_special_tokens = {}
26
+ for token in special_tokens:
27
+ self.special_tokens[token] = self.n_words
28
+ self.index_special_tokens[self.n_words] = token
29
+ self.n_words += 1
30
+
31
+ def tokenize(self, s: str):
32
+ return self.sp_model.EncodeAsPieces(s)
33
+
34
+ def encode(self, s: str, bos: bool = False, eos: bool = False) -> List[int]:
35
+ assert type(s) is str
36
+ t = self.sp_model.encode(s)
37
+ if bos:
38
+ t = [self.bos_id] + t
39
+ if eos:
40
+ t = t + [self.eos_id]
41
+ return t
42
+
43
+ def decode(self, t: List[int]) -> str:
44
+ return self.sp_model.decode(t)
45
+
46
+ def decode_tokens(self, tokens: List[str]) -> str:
47
+ text = self.sp_model.DecodePieces(tokens)
48
+ return text
49
+
50
+ def convert_token_to_id(self, token):
51
+ """ Converts a token (str) in an id using the vocab. """
52
+ if token in self.special_tokens:
53
+ return self.special_tokens[token]
54
+ return self.sp_model.PieceToId(token)
55
+
56
+ def convert_id_to_token(self, index):
57
+ """Converts an index (integer) in a token (str) using the vocab."""
58
+ if index in self.index_special_tokens or index in [self.eos_id, self.bos_id, self.pad_id] or index < 0:
59
+ return ""
60
+ return self.sp_model.IdToPiece(index)
61
+
62
+
63
+ class ChatGLMTokenizer(PreTrainedTokenizer):
64
+ vocab_files_names = {"vocab_file": "tokenizer.model"}
65
+
66
+ model_input_names = ["input_ids", "attention_mask", "position_ids"]
67
+
68
+ def __init__(self, vocab_file, padding_side="left", **kwargs):
69
+ super().__init__(padding_side=padding_side, **kwargs)
70
+ self.name = "GLMTokenizer"
71
+
72
+ self.vocab_file = vocab_file
73
+ self.tokenizer = SPTokenizer(vocab_file)
74
+ self.special_tokens = {
75
+ "<bos>": self.tokenizer.bos_id,
76
+ "<eos>": self.tokenizer.eos_id,
77
+ "<pad>": self.tokenizer.pad_id
78
+ }
79
+
80
+ def get_command(self, token):
81
+ if token in self.special_tokens:
82
+ return self.special_tokens[token]
83
+ assert token in self.tokenizer.special_tokens, f"{token} is not a special token for {self.name}"
84
+ return self.tokenizer.special_tokens[token]
85
+
86
+ @property
87
+ def pad_token(self) -> str:
88
+ return "<unk>"
89
+
90
+ @property
91
+ def pad_token_id(self):
92
+ return self.get_command("<pad>")
93
+
94
+ @property
95
+ def eos_token(self) -> str:
96
+ return "</s>"
97
+
98
+ @property
99
+ def eos_token_id(self):
100
+ return self.get_command("<eos>")
101
+
102
+ @property
103
+ def vocab_size(self):
104
+ return self.tokenizer.n_words
105
+
106
+ def get_vocab(self):
107
+ """ Returns vocab as a dict """
108
+ vocab = {self._convert_id_to_token(i): i for i in range(self.vocab_size)}
109
+ vocab.update(self.added_tokens_encoder)
110
+ return vocab
111
+
112
+ def _tokenize(self, text, **kwargs):
113
+ return self.tokenizer.tokenize(text)
114
+
115
+ def _convert_token_to_id(self, token):
116
+ """ Converts a token (str) in an id using the vocab. """
117
+ return self.tokenizer.convert_token_to_id(token)
118
+
119
+ def _convert_id_to_token(self, index):
120
+ """Converts an index (integer) in a token (str) using the vocab."""
121
+ return self.tokenizer.convert_id_to_token(index)
122
+
123
+ def convert_tokens_to_string(self, tokens: List[str]) -> str:
124
+ return self.tokenizer.decode_tokens(tokens)
125
+
126
+ def save_vocabulary(self, save_directory, filename_prefix=None):
127
+ """
128
+ Save the vocabulary and special tokens file to a directory.
129
+
130
+ Args:
131
+ save_directory (`str`):
132
+ The directory in which to save the vocabulary.
133
+ filename_prefix (`str`, *optional*):
134
+ An optional prefix to add to the named of the saved files.
135
+
136
+ Returns:
137
+ `Tuple(str)`: Paths to the files saved.
138
+ """
139
+ if os.path.isdir(save_directory):
140
+ vocab_file = os.path.join(
141
+ save_directory, self.vocab_files_names["vocab_file"]
142
+ )
143
+ else:
144
+ vocab_file = save_directory
145
+
146
+ with open(self.vocab_file, 'rb') as fin:
147
+ proto_str = fin.read()
148
+
149
+ with open(vocab_file, "wb") as writer:
150
+ writer.write(proto_str)
151
+
152
+ return (vocab_file,)
153
+
154
+ def get_prefix_tokens(self):
155
+ prefix_tokens = [self.get_command("[gMASK]"), self.get_command("sop")]
156
+ return prefix_tokens
157
+
158
+ def build_prompt(self, query, history=None):
159
+ if history is None:
160
+ history = []
161
+ prompt = ""
162
+ for i, (old_query, response) in enumerate(history):
163
+ prompt += "[Round {}]\n\n问:{}\n\n答:{}\n\n".format(i + 1, old_query, response)
164
+ prompt += "[Round {}]\n\n问:{}\n\n答:".format(len(history) + 1, query)
165
+ return prompt
166
+
167
+ def build_inputs_with_special_tokens(
168
+ self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
169
+ ) -> List[int]:
170
+ """
171
+ Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
172
+ adding special tokens. A BERT sequence has the following format:
173
+
174
+ - single sequence: `[CLS] X [SEP]`
175
+ - pair of sequences: `[CLS] A [SEP] B [SEP]`
176
+
177
+ Args:
178
+ token_ids_0 (`List[int]`):
179
+ List of IDs to which the special tokens will be added.
180
+ token_ids_1 (`List[int]`, *optional*):
181
+ Optional second list of IDs for sequence pairs.
182
+
183
+ Returns:
184
+ `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
185
+ """
186
+ prefix_tokens = self.get_prefix_tokens()
187
+ token_ids_0 = prefix_tokens + token_ids_0
188
+ if token_ids_1 is not None:
189
+ token_ids_0 = token_ids_0 + token_ids_1 + [self.get_command("<eos>")]
190
+ return token_ids_0
191
+
192
+ def _pad(
193
+ self,
194
+ encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding],
195
+ max_length: Optional[int] = None,
196
+ padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
197
+ pad_to_multiple_of: Optional[int] = None,
198
+ return_attention_mask: Optional[bool] = None,
199
+ ) -> dict:
200
+ """
201
+ Pad encoded inputs (on left/right and up to predefined length or max length in the batch)
202
+
203
+ Args:
204
+ encoded_inputs:
205
+ Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`).
206
+ max_length: maximum length of the returned list and optionally padding length (see below).
207
+ Will truncate by taking into account the special tokens.
208
+ padding_strategy: PaddingStrategy to use for padding.
209
+
210
+ - PaddingStrategy.LONGEST Pad to the longest sequence in the batch
211
+ - PaddingStrategy.MAX_LENGTH: Pad to the max length (default)
212
+ - PaddingStrategy.DO_NOT_PAD: Do not pad
213
+ The tokenizer padding sides are defined in self.padding_side:
214
+
215
+ - 'left': pads on the left of the sequences
216
+ - 'right': pads on the right of the sequences
217
+ pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value.
218
+ This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability
219
+ `>= 7.5` (Volta).
220
+ return_attention_mask:
221
+ (optional) Set to False to avoid returning attention mask (default: set to model specifics)
222
+ """
223
+ # Load from model defaults
224
+ assert self.padding_side == "left"
225
+
226
+ required_input = encoded_inputs[self.model_input_names[0]]
227
+ seq_length = len(required_input)
228
+
229
+ if padding_strategy == PaddingStrategy.LONGEST:
230
+ max_length = len(required_input)
231
+
232
+ if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
233
+ max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of
234
+
235
+ needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(required_input) != max_length
236
+
237
+ # Initialize attention mask if not present.
238
+ if "attention_mask" not in encoded_inputs:
239
+ encoded_inputs["attention_mask"] = [1] * seq_length
240
+
241
+ if "position_ids" not in encoded_inputs:
242
+ encoded_inputs["position_ids"] = list(range(seq_length))
243
+
244
+ if needs_to_be_padded:
245
+ difference = max_length - len(required_input)
246
+
247
+ if "attention_mask" in encoded_inputs:
248
+ encoded_inputs["attention_mask"] = [0] * difference + encoded_inputs["attention_mask"]
249
+ if "position_ids" in encoded_inputs:
250
+ encoded_inputs["position_ids"] = [0] * difference + encoded_inputs["position_ids"]
251
+ encoded_inputs[self.model_input_names[0]] = [self.pad_token_id] * difference + required_input
252
+
253
+ return encoded_inputs
checkpoint-1000/tokenizer.model ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e7dc4c393423b76e4373e5157ddc34803a0189ba96b21ddbb40269d31468a6f2
3
+ size 1018370
checkpoint-1000/tokenizer_config.json ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "auto_map": {
3
+ "AutoTokenizer": [
4
+ "tokenization_chatglm.ChatGLMTokenizer",
5
+ null
6
+ ]
7
+ },
8
+ "clean_up_tokenization_spaces": true,
9
+ "do_lower_case": false,
10
+ "model_max_length": 1000000000000000019884624838656,
11
+ "padding_side": "left",
12
+ "remove_space": false,
13
+ "tokenizer_class": "ChatGLMTokenizer"
14
+ }
checkpoint-1000/trainer_state.json ADDED
@@ -0,0 +1,616 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "best_metric": null,
3
+ "best_model_checkpoint": null,
4
+ "epoch": 1.3767783386874712,
5
+ "global_step": 1000,
6
+ "is_hyper_param_search": false,
7
+ "is_local_process_zero": true,
8
+ "is_world_process_zero": true,
9
+ "log_history": [
10
+ {
11
+ "epoch": 0.01,
12
+ "learning_rate": 0.00019999739928881725,
13
+ "loss": 2.1906,
14
+ "step": 10
15
+ },
16
+ {
17
+ "epoch": 0.03,
18
+ "learning_rate": 0.00019998959729054295,
19
+ "loss": 1.8919,
20
+ "step": 20
21
+ },
22
+ {
23
+ "epoch": 0.04,
24
+ "learning_rate": 0.00019997659441099206,
25
+ "loss": 1.7855,
26
+ "step": 30
27
+ },
28
+ {
29
+ "epoch": 0.06,
30
+ "learning_rate": 0.00019995839132649917,
31
+ "loss": 1.7235,
32
+ "step": 40
33
+ },
34
+ {
35
+ "epoch": 0.07,
36
+ "learning_rate": 0.0001999349889838836,
37
+ "loss": 1.7027,
38
+ "step": 50
39
+ },
40
+ {
41
+ "epoch": 0.08,
42
+ "learning_rate": 0.00019990638860040006,
43
+ "loss": 1.7118,
44
+ "step": 60
45
+ },
46
+ {
47
+ "epoch": 0.1,
48
+ "learning_rate": 0.00019987259166367533,
49
+ "loss": 1.689,
50
+ "step": 70
51
+ },
52
+ {
53
+ "epoch": 0.11,
54
+ "learning_rate": 0.00019983359993163078,
55
+ "loss": 1.6845,
56
+ "step": 80
57
+ },
58
+ {
59
+ "epoch": 0.12,
60
+ "learning_rate": 0.0001997894154323911,
61
+ "loss": 1.6655,
62
+ "step": 90
63
+ },
64
+ {
65
+ "epoch": 0.14,
66
+ "learning_rate": 0.0001997400404641787,
67
+ "loss": 1.6518,
68
+ "step": 100
69
+ },
70
+ {
71
+ "epoch": 0.15,
72
+ "learning_rate": 0.00019968547759519425,
73
+ "loss": 1.6571,
74
+ "step": 110
75
+ },
76
+ {
77
+ "epoch": 0.17,
78
+ "learning_rate": 0.000199625729663483,
79
+ "loss": 1.6559,
80
+ "step": 120
81
+ },
82
+ {
83
+ "epoch": 0.18,
84
+ "learning_rate": 0.00019956079977678722,
85
+ "loss": 1.6417,
86
+ "step": 130
87
+ },
88
+ {
89
+ "epoch": 0.19,
90
+ "learning_rate": 0.0001994906913123846,
91
+ "loss": 1.6313,
92
+ "step": 140
93
+ },
94
+ {
95
+ "epoch": 0.21,
96
+ "learning_rate": 0.00019941540791691245,
97
+ "loss": 1.6243,
98
+ "step": 150
99
+ },
100
+ {
101
+ "epoch": 0.22,
102
+ "learning_rate": 0.00019933495350617813,
103
+ "loss": 1.6189,
104
+ "step": 160
105
+ },
106
+ {
107
+ "epoch": 0.23,
108
+ "learning_rate": 0.0001992493322649554,
109
+ "loss": 1.628,
110
+ "step": 170
111
+ },
112
+ {
113
+ "epoch": 0.25,
114
+ "learning_rate": 0.00019915854864676664,
115
+ "loss": 1.6149,
116
+ "step": 180
117
+ },
118
+ {
119
+ "epoch": 0.26,
120
+ "learning_rate": 0.00019906260737365122,
121
+ "loss": 1.6032,
122
+ "step": 190
123
+ },
124
+ {
125
+ "epoch": 0.28,
126
+ "learning_rate": 0.00019896151343592008,
127
+ "loss": 1.626,
128
+ "step": 200
129
+ },
130
+ {
131
+ "epoch": 0.29,
132
+ "learning_rate": 0.0001988552720918958,
133
+ "loss": 1.6369,
134
+ "step": 210
135
+ },
136
+ {
137
+ "epoch": 0.3,
138
+ "learning_rate": 0.00019874388886763944,
139
+ "loss": 1.5757,
140
+ "step": 220
141
+ },
142
+ {
143
+ "epoch": 0.32,
144
+ "learning_rate": 0.00019862736955666296,
145
+ "loss": 1.5873,
146
+ "step": 230
147
+ },
148
+ {
149
+ "epoch": 0.33,
150
+ "learning_rate": 0.00019850572021962788,
151
+ "loss": 1.6041,
152
+ "step": 240
153
+ },
154
+ {
155
+ "epoch": 0.34,
156
+ "learning_rate": 0.00019837894718402997,
157
+ "loss": 1.5905,
158
+ "step": 250
159
+ },
160
+ {
161
+ "epoch": 0.36,
162
+ "learning_rate": 0.00019824705704387028,
163
+ "loss": 1.6131,
164
+ "step": 260
165
+ },
166
+ {
167
+ "epoch": 0.37,
168
+ "learning_rate": 0.00019811005665931205,
169
+ "loss": 1.6045,
170
+ "step": 270
171
+ },
172
+ {
173
+ "epoch": 0.39,
174
+ "learning_rate": 0.00019796795315632395,
175
+ "loss": 1.5942,
176
+ "step": 280
177
+ },
178
+ {
179
+ "epoch": 0.4,
180
+ "learning_rate": 0.00019782075392630935,
181
+ "loss": 1.5768,
182
+ "step": 290
183
+ },
184
+ {
185
+ "epoch": 0.41,
186
+ "learning_rate": 0.00019766846662572191,
187
+ "loss": 1.5967,
188
+ "step": 300
189
+ },
190
+ {
191
+ "epoch": 0.43,
192
+ "learning_rate": 0.00019751109917566737,
193
+ "loss": 1.5606,
194
+ "step": 310
195
+ },
196
+ {
197
+ "epoch": 0.44,
198
+ "learning_rate": 0.00019734865976149145,
199
+ "loss": 1.58,
200
+ "step": 320
201
+ },
202
+ {
203
+ "epoch": 0.45,
204
+ "learning_rate": 0.00019718115683235417,
205
+ "loss": 1.5809,
206
+ "step": 330
207
+ },
208
+ {
209
+ "epoch": 0.47,
210
+ "learning_rate": 0.00019700859910079036,
211
+ "loss": 1.5758,
212
+ "step": 340
213
+ },
214
+ {
215
+ "epoch": 0.48,
216
+ "learning_rate": 0.00019683099554225649,
217
+ "loss": 1.57,
218
+ "step": 350
219
+ },
220
+ {
221
+ "epoch": 0.5,
222
+ "learning_rate": 0.0001966483553946637,
223
+ "loss": 1.5783,
224
+ "step": 360
225
+ },
226
+ {
227
+ "epoch": 0.51,
228
+ "learning_rate": 0.00019646068815789755,
229
+ "loss": 1.5699,
230
+ "step": 370
231
+ },
232
+ {
233
+ "epoch": 0.52,
234
+ "learning_rate": 0.00019626800359332362,
235
+ "loss": 1.5767,
236
+ "step": 380
237
+ },
238
+ {
239
+ "epoch": 0.54,
240
+ "learning_rate": 0.00019607031172327996,
241
+ "loss": 1.5635,
242
+ "step": 390
243
+ },
244
+ {
245
+ "epoch": 0.55,
246
+ "learning_rate": 0.00019586762283055573,
247
+ "loss": 1.5763,
248
+ "step": 400
249
+ },
250
+ {
251
+ "epoch": 0.56,
252
+ "learning_rate": 0.0001956599474578563,
253
+ "loss": 1.5719,
254
+ "step": 410
255
+ },
256
+ {
257
+ "epoch": 0.58,
258
+ "learning_rate": 0.00019544729640725498,
259
+ "loss": 1.5712,
260
+ "step": 420
261
+ },
262
+ {
263
+ "epoch": 0.59,
264
+ "learning_rate": 0.00019522968073963106,
265
+ "loss": 1.5699,
266
+ "step": 430
267
+ },
268
+ {
269
+ "epoch": 0.61,
270
+ "learning_rate": 0.00019500711177409454,
271
+ "loss": 1.539,
272
+ "step": 440
273
+ },
274
+ {
275
+ "epoch": 0.62,
276
+ "learning_rate": 0.0001947796010873974,
277
+ "loss": 1.5515,
278
+ "step": 450
279
+ },
280
+ {
281
+ "epoch": 0.63,
282
+ "learning_rate": 0.00019454716051333135,
283
+ "loss": 1.5691,
284
+ "step": 460
285
+ },
286
+ {
287
+ "epoch": 0.65,
288
+ "learning_rate": 0.0001943098021421124,
289
+ "loss": 1.573,
290
+ "step": 470
291
+ },
292
+ {
293
+ "epoch": 0.66,
294
+ "learning_rate": 0.00019406753831975203,
295
+ "loss": 1.5532,
296
+ "step": 480
297
+ },
298
+ {
299
+ "epoch": 0.67,
300
+ "learning_rate": 0.00019382038164741477,
301
+ "loss": 1.5778,
302
+ "step": 490
303
+ },
304
+ {
305
+ "epoch": 0.69,
306
+ "learning_rate": 0.0001935683449807631,
307
+ "loss": 1.5514,
308
+ "step": 500
309
+ },
310
+ {
311
+ "epoch": 0.7,
312
+ "learning_rate": 0.00019331144142928854,
313
+ "loss": 1.5493,
314
+ "step": 510
315
+ },
316
+ {
317
+ "epoch": 0.72,
318
+ "learning_rate": 0.00019304968435562993,
319
+ "loss": 1.5698,
320
+ "step": 520
321
+ },
322
+ {
323
+ "epoch": 0.73,
324
+ "learning_rate": 0.00019278308737487822,
325
+ "loss": 1.5699,
326
+ "step": 530
327
+ },
328
+ {
329
+ "epoch": 0.74,
330
+ "learning_rate": 0.0001925116643538684,
331
+ "loss": 1.554,
332
+ "step": 540
333
+ },
334
+ {
335
+ "epoch": 0.76,
336
+ "learning_rate": 0.00019223542941045817,
337
+ "loss": 1.5527,
338
+ "step": 550
339
+ },
340
+ {
341
+ "epoch": 0.77,
342
+ "learning_rate": 0.00019195439691279363,
343
+ "loss": 1.5517,
344
+ "step": 560
345
+ },
346
+ {
347
+ "epoch": 0.78,
348
+ "learning_rate": 0.00019166858147856203,
349
+ "loss": 1.5673,
350
+ "step": 570
351
+ },
352
+ {
353
+ "epoch": 0.8,
354
+ "learning_rate": 0.00019137799797423126,
355
+ "loss": 1.5469,
356
+ "step": 580
357
+ },
358
+ {
359
+ "epoch": 0.81,
360
+ "learning_rate": 0.00019108266151427673,
361
+ "loss": 1.5499,
362
+ "step": 590
363
+ },
364
+ {
365
+ "epoch": 0.83,
366
+ "learning_rate": 0.00019078258746039507,
367
+ "loss": 1.5502,
368
+ "step": 600
369
+ },
370
+ {
371
+ "epoch": 0.84,
372
+ "learning_rate": 0.00019047779142070527,
373
+ "loss": 1.5402,
374
+ "step": 610
375
+ },
376
+ {
377
+ "epoch": 0.85,
378
+ "learning_rate": 0.0001901682892489367,
379
+ "loss": 1.5444,
380
+ "step": 620
381
+ },
382
+ {
383
+ "epoch": 0.87,
384
+ "learning_rate": 0.00018985409704360456,
385
+ "loss": 1.5524,
386
+ "step": 630
387
+ },
388
+ {
389
+ "epoch": 0.88,
390
+ "learning_rate": 0.00018953523114717245,
391
+ "loss": 1.5338,
392
+ "step": 640
393
+ },
394
+ {
395
+ "epoch": 0.89,
396
+ "learning_rate": 0.00018921170814520247,
397
+ "loss": 1.5416,
398
+ "step": 650
399
+ },
400
+ {
401
+ "epoch": 0.91,
402
+ "learning_rate": 0.00018888354486549237,
403
+ "loss": 1.5524,
404
+ "step": 660
405
+ },
406
+ {
407
+ "epoch": 0.92,
408
+ "learning_rate": 0.00018855075837720034,
409
+ "loss": 1.544,
410
+ "step": 670
411
+ },
412
+ {
413
+ "epoch": 0.94,
414
+ "learning_rate": 0.0001882133659899573,
415
+ "loss": 1.5317,
416
+ "step": 680
417
+ },
418
+ {
419
+ "epoch": 0.95,
420
+ "learning_rate": 0.0001878713852529663,
421
+ "loss": 1.5516,
422
+ "step": 690
423
+ },
424
+ {
425
+ "epoch": 0.96,
426
+ "learning_rate": 0.00018752483395408987,
427
+ "loss": 1.5479,
428
+ "step": 700
429
+ },
430
+ {
431
+ "epoch": 0.98,
432
+ "learning_rate": 0.00018717373011892474,
433
+ "loss": 1.5328,
434
+ "step": 710
435
+ },
436
+ {
437
+ "epoch": 0.99,
438
+ "learning_rate": 0.0001868180920098644,
439
+ "loss": 1.524,
440
+ "step": 720
441
+ },
442
+ {
443
+ "epoch": 1.01,
444
+ "learning_rate": 0.00018645793812514894,
445
+ "loss": 1.4712,
446
+ "step": 730
447
+ },
448
+ {
449
+ "epoch": 1.02,
450
+ "learning_rate": 0.0001860932871979031,
451
+ "loss": 1.3509,
452
+ "step": 740
453
+ },
454
+ {
455
+ "epoch": 1.03,
456
+ "learning_rate": 0.00018572415819516174,
457
+ "loss": 1.3227,
458
+ "step": 750
459
+ },
460
+ {
461
+ "epoch": 1.05,
462
+ "learning_rate": 0.00018535057031688335,
463
+ "loss": 1.3216,
464
+ "step": 760
465
+ },
466
+ {
467
+ "epoch": 1.06,
468
+ "learning_rate": 0.00018497254299495146,
469
+ "loss": 1.3022,
470
+ "step": 770
471
+ },
472
+ {
473
+ "epoch": 1.07,
474
+ "learning_rate": 0.00018459009589216364,
475
+ "loss": 1.308,
476
+ "step": 780
477
+ },
478
+ {
479
+ "epoch": 1.09,
480
+ "learning_rate": 0.00018420324890120916,
481
+ "loss": 1.3295,
482
+ "step": 790
483
+ },
484
+ {
485
+ "epoch": 1.1,
486
+ "learning_rate": 0.0001838120221436338,
487
+ "loss": 1.3249,
488
+ "step": 800
489
+ },
490
+ {
491
+ "epoch": 1.12,
492
+ "learning_rate": 0.00018341643596879367,
493
+ "loss": 1.3109,
494
+ "step": 810
495
+ },
496
+ {
497
+ "epoch": 1.13,
498
+ "learning_rate": 0.00018301651095279655,
499
+ "loss": 1.328,
500
+ "step": 820
501
+ },
502
+ {
503
+ "epoch": 1.14,
504
+ "learning_rate": 0.00018261226789743172,
505
+ "loss": 1.3215,
506
+ "step": 830
507
+ },
508
+ {
509
+ "epoch": 1.16,
510
+ "learning_rate": 0.00018220372782908777,
511
+ "loss": 1.3378,
512
+ "step": 840
513
+ },
514
+ {
515
+ "epoch": 1.17,
516
+ "learning_rate": 0.00018179091199765926,
517
+ "loss": 1.3289,
518
+ "step": 850
519
+ },
520
+ {
521
+ "epoch": 1.18,
522
+ "learning_rate": 0.00018137384187544116,
523
+ "loss": 1.3362,
524
+ "step": 860
525
+ },
526
+ {
527
+ "epoch": 1.2,
528
+ "learning_rate": 0.00018095253915601206,
529
+ "loss": 1.3263,
530
+ "step": 870
531
+ },
532
+ {
533
+ "epoch": 1.21,
534
+ "learning_rate": 0.00018052702575310588,
535
+ "loss": 1.3238,
536
+ "step": 880
537
+ },
538
+ {
539
+ "epoch": 1.23,
540
+ "learning_rate": 0.00018009732379947188,
541
+ "loss": 1.3278,
542
+ "step": 890
543
+ },
544
+ {
545
+ "epoch": 1.24,
546
+ "learning_rate": 0.0001796634556457236,
547
+ "loss": 1.3277,
548
+ "step": 900
549
+ },
550
+ {
551
+ "epoch": 1.25,
552
+ "learning_rate": 0.00017922544385917628,
553
+ "loss": 1.3332,
554
+ "step": 910
555
+ },
556
+ {
557
+ "epoch": 1.27,
558
+ "learning_rate": 0.00017878331122267284,
559
+ "loss": 1.3268,
560
+ "step": 920
561
+ },
562
+ {
563
+ "epoch": 1.28,
564
+ "learning_rate": 0.00017833708073339922,
565
+ "loss": 1.3316,
566
+ "step": 930
567
+ },
568
+ {
569
+ "epoch": 1.29,
570
+ "learning_rate": 0.00017788677560168784,
571
+ "loss": 1.337,
572
+ "step": 940
573
+ },
574
+ {
575
+ "epoch": 1.31,
576
+ "learning_rate": 0.0001774324192498105,
577
+ "loss": 1.3307,
578
+ "step": 950
579
+ },
580
+ {
581
+ "epoch": 1.32,
582
+ "learning_rate": 0.0001769740353107602,
583
+ "loss": 1.3295,
584
+ "step": 960
585
+ },
586
+ {
587
+ "epoch": 1.34,
588
+ "learning_rate": 0.0001765116476270216,
589
+ "loss": 1.3148,
590
+ "step": 970
591
+ },
592
+ {
593
+ "epoch": 1.35,
594
+ "learning_rate": 0.00017604528024933115,
595
+ "loss": 1.3269,
596
+ "step": 980
597
+ },
598
+ {
599
+ "epoch": 1.36,
600
+ "learning_rate": 0.00017557495743542585,
601
+ "loss": 1.3488,
602
+ "step": 990
603
+ },
604
+ {
605
+ "epoch": 1.38,
606
+ "learning_rate": 0.00017510070364878177,
607
+ "loss": 1.3377,
608
+ "step": 1000
609
+ }
610
+ ],
611
+ "max_steps": 4356,
612
+ "num_train_epochs": 6,
613
+ "total_flos": 1.2097274818718597e+18,
614
+ "trial_name": null,
615
+ "trial_params": null
616
+ }
checkpoint-1000/training_args.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:326a38cf0f9916b8ede07877300460617478f67a798b1bed6b7d9bf33cb7066b
3
+ size 3289
checkpoint-2000/config.json ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "/home/ubuntu/chatglm2-6b",
3
+ "add_bias_linear": false,
4
+ "add_qkv_bias": true,
5
+ "apply_query_key_layer_scaling": true,
6
+ "apply_residual_connection_post_layernorm": false,
7
+ "architectures": [
8
+ "ChatGLMForConditionalGeneration"
9
+ ],
10
+ "attention_dropout": 0.0,
11
+ "attention_softmax_in_fp32": true,
12
+ "auto_map": {
13
+ "AutoConfig": "configuration_chatglm.ChatGLMConfig",
14
+ "AutoModel": "modeling_chatglm.ChatGLMForConditionalGeneration",
15
+ "AutoModelForCausalLM": "modeling_chatglm.ChatGLMForConditionalGeneration",
16
+ "AutoModelForSeq2SeqLM": "modeling_chatglm.ChatGLMForConditionalGeneration"
17
+ },
18
+ "bias_dropout_fusion": true,
19
+ "eos_token_id": 2,
20
+ "ffn_hidden_size": 13696,
21
+ "fp32_residual_connection": false,
22
+ "hidden_dropout": 0.0,
23
+ "hidden_size": 4096,
24
+ "kv_channels": 128,
25
+ "layernorm_epsilon": 1e-05,
26
+ "model_type": "chatglm",
27
+ "multi_query_attention": true,
28
+ "multi_query_group_num": 2,
29
+ "num_attention_heads": 32,
30
+ "num_layers": 28,
31
+ "original_rope": true,
32
+ "pad_token_id": 0,
33
+ "padded_vocab_size": 65024,
34
+ "post_layer_norm": true,
35
+ "pre_seq_len": null,
36
+ "prefix_projection": false,
37
+ "quantization_bit": 0,
38
+ "rmsnorm": true,
39
+ "seq_length": 32768,
40
+ "tie_word_embeddings": false,
41
+ "torch_dtype": "float16",
42
+ "transformers_version": "4.30.2",
43
+ "use_cache": true,
44
+ "vocab_size": 65024
45
+ }
checkpoint-2000/configuration_chatglm.py ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import PretrainedConfig
2
+
3
+
4
+ class ChatGLMConfig(PretrainedConfig):
5
+ model_type = "chatglm"
6
+ def __init__(
7
+ self,
8
+ num_layers=28,
9
+ padded_vocab_size=65024,
10
+ hidden_size=4096,
11
+ ffn_hidden_size=13696,
12
+ kv_channels=128,
13
+ num_attention_heads=32,
14
+ seq_length=2048,
15
+ hidden_dropout=0.0,
16
+ attention_dropout=0.0,
17
+ layernorm_epsilon=1e-5,
18
+ rmsnorm=True,
19
+ apply_residual_connection_post_layernorm=False,
20
+ post_layer_norm=True,
21
+ add_bias_linear=False,
22
+ add_qkv_bias=False,
23
+ bias_dropout_fusion=True,
24
+ multi_query_attention=False,
25
+ multi_query_group_num=1,
26
+ apply_query_key_layer_scaling=True,
27
+ attention_softmax_in_fp32=True,
28
+ fp32_residual_connection=False,
29
+ quantization_bit=0,
30
+ pre_seq_len=None,
31
+ prefix_projection=False,
32
+ **kwargs
33
+ ):
34
+ self.num_layers = num_layers
35
+ self.vocab_size = padded_vocab_size
36
+ self.padded_vocab_size = padded_vocab_size
37
+ self.hidden_size = hidden_size
38
+ self.ffn_hidden_size = ffn_hidden_size
39
+ self.kv_channels = kv_channels
40
+ self.num_attention_heads = num_attention_heads
41
+ self.seq_length = seq_length
42
+ self.hidden_dropout = hidden_dropout
43
+ self.attention_dropout = attention_dropout
44
+ self.layernorm_epsilon = layernorm_epsilon
45
+ self.rmsnorm = rmsnorm
46
+ self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm
47
+ self.post_layer_norm = post_layer_norm
48
+ self.add_bias_linear = add_bias_linear
49
+ self.add_qkv_bias = add_qkv_bias
50
+ self.bias_dropout_fusion = bias_dropout_fusion
51
+ self.multi_query_attention = multi_query_attention
52
+ self.multi_query_group_num = multi_query_group_num
53
+ self.apply_query_key_layer_scaling = apply_query_key_layer_scaling
54
+ self.attention_softmax_in_fp32 = attention_softmax_in_fp32
55
+ self.fp32_residual_connection = fp32_residual_connection
56
+ self.quantization_bit = quantization_bit
57
+ self.pre_seq_len = pre_seq_len
58
+ self.prefix_projection = prefix_projection
59
+ super().__init__(**kwargs)
checkpoint-2000/finetuning_args.json ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "finetuning_type": "freeze",
3
+ "lora_alpha": 9.0,
4
+ "lora_dropout": 0.05,
5
+ "lora_rank": 100,
6
+ "lora_target": [
7
+ "query_key_value"
8
+ ],
9
+ "name_module_trainable": "mlp",
10
+ "num_layer_trainable": 1,
11
+ "pre_seq_len": 64,
12
+ "prefix_projection": false
13
+ }
checkpoint-2000/generation_config.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "eos_token_id": 2,
4
+ "pad_token_id": 0,
5
+ "transformers_version": "4.30.2"
6
+ }
checkpoint-2000/modeling_chatglm.py ADDED
@@ -0,0 +1,1193 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """ PyTorch ChatGLM model. """
2
+
3
+ import math
4
+ import copy
5
+ import warnings
6
+ import re
7
+ import sys
8
+
9
+ import torch
10
+ import torch.utils.checkpoint
11
+ import torch.nn.functional as F
12
+ from torch import nn
13
+ from torch.nn import CrossEntropyLoss, LayerNorm
14
+ from torch.nn.utils import skip_init
15
+ from typing import Optional, Tuple, Union, List, Callable, Dict, Any
16
+
17
+ from transformers.modeling_outputs import (
18
+ BaseModelOutputWithPast,
19
+ CausalLMOutputWithPast,
20
+ )
21
+ from transformers.modeling_utils import PreTrainedModel
22
+ from transformers.utils import logging
23
+ from transformers.generation.logits_process import LogitsProcessor
24
+ from transformers.generation.utils import LogitsProcessorList, StoppingCriteriaList, GenerationConfig, ModelOutput
25
+
26
+ from .configuration_chatglm import ChatGLMConfig
27
+
28
+ # flags required to enable jit fusion kernels
29
+
30
+ if sys.platform != 'darwin':
31
+ torch._C._jit_set_profiling_mode(False)
32
+ torch._C._jit_set_profiling_executor(False)
33
+ torch._C._jit_override_can_fuse_on_cpu(True)
34
+ torch._C._jit_override_can_fuse_on_gpu(True)
35
+
36
+ logger = logging.get_logger(__name__)
37
+
38
+ _CHECKPOINT_FOR_DOC = "THUDM/ChatGLM2-6B"
39
+ _CONFIG_FOR_DOC = "ChatGLM6BConfig"
40
+
41
+ CHATGLM_6B_PRETRAINED_MODEL_ARCHIVE_LIST = [
42
+ "THUDM/chatglm2-6b",
43
+ # See all ChatGLM models at https://huggingface.co/models?filter=chatglm
44
+ ]
45
+
46
+
47
+ def default_init(cls, *args, **kwargs):
48
+ return cls(*args, **kwargs)
49
+
50
+
51
+ class InvalidScoreLogitsProcessor(LogitsProcessor):
52
+ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
53
+ if torch.isnan(scores).any() or torch.isinf(scores).any():
54
+ scores.zero_()
55
+ scores[..., 5] = 5e4
56
+ return scores
57
+
58
+
59
+ class PrefixEncoder(torch.nn.Module):
60
+ """
61
+ The torch.nn model to encode the prefix
62
+ Input shape: (batch-size, prefix-length)
63
+ Output shape: (batch-size, prefix-length, 2*layers*hidden)
64
+ """
65
+
66
+ def __init__(self, config: ChatGLMConfig):
67
+ super().__init__()
68
+ self.prefix_projection = config.prefix_projection
69
+ if self.prefix_projection:
70
+ # Use a two-layer MLP to encode the prefix
71
+ kv_size = config.num_layers * config.kv_channels * config.multi_query_group_num * 2
72
+ self.embedding = torch.nn.Embedding(config.pre_seq_len, kv_size)
73
+ self.trans = torch.nn.Sequential(
74
+ torch.nn.Linear(kv_size, config.hidden_size),
75
+ torch.nn.Tanh(),
76
+ torch.nn.Linear(config.hidden_size, kv_size)
77
+ )
78
+ else:
79
+ self.embedding = torch.nn.Embedding(config.pre_seq_len,
80
+ config.num_layers * config.kv_channels * config.multi_query_group_num * 2)
81
+
82
+ def forward(self, prefix: torch.Tensor):
83
+ if self.prefix_projection:
84
+ prefix_tokens = self.embedding(prefix)
85
+ past_key_values = self.trans(prefix_tokens)
86
+ else:
87
+ past_key_values = self.embedding(prefix)
88
+ return past_key_values
89
+
90
+
91
+ def split_tensor_along_last_dim(
92
+ tensor: torch.Tensor,
93
+ num_partitions: int,
94
+ contiguous_split_chunks: bool = False,
95
+ ) -> List[torch.Tensor]:
96
+ """Split a tensor along its last dimension.
97
+
98
+ Arguments:
99
+ tensor: input tensor.
100
+ num_partitions: number of partitions to split the tensor
101
+ contiguous_split_chunks: If True, make each chunk contiguous
102
+ in memory.
103
+
104
+ Returns:
105
+ A list of Tensors
106
+ """
107
+ # Get the size and dimension.
108
+ last_dim = tensor.dim() - 1
109
+ last_dim_size = tensor.size()[last_dim] // num_partitions
110
+ # Split.
111
+ tensor_list = torch.split(tensor, last_dim_size, dim=last_dim)
112
+ # Note: torch.split does not create contiguous tensors by default.
113
+ if contiguous_split_chunks:
114
+ return tuple(chunk.contiguous() for chunk in tensor_list)
115
+
116
+ return tensor_list
117
+
118
+
119
+ class RotaryEmbedding(nn.Module):
120
+ def __init__(self, dim, original_impl=False, device=None, dtype=None):
121
+ super().__init__()
122
+ inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, device=device).to(dtype=dtype) / dim))
123
+ self.register_buffer("inv_freq", inv_freq)
124
+ self.dim = dim
125
+ self.original_impl = original_impl
126
+
127
+ def forward_impl(
128
+ self, seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000
129
+ ):
130
+ """Enhanced Transformer with Rotary Position Embedding.
131
+
132
+ Derived from: https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/
133
+ transformers/rope/__init__.py. MIT License:
134
+ https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license.
135
+ """
136
+ # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
137
+ theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, dtype=dtype, device=device) / n_elem))
138
+
139
+ # Create position indexes `[0, 1, ..., seq_len - 1]`
140
+ seq_idx = torch.arange(seq_len, dtype=dtype, device=device)
141
+
142
+ # Calculate the product of position index and $\theta_i$
143
+ idx_theta = torch.outer(seq_idx, theta).float()
144
+
145
+ cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1)
146
+
147
+ # this is to mimic the behaviour of complex32, else we will get different results
148
+ if dtype in (torch.float16, torch.bfloat16, torch.int8):
149
+ cache = cache.bfloat16() if dtype == torch.bfloat16 else cache.half()
150
+ return cache
151
+
152
+ def forward(self, max_seq_len, offset=0):
153
+ return self.forward_impl(
154
+ max_seq_len, self.dim, dtype=self.inv_freq.dtype, device=self.inv_freq.device
155
+ )
156
+
157
+
158
+ @torch.jit.script
159
+ def apply_rotary_pos_emb(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor:
160
+ # x: [sq, b, np, hn]
161
+ sq, b, np, hn = x.size(0), x.size(1), x.size(2), x.size(3)
162
+ rot_dim = rope_cache.shape[-2] * 2
163
+ x, x_pass = x[..., :rot_dim], x[..., rot_dim:]
164
+ # truncate to support variable sizes
165
+ rope_cache = rope_cache[:sq]
166
+ xshaped = x.reshape(sq, -1, np, rot_dim // 2, 2)
167
+ rope_cache = rope_cache.view(sq, -1, 1, xshaped.size(3), 2)
168
+ x_out2 = torch.stack(
169
+ [
170
+ xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1],
171
+ xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1],
172
+ ],
173
+ -1,
174
+ )
175
+ x_out2 = x_out2.flatten(3)
176
+ return torch.cat((x_out2, x_pass), dim=-1)
177
+
178
+
179
+ class RMSNorm(torch.nn.Module):
180
+ def __init__(self, normalized_shape, eps=1e-5, device=None, dtype=None, **kwargs):
181
+ super().__init__()
182
+ self.weight = torch.nn.Parameter(torch.empty(normalized_shape, device=device, dtype=dtype))
183
+ self.eps = eps
184
+
185
+ def forward(self, hidden_states: torch.Tensor):
186
+ input_dtype = hidden_states.dtype
187
+ variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
188
+ hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
189
+
190
+ return (self.weight * hidden_states).to(input_dtype)
191
+
192
+
193
+ class CoreAttention(torch.nn.Module):
194
+ def __init__(self, config: ChatGLMConfig, layer_number):
195
+ super(CoreAttention, self).__init__()
196
+
197
+ self.apply_query_key_layer_scaling = config.apply_query_key_layer_scaling
198
+ self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32
199
+ if self.apply_query_key_layer_scaling:
200
+ self.attention_softmax_in_fp32 = True
201
+ self.layer_number = max(1, layer_number)
202
+
203
+ projection_size = config.kv_channels * config.num_attention_heads
204
+
205
+ # Per attention head and per partition values.
206
+ self.hidden_size_per_partition = projection_size
207
+ self.hidden_size_per_attention_head = projection_size // config.num_attention_heads
208
+ self.num_attention_heads_per_partition = config.num_attention_heads
209
+
210
+ coeff = None
211
+ self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)
212
+ if self.apply_query_key_layer_scaling:
213
+ coeff = self.layer_number
214
+ self.norm_factor *= coeff
215
+ self.coeff = coeff
216
+
217
+ self.attention_dropout = torch.nn.Dropout(config.attention_dropout)
218
+
219
+ def forward(self, query_layer, key_layer, value_layer, attention_mask):
220
+ pytorch_major_version = int(torch.__version__.split('.')[0])
221
+ if pytorch_major_version >= 2:
222
+ query_layer, key_layer, value_layer = [k.permute(1, 2, 0, 3) for k in [query_layer, key_layer, value_layer]]
223
+ if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]:
224
+ context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
225
+ is_causal=True)
226
+ else:
227
+ if attention_mask is not None:
228
+ attention_mask = ~attention_mask
229
+ context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
230
+ attention_mask)
231
+ context_layer = context_layer.permute(2, 0, 1, 3)
232
+ new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
233
+ context_layer = context_layer.reshape(*new_context_layer_shape)
234
+ else:
235
+ # Raw attention scores
236
+
237
+ # [b, np, sq, sk]
238
+ output_size = (query_layer.size(1), query_layer.size(2), query_layer.size(0), key_layer.size(0))
239
+
240
+ # [sq, b, np, hn] -> [sq, b * np, hn]
241
+ query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1)
242
+ # [sk, b, np, hn] -> [sk, b * np, hn]
243
+ key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1)
244
+
245
+ # preallocting input tensor: [b * np, sq, sk]
246
+ matmul_input_buffer = torch.empty(
247
+ output_size[0] * output_size[1], output_size[2], output_size[3], dtype=query_layer.dtype,
248
+ device=query_layer.device
249
+ )
250
+
251
+ # Raw attention scores. [b * np, sq, sk]
252
+ matmul_result = torch.baddbmm(
253
+ matmul_input_buffer,
254
+ query_layer.transpose(0, 1), # [b * np, sq, hn]
255
+ key_layer.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk]
256
+ beta=0.0,
257
+ alpha=(1.0 / self.norm_factor),
258
+ )
259
+
260
+ # change view to [b, np, sq, sk]
261
+ attention_scores = matmul_result.view(*output_size)
262
+
263
+ # ===========================
264
+ # Attention probs and dropout
265
+ # ===========================
266
+
267
+ # attention scores and attention mask [b, np, sq, sk]
268
+ if self.attention_softmax_in_fp32:
269
+ attention_scores = attention_scores.float()
270
+ if self.coeff is not None:
271
+ attention_scores = attention_scores * self.coeff
272
+ if attention_mask is None and attention_scores.shape[2] == attention_scores.shape[3]:
273
+ attention_mask = torch.ones(output_size[0], 1, output_size[2], output_size[3],
274
+ device=attention_scores.device, dtype=torch.bool)
275
+ attention_mask.tril_()
276
+ attention_mask = ~attention_mask
277
+ if attention_mask is not None:
278
+ attention_scores = attention_scores.masked_fill(attention_mask, float("-inf"))
279
+ attention_probs = F.softmax(attention_scores, dim=-1)
280
+ attention_probs = attention_probs.type_as(value_layer)
281
+
282
+ # This is actually dropping out entire tokens to attend to, which might
283
+ # seem a bit unusual, but is taken from the original Transformer paper.
284
+ attention_probs = self.attention_dropout(attention_probs)
285
+ # =========================
286
+ # Context layer. [sq, b, hp]
287
+ # =========================
288
+
289
+ # value_layer -> context layer.
290
+ # [sk, b, np, hn] --> [b, np, sq, hn]
291
+
292
+ # context layer shape: [b, np, sq, hn]
293
+ output_size = (value_layer.size(1), value_layer.size(2), query_layer.size(0), value_layer.size(3))
294
+ # change view [sk, b * np, hn]
295
+ value_layer = value_layer.view(value_layer.size(0), output_size[0] * output_size[1], -1)
296
+ # change view [b * np, sq, sk]
297
+ attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
298
+ # matmul: [b * np, sq, hn]
299
+ context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1))
300
+ # change view [b, np, sq, hn]
301
+ context_layer = context_layer.view(*output_size)
302
+ # [b, np, sq, hn] --> [sq, b, np, hn]
303
+ context_layer = context_layer.permute(2, 0, 1, 3).contiguous()
304
+ # [sq, b, np, hn] --> [sq, b, hp]
305
+ new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
306
+ context_layer = context_layer.view(*new_context_layer_shape)
307
+
308
+ return context_layer
309
+
310
+
311
+ class SelfAttention(torch.nn.Module):
312
+ """Parallel self-attention layer abstract class.
313
+
314
+ Self-attention layer takes input with size [s, b, h]
315
+ and returns output of the same size.
316
+ """
317
+
318
+ def __init__(self, config: ChatGLMConfig, layer_number, device=None):
319
+ super(SelfAttention, self).__init__()
320
+ self.layer_number = max(1, layer_number)
321
+
322
+ self.projection_size = config.kv_channels * config.num_attention_heads
323
+
324
+ # Per attention head and per partition values.
325
+ self.hidden_size_per_attention_head = self.projection_size // config.num_attention_heads
326
+ self.num_attention_heads_per_partition = config.num_attention_heads
327
+
328
+ self.multi_query_attention = config.multi_query_attention
329
+ self.qkv_hidden_size = 3 * self.projection_size
330
+ if self.multi_query_attention:
331
+ self.num_multi_query_groups_per_partition = config.multi_query_group_num
332
+ self.qkv_hidden_size = (
333
+ self.projection_size + 2 * self.hidden_size_per_attention_head * config.multi_query_group_num
334
+ )
335
+ self.query_key_value = nn.Linear(config.hidden_size, self.qkv_hidden_size,
336
+ bias=config.add_bias_linear or config.add_qkv_bias,
337
+ device=device, **_config_to_kwargs(config)
338
+ )
339
+
340
+ self.core_attention = CoreAttention(config, self.layer_number)
341
+
342
+ # Output.
343
+ self.dense = nn.Linear(self.projection_size, config.hidden_size, bias=config.add_bias_linear,
344
+ device=device, **_config_to_kwargs(config)
345
+ )
346
+
347
+ def _allocate_memory(self, inference_max_sequence_len, batch_size, device=None, dtype=None):
348
+ if self.multi_query_attention:
349
+ num_attention_heads = self.num_multi_query_groups_per_partition
350
+ else:
351
+ num_attention_heads = self.num_attention_heads_per_partition
352
+ return torch.empty(
353
+ inference_max_sequence_len,
354
+ batch_size,
355
+ num_attention_heads,
356
+ self.hidden_size_per_attention_head,
357
+ dtype=dtype,
358
+ device=device,
359
+ )
360
+
361
+ def forward(
362
+ self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True
363
+ ):
364
+ # hidden_states: [sq, b, h]
365
+
366
+ # =================================================
367
+ # Pre-allocate memory for key-values for inference.
368
+ # =================================================
369
+ # =====================
370
+ # Query, Key, and Value
371
+ # =====================
372
+
373
+ # Attention heads [sq, b, h] --> [sq, b, (np * 3 * hn)]
374
+ mixed_x_layer = self.query_key_value(hidden_states)
375
+
376
+ if self.multi_query_attention:
377
+ (query_layer, key_layer, value_layer) = mixed_x_layer.split(
378
+ [
379
+ self.num_attention_heads_per_partition * self.hidden_size_per_attention_head,
380
+ self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
381
+ self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
382
+ ],
383
+ dim=-1,
384
+ )
385
+ query_layer = query_layer.view(
386
+ query_layer.size()[:-1] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)
387
+ )
388
+ key_layer = key_layer.view(
389
+ key_layer.size()[:-1] + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)
390
+ )
391
+ value_layer = value_layer.view(
392
+ value_layer.size()[:-1]
393
+ + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)
394
+ )
395
+ else:
396
+ new_tensor_shape = mixed_x_layer.size()[:-1] + \
397
+ (self.num_attention_heads_per_partition,
398
+ 3 * self.hidden_size_per_attention_head)
399
+ mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)
400
+
401
+ # [sq, b, np, 3 * hn] --> 3 [sq, b, np, hn]
402
+ (query_layer, key_layer, value_layer) = split_tensor_along_last_dim(mixed_x_layer, 3)
403
+
404
+ # apply relative positional encoding (rotary embedding)
405
+ if rotary_pos_emb is not None:
406
+ query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb)
407
+ key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb)
408
+
409
+ # adjust key and value for inference
410
+ if kv_cache is not None:
411
+ cache_k, cache_v = kv_cache
412
+ key_layer = torch.cat((cache_k, key_layer), dim=0)
413
+ value_layer = torch.cat((cache_v, value_layer), dim=0)
414
+ if use_cache:
415
+ kv_cache = (key_layer, value_layer)
416
+ else:
417
+ kv_cache = None
418
+
419
+ if self.multi_query_attention:
420
+ key_layer = key_layer.unsqueeze(-2)
421
+ key_layer = key_layer.expand(
422
+ -1, -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1
423
+ )
424
+ key_layer = key_layer.contiguous().view(
425
+ key_layer.size()[:2] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)
426
+ )
427
+ value_layer = value_layer.unsqueeze(-2)
428
+ value_layer = value_layer.expand(
429
+ -1, -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1
430
+ )
431
+ value_layer = value_layer.contiguous().view(
432
+ value_layer.size()[:2] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)
433
+ )
434
+
435
+ # ==================================
436
+ # core attention computation
437
+ # ==================================
438
+
439
+ context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask)
440
+
441
+ # =================
442
+ # Output. [sq, b, h]
443
+ # =================
444
+
445
+ output = self.dense(context_layer)
446
+
447
+ return output, kv_cache
448
+
449
+
450
+ def _config_to_kwargs(args):
451
+ common_kwargs = {
452
+ "dtype": args.torch_dtype,
453
+ }
454
+ return common_kwargs
455
+
456
+
457
+ class MLP(torch.nn.Module):
458
+ """MLP.
459
+
460
+ MLP will take the input with h hidden state, project it to 4*h
461
+ hidden dimension, perform nonlinear transformation, and project the
462
+ state back into h hidden dimension.
463
+ """
464
+
465
+ def __init__(self, config: ChatGLMConfig, device=None):
466
+ super(MLP, self).__init__()
467
+
468
+ self.add_bias = config.add_bias_linear
469
+
470
+ # Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf
471
+ self.dense_h_to_4h = nn.Linear(
472
+ config.hidden_size,
473
+ config.ffn_hidden_size * 2,
474
+ bias=self.add_bias,
475
+ device=device,
476
+ **_config_to_kwargs(config)
477
+ )
478
+
479
+ def swiglu(x):
480
+ x = torch.chunk(x, 2, dim=-1)
481
+ return F.silu(x[0]) * x[1]
482
+
483
+ self.activation_func = swiglu
484
+
485
+ # Project back to h.
486
+ self.dense_4h_to_h = nn.Linear(
487
+ config.ffn_hidden_size,
488
+ config.hidden_size,
489
+ bias=self.add_bias,
490
+ device=device,
491
+ **_config_to_kwargs(config)
492
+ )
493
+
494
+ def forward(self, hidden_states):
495
+ # [s, b, 4hp]
496
+ intermediate_parallel = self.dense_h_to_4h(hidden_states)
497
+ intermediate_parallel = self.activation_func(intermediate_parallel)
498
+ # [s, b, h]
499
+ output = self.dense_4h_to_h(intermediate_parallel)
500
+ return output
501
+
502
+
503
+ class GLMBlock(torch.nn.Module):
504
+ """A single transformer layer.
505
+
506
+ Transformer layer takes input with size [s, b, h] and returns an
507
+ output of the same size.
508
+ """
509
+
510
+ def __init__(self, config: ChatGLMConfig, layer_number, device=None):
511
+ super(GLMBlock, self).__init__()
512
+ self.layer_number = layer_number
513
+
514
+ self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm
515
+
516
+ self.fp32_residual_connection = config.fp32_residual_connection
517
+
518
+ LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm
519
+ # Layernorm on the input data.
520
+ self.input_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
521
+ dtype=config.torch_dtype)
522
+
523
+ # Self attention.
524
+ self.self_attention = SelfAttention(config, layer_number, device=device)
525
+ self.hidden_dropout = config.hidden_dropout
526
+
527
+ # Layernorm on the attention output
528
+ self.post_attention_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
529
+ dtype=config.torch_dtype)
530
+
531
+ # MLP
532
+ self.mlp = MLP(config, device=device)
533
+
534
+ def forward(
535
+ self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True,
536
+ ):
537
+ # hidden_states: [s, b, h]
538
+
539
+ # Layer norm at the beginning of the transformer layer.
540
+ layernorm_output = self.input_layernorm(hidden_states)
541
+ # Self attention.
542
+ attention_output, kv_cache = self.self_attention(
543
+ layernorm_output,
544
+ attention_mask,
545
+ rotary_pos_emb,
546
+ kv_cache=kv_cache,
547
+ use_cache=use_cache
548
+ )
549
+
550
+ # Residual connection.
551
+ if self.apply_residual_connection_post_layernorm:
552
+ residual = layernorm_output
553
+ else:
554
+ residual = hidden_states
555
+
556
+ layernorm_input = torch.nn.functional.dropout(attention_output, p=self.hidden_dropout, training=self.training)
557
+ layernorm_input = residual + layernorm_input
558
+
559
+ # Layer norm post the self attention.
560
+ layernorm_output = self.post_attention_layernorm(layernorm_input)
561
+
562
+ # MLP.
563
+ mlp_output = self.mlp(layernorm_output)
564
+
565
+ # Second residual connection.
566
+ if self.apply_residual_connection_post_layernorm:
567
+ residual = layernorm_output
568
+ else:
569
+ residual = layernorm_input
570
+
571
+ output = torch.nn.functional.dropout(mlp_output, p=self.hidden_dropout, training=self.training)
572
+ output = residual + output
573
+
574
+ return output, kv_cache
575
+
576
+
577
+ class GLMTransformer(torch.nn.Module):
578
+ """Transformer class."""
579
+
580
+ def __init__(self, config: ChatGLMConfig, device=None):
581
+ super(GLMTransformer, self).__init__()
582
+
583
+ self.fp32_residual_connection = config.fp32_residual_connection
584
+ self.post_layer_norm = config.post_layer_norm
585
+
586
+ # Number of layers.
587
+ self.num_layers = config.num_layers
588
+
589
+ # Transformer layers.
590
+ def build_layer(layer_number):
591
+ return GLMBlock(config, layer_number, device=device)
592
+
593
+ self.layers = torch.nn.ModuleList([build_layer(i + 1) for i in range(self.num_layers)])
594
+
595
+ if self.post_layer_norm:
596
+ LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm
597
+ # Final layer norm before output.
598
+ self.final_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
599
+ dtype=config.torch_dtype)
600
+
601
+ self.gradient_checkpointing = False
602
+
603
+ def _get_layer(self, layer_number):
604
+ return self.layers[layer_number]
605
+
606
+ def forward(
607
+ self, hidden_states, attention_mask, rotary_pos_emb, kv_caches=None,
608
+ use_cache: Optional[bool] = True,
609
+ output_hidden_states: Optional[bool] = False,
610
+ ):
611
+ if not kv_caches:
612
+ kv_caches = [None for _ in range(self.num_layers)]
613
+ presents = () if use_cache else None
614
+ if self.gradient_checkpointing and self.training:
615
+ if use_cache:
616
+ logger.warning_once(
617
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
618
+ )
619
+ use_cache = False
620
+
621
+ all_self_attentions = None
622
+ all_hidden_states = () if output_hidden_states else None
623
+ for index in range(self.num_layers):
624
+ if output_hidden_states:
625
+ all_hidden_states = all_hidden_states + (hidden_states,)
626
+
627
+ layer = self._get_layer(index)
628
+ if self.gradient_checkpointing and self.training:
629
+ layer_ret = torch.utils.checkpoint.checkpoint(
630
+ layer,
631
+ hidden_states,
632
+ attention_mask,
633
+ rotary_pos_emb,
634
+ kv_caches[index],
635
+ use_cache
636
+ )
637
+ else:
638
+ layer_ret = layer(
639
+ hidden_states,
640
+ attention_mask,
641
+ rotary_pos_emb,
642
+ kv_cache=kv_caches[index],
643
+ use_cache=use_cache
644
+ )
645
+ hidden_states, kv_cache = layer_ret
646
+ if use_cache:
647
+ presents = presents + (kv_cache,)
648
+
649
+ if output_hidden_states:
650
+ all_hidden_states = all_hidden_states + (hidden_states,)
651
+
652
+ # Final layer norm.
653
+ if self.post_layer_norm:
654
+ hidden_states = self.final_layernorm(hidden_states)
655
+
656
+ return hidden_states, presents, all_hidden_states, all_self_attentions
657
+
658
+
659
+ class ChatGLMPreTrainedModel(PreTrainedModel):
660
+ """
661
+ An abstract class to handle weights initialization and
662
+ a simple interface for downloading and loading pretrained models.
663
+ """
664
+
665
+ is_parallelizable = False
666
+ supports_gradient_checkpointing = True
667
+ config_class = ChatGLMConfig
668
+ base_model_prefix = "transformer"
669
+ _no_split_modules = ["GLMBlock"]
670
+
671
+ def _init_weights(self, module: nn.Module):
672
+ """Initialize the weights."""
673
+ return
674
+
675
+ def get_masks(self, input_ids, past_key_values, padding_mask=None):
676
+ batch_size, seq_length = input_ids.shape
677
+ full_attention_mask = torch.ones(batch_size, seq_length, seq_length, device=input_ids.device)
678
+ full_attention_mask.tril_()
679
+ past_length = 0
680
+ if past_key_values:
681
+ past_length = past_key_values[0][0].shape[0]
682
+ if past_length:
683
+ full_attention_mask = torch.cat((torch.ones(batch_size, seq_length, past_length,
684
+ device=input_ids.device), full_attention_mask), dim=-1)
685
+ if padding_mask is not None:
686
+ full_attention_mask = full_attention_mask * padding_mask.unsqueeze(1)
687
+ if not past_length and padding_mask is not None:
688
+ full_attention_mask -= padding_mask.unsqueeze(-1) - 1
689
+ full_attention_mask = (full_attention_mask < 0.5).bool()
690
+ full_attention_mask.unsqueeze_(1)
691
+ return full_attention_mask
692
+
693
+ def get_position_ids(self, input_ids, device):
694
+ batch_size, seq_length = input_ids.shape
695
+ position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1)
696
+ return position_ids
697
+
698
+ def _set_gradient_checkpointing(self, module, value=False):
699
+ if isinstance(module, GLMTransformer):
700
+ module.gradient_checkpointing = value
701
+
702
+
703
+ class Embedding(torch.nn.Module):
704
+ """Language model embeddings."""
705
+
706
+ def __init__(self, config: ChatGLMConfig, device=None):
707
+ super(Embedding, self).__init__()
708
+
709
+ self.hidden_size = config.hidden_size
710
+ # Word embeddings (parallel).
711
+ self.word_embeddings = nn.Embedding(
712
+ config.padded_vocab_size,
713
+ self.hidden_size,
714
+ dtype=config.torch_dtype,
715
+ device=device
716
+ )
717
+ self.fp32_residual_connection = config.fp32_residual_connection
718
+
719
+ def forward(self, input_ids):
720
+ # Embeddings.
721
+ words_embeddings = self.word_embeddings(input_ids)
722
+ embeddings = words_embeddings
723
+ # Data format change to avoid explicit tranposes : [b s h] --> [s b h].
724
+ embeddings = embeddings.transpose(0, 1).contiguous()
725
+ # If the input flag for fp32 residual connection is set, convert for float.
726
+ if self.fp32_residual_connection:
727
+ embeddings = embeddings.float()
728
+ return embeddings
729
+
730
+
731
+ class ChatGLMModel(ChatGLMPreTrainedModel):
732
+ def __init__(self, config: ChatGLMConfig, device=None, empty_init=True):
733
+ super().__init__(config)
734
+ if empty_init:
735
+ init_method = skip_init
736
+ else:
737
+ init_method = default_init
738
+ init_kwargs = {}
739
+ if device is not None:
740
+ init_kwargs["device"] = device
741
+ self.embedding = init_method(Embedding, config, **init_kwargs)
742
+ self.num_layers = config.num_layers
743
+ self.multi_query_group_num = config.multi_query_group_num
744
+ self.kv_channels = config.kv_channels
745
+
746
+ # Rotary positional embeddings
747
+ self.seq_length = config.seq_length
748
+ rotary_dim = (
749
+ config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels
750
+ )
751
+
752
+ self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, original_impl=config.original_rope, device=device,
753
+ dtype=config.torch_dtype)
754
+ self.encoder = init_method(GLMTransformer, config, **init_kwargs)
755
+ self.output_layer = init_method(nn.Linear, config.hidden_size, config.padded_vocab_size, bias=False,
756
+ dtype=config.torch_dtype, **init_kwargs)
757
+ self.pre_seq_len = config.pre_seq_len
758
+ self.prefix_projection = config.prefix_projection
759
+ if self.pre_seq_len is not None:
760
+ for param in self.parameters():
761
+ param.requires_grad = False
762
+ self.prefix_tokens = torch.arange(self.pre_seq_len).long()
763
+ self.prefix_encoder = PrefixEncoder(config)
764
+ self.dropout = torch.nn.Dropout(0.1)
765
+
766
+ def get_input_embeddings(self):
767
+ return self.embedding.word_embeddings
768
+
769
+ def get_prompt(self, batch_size, device, dtype=torch.half):
770
+ prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(device)
771
+ past_key_values = self.prefix_encoder(prefix_tokens).type(dtype)
772
+ past_key_values = past_key_values.view(
773
+ batch_size,
774
+ self.pre_seq_len,
775
+ self.num_layers * 2,
776
+ self.multi_query_group_num,
777
+ self.kv_channels
778
+ )
779
+ # seq_len, b, nh, hidden_size
780
+ past_key_values = self.dropout(past_key_values)
781
+ past_key_values = past_key_values.permute([2, 1, 0, 3, 4]).split(2)
782
+ return past_key_values
783
+
784
+ def forward(
785
+ self,
786
+ input_ids,
787
+ position_ids: Optional[torch.Tensor] = None,
788
+ attention_mask: Optional[torch.BoolTensor] = None,
789
+ full_attention_mask: Optional[torch.BoolTensor] = None,
790
+ past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
791
+ inputs_embeds: Optional[torch.Tensor] = None,
792
+ use_cache: Optional[bool] = None,
793
+ output_hidden_states: Optional[bool] = None,
794
+ return_dict: Optional[bool] = None,
795
+ ):
796
+ output_hidden_states = (
797
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
798
+ )
799
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
800
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
801
+
802
+ batch_size, seq_length = input_ids.shape
803
+
804
+ if inputs_embeds is None:
805
+ inputs_embeds = self.embedding(input_ids)
806
+
807
+ if self.pre_seq_len is not None:
808
+ if past_key_values is None:
809
+ past_key_values = self.get_prompt(batch_size=batch_size, device=input_ids.device,
810
+ dtype=inputs_embeds.dtype)
811
+ if attention_mask is not None:
812
+ attention_mask = torch.cat([attention_mask.new_ones((batch_size, self.pre_seq_len)),
813
+ attention_mask], dim=-1)
814
+
815
+ if full_attention_mask is None:
816
+ if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1):
817
+ full_attention_mask = self.get_masks(input_ids, past_key_values, padding_mask=attention_mask)
818
+
819
+ # Rotary positional embeddings
820
+ rotary_pos_emb = self.rotary_pos_emb(self.seq_length)
821
+ if position_ids is not None:
822
+ rotary_pos_emb = rotary_pos_emb[position_ids]
823
+ else:
824
+ rotary_pos_emb = rotary_pos_emb[None, :seq_length]
825
+ rotary_pos_emb = rotary_pos_emb.transpose(0, 1).contiguous()
826
+
827
+ # Run encoder.
828
+ hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder(
829
+ inputs_embeds, full_attention_mask, rotary_pos_emb=rotary_pos_emb,
830
+ kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states
831
+ )
832
+
833
+ if not return_dict:
834
+ return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)
835
+
836
+ return BaseModelOutputWithPast(
837
+ last_hidden_state=hidden_states,
838
+ past_key_values=presents,
839
+ hidden_states=all_hidden_states,
840
+ attentions=all_self_attentions,
841
+ )
842
+
843
+ def quantize(self, weight_bit_width: int):
844
+ from .quantization import quantize
845
+ quantize(self.encoder, weight_bit_width)
846
+ return self
847
+
848
+
849
+ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
850
+ def __init__(self, config: ChatGLMConfig, empty_init=True, device=None):
851
+ super().__init__(config)
852
+
853
+ self.max_sequence_length = config.max_length
854
+ self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device)
855
+ self.config = config
856
+ self.quantized = False
857
+
858
+ if self.config.quantization_bit:
859
+ self.quantize(self.config.quantization_bit, empty_init=True)
860
+
861
+ def _update_model_kwargs_for_generation(
862
+ self,
863
+ outputs: ModelOutput,
864
+ model_kwargs: Dict[str, Any],
865
+ is_encoder_decoder: bool = False,
866
+ standardize_cache_format: bool = False,
867
+ ) -> Dict[str, Any]:
868
+ # update past_key_values
869
+ model_kwargs["past_key_values"] = self._extract_past_from_model_output(
870
+ outputs, standardize_cache_format=standardize_cache_format
871
+ )
872
+
873
+ # update attention mask
874
+ if "attention_mask" in model_kwargs:
875
+ attention_mask = model_kwargs["attention_mask"]
876
+ model_kwargs["attention_mask"] = torch.cat(
877
+ [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
878
+ )
879
+
880
+ # update position ids
881
+ if "position_ids" in model_kwargs:
882
+ position_ids = model_kwargs["position_ids"]
883
+ new_position_id = position_ids[..., -1:].clone()
884
+ new_position_id += 1
885
+ model_kwargs["position_ids"] = torch.cat(
886
+ [position_ids, new_position_id], dim=-1
887
+ )
888
+
889
+ model_kwargs["is_first_forward"] = False
890
+ return model_kwargs
891
+
892
+ def prepare_inputs_for_generation(
893
+ self,
894
+ input_ids: torch.LongTensor,
895
+ past_key_values: Optional[torch.Tensor] = None,
896
+ attention_mask: Optional[torch.Tensor] = None,
897
+ position_ids: Optional[torch.Tensor] = None,
898
+ is_first_forward: bool = True,
899
+ **kwargs
900
+ ) -> dict:
901
+ # only last token for input_ids if past is not None
902
+ if position_ids is None:
903
+ position_ids = self.get_position_ids(input_ids, device=input_ids.device)
904
+ if not is_first_forward:
905
+ position_ids = position_ids[..., -1:]
906
+ input_ids = input_ids[:, -1:]
907
+ return {
908
+ "input_ids": input_ids,
909
+ "past_key_values": past_key_values,
910
+ "position_ids": position_ids,
911
+ "attention_mask": attention_mask,
912
+ "return_last_logit": True
913
+ }
914
+
915
+ def forward(
916
+ self,
917
+ input_ids: Optional[torch.Tensor] = None,
918
+ position_ids: Optional[torch.Tensor] = None,
919
+ attention_mask: Optional[torch.Tensor] = None,
920
+ past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
921
+ inputs_embeds: Optional[torch.Tensor] = None,
922
+ labels: Optional[torch.Tensor] = None,
923
+ use_cache: Optional[bool] = None,
924
+ output_attentions: Optional[bool] = None,
925
+ output_hidden_states: Optional[bool] = None,
926
+ return_dict: Optional[bool] = None,
927
+ return_last_logit: Optional[bool] = False,
928
+ ):
929
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
930
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
931
+
932
+ transformer_outputs = self.transformer(
933
+ input_ids=input_ids,
934
+ position_ids=position_ids,
935
+ attention_mask=attention_mask,
936
+ past_key_values=past_key_values,
937
+ inputs_embeds=inputs_embeds,
938
+ use_cache=use_cache,
939
+ output_hidden_states=output_hidden_states,
940
+ return_dict=return_dict,
941
+ )
942
+
943
+ hidden_states = transformer_outputs[0]
944
+ if return_last_logit:
945
+ hidden_states = hidden_states[-1:]
946
+ lm_logits = self.transformer.output_layer(hidden_states)
947
+ lm_logits = lm_logits.transpose(0, 1).contiguous()
948
+
949
+ loss = None
950
+ if labels is not None:
951
+ lm_logits = lm_logits.to(torch.float32)
952
+
953
+ # Shift so that tokens < n predict n
954
+ shift_logits = lm_logits[..., :-1, :].contiguous()
955
+ shift_labels = labels[..., 1:].contiguous()
956
+ # Flatten the tokens
957
+ loss_fct = CrossEntropyLoss(ignore_index=-100)
958
+ loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
959
+
960
+ lm_logits = lm_logits.to(hidden_states.dtype)
961
+ loss = loss.to(hidden_states.dtype)
962
+
963
+ if not return_dict:
964
+ output = (lm_logits,) + transformer_outputs[1:]
965
+ return ((loss,) + output) if loss is not None else output
966
+
967
+ return CausalLMOutputWithPast(
968
+ loss=loss,
969
+ logits=lm_logits,
970
+ past_key_values=transformer_outputs.past_key_values,
971
+ hidden_states=transformer_outputs.hidden_states,
972
+ attentions=transformer_outputs.attentions,
973
+ )
974
+
975
+ @staticmethod
976
+ def _reorder_cache(
977
+ past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor
978
+ ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:
979
+ """
980
+ This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
981
+ [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
982
+ beam_idx at every generation step.
983
+
984
+ Output shares the same memory storage as `past`.
985
+ """
986
+ return tuple(
987
+ (
988
+ layer_past[0].index_select(1, beam_idx.to(layer_past[0].device)),
989
+ layer_past[1].index_select(1, beam_idx.to(layer_past[1].device)),
990
+ )
991
+ for layer_past in past
992
+ )
993
+
994
+ def process_response(self, response):
995
+ response = response.strip()
996
+ response = response.replace("[[训练时间]]", "2023年")
997
+ return response
998
+
999
+ def build_inputs(self, tokenizer, query: str, history: List[Tuple[str, str]] = None):
1000
+ prompt = tokenizer.build_prompt(query, history=history)
1001
+ inputs = tokenizer([prompt], return_tensors="pt")
1002
+ inputs = inputs.to(self.device)
1003
+ return inputs
1004
+
1005
+ def build_stream_inputs(self, tokenizer, query: str, history: List[Tuple[str, str]] = None):
1006
+ if history:
1007
+ prompt = "\n\n[Round {}]\n\n问:{}\n\n答:".format(len(history) + 1, query)
1008
+ input_ids = tokenizer.encode(prompt, add_special_tokens=False)
1009
+ input_ids = input_ids[1:]
1010
+ inputs = tokenizer.batch_encode_plus([(input_ids, None)], return_tensors="pt", add_special_tokens=False)
1011
+ else:
1012
+ prompt = "[Round {}]\n\n问:{}\n\n答:".format(len(history) + 1, query)
1013
+ inputs = tokenizer([prompt], return_tensors="pt")
1014
+ inputs = inputs.to(self.device)
1015
+ return inputs
1016
+
1017
+ @torch.no_grad()
1018
+ def chat(self, tokenizer, query: str, history: List[Tuple[str, str]] = None, max_length: int = 8192, num_beams=1,
1019
+ do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None, **kwargs):
1020
+ if history is None:
1021
+ history = []
1022
+ if logits_processor is None:
1023
+ logits_processor = LogitsProcessorList()
1024
+ logits_processor.append(InvalidScoreLogitsProcessor())
1025
+ gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "do_sample": do_sample, "top_p": top_p,
1026
+ "temperature": temperature, "logits_processor": logits_processor, **kwargs}
1027
+ inputs = self.build_inputs(tokenizer, query, history=history)
1028
+ outputs = self.generate(**inputs, **gen_kwargs)
1029
+ outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):]
1030
+ response = tokenizer.decode(outputs)
1031
+ response = self.process_response(response)
1032
+ history = history + [(query, response)]
1033
+ return response, history
1034
+
1035
+ @torch.no_grad()
1036
+ def stream_chat(self, tokenizer, query: str, history: List[Tuple[str, str]] = None, past_key_values=None,
1037
+ max_length: int = 8192, do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None,
1038
+ return_past_key_values=False, **kwargs):
1039
+ if history is None:
1040
+ history = []
1041
+ if logits_processor is None:
1042
+ logits_processor = LogitsProcessorList()
1043
+ logits_processor.append(InvalidScoreLogitsProcessor())
1044
+ gen_kwargs = {"max_length": max_length, "do_sample": do_sample, "top_p": top_p,
1045
+ "temperature": temperature, "logits_processor": logits_processor, **kwargs}
1046
+ if past_key_values is None and not return_past_key_values:
1047
+ inputs = self.build_inputs(tokenizer, query, history=history)
1048
+ else:
1049
+ inputs = self.build_stream_inputs(tokenizer, query, history=history)
1050
+ if past_key_values is not None:
1051
+ past_length = past_key_values[0][0].shape[0]
1052
+ if self.transformer.pre_seq_len is not None:
1053
+ past_length -= self.transformer.pre_seq_len
1054
+ inputs.position_ids += past_length
1055
+ attention_mask = inputs.attention_mask
1056
+ attention_mask = torch.cat((attention_mask.new_ones(1, past_length), attention_mask), dim=1)
1057
+ inputs['attention_mask'] = attention_mask
1058
+ for outputs in self.stream_generate(**inputs, past_key_values=past_key_values,
1059
+ return_past_key_values=return_past_key_values, **gen_kwargs):
1060
+ if return_past_key_values:
1061
+ outputs, past_key_values = outputs
1062
+ outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):]
1063
+ response = tokenizer.decode(outputs)
1064
+ if response and response[-1] != "�":
1065
+ response = self.process_response(response)
1066
+ new_history = history + [(query, response)]
1067
+ if return_past_key_values:
1068
+ yield response, new_history, past_key_values
1069
+ else:
1070
+ yield response, new_history
1071
+
1072
+ @torch.no_grad()
1073
+ def stream_generate(
1074
+ self,
1075
+ input_ids,
1076
+ generation_config: Optional[GenerationConfig] = None,
1077
+ logits_processor: Optional[LogitsProcessorList] = None,
1078
+ stopping_criteria: Optional[StoppingCriteriaList] = None,
1079
+ prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
1080
+ return_past_key_values=False,
1081
+ **kwargs,
1082
+ ):
1083
+ batch_size, input_ids_seq_length = input_ids.shape[0], input_ids.shape[-1]
1084
+
1085
+ if generation_config is None:
1086
+ generation_config = self.generation_config
1087
+ generation_config = copy.deepcopy(generation_config)
1088
+ model_kwargs = generation_config.update(**kwargs)
1089
+ bos_token_id, eos_token_id = generation_config.bos_token_id, generation_config.eos_token_id
1090
+
1091
+ if isinstance(eos_token_id, int):
1092
+ eos_token_id = [eos_token_id]
1093
+
1094
+ has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
1095
+ if has_default_max_length and generation_config.max_new_tokens is None:
1096
+ warnings.warn(
1097
+ f"Using `max_length`'s default ({generation_config.max_length}) to control the generation length. "
1098
+ "This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we"
1099
+ " recommend using `max_new_tokens` to control the maximum length of the generation.",
1100
+ UserWarning,
1101
+ )
1102
+ elif generation_config.max_new_tokens is not None:
1103
+ generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length
1104
+ if not has_default_max_length:
1105
+ logger.warn(
1106
+ f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="
1107
+ f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "
1108
+ "Please refer to the documentation for more information. "
1109
+ "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)",
1110
+ UserWarning,
1111
+ )
1112
+
1113
+ if input_ids_seq_length >= generation_config.max_length:
1114
+ input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"
1115
+ logger.warning(
1116
+ f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to"
1117
+ f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
1118
+ " increasing `max_new_tokens`."
1119
+ )
1120
+
1121
+ # 2. Set generation parameters if not already defined
1122
+ logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
1123
+ stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
1124
+
1125
+ logits_processor = self._get_logits_processor(
1126
+ generation_config=generation_config,
1127
+ input_ids_seq_length=input_ids_seq_length,
1128
+ encoder_input_ids=input_ids,
1129
+ prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
1130
+ logits_processor=logits_processor,
1131
+ )
1132
+
1133
+ stopping_criteria = self._get_stopping_criteria(
1134
+ generation_config=generation_config, stopping_criteria=stopping_criteria
1135
+ )
1136
+ logits_warper = self._get_logits_warper(generation_config)
1137
+
1138
+ unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1)
1139
+ scores = None
1140
+ while True:
1141
+ model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
1142
+ # forward pass to get next token
1143
+ outputs = self(
1144
+ **model_inputs,
1145
+ return_dict=True,
1146
+ output_attentions=False,
1147
+ output_hidden_states=False,
1148
+ )
1149
+
1150
+ next_token_logits = outputs.logits[:, -1, :]
1151
+
1152
+ # pre-process distribution
1153
+ next_token_scores = logits_processor(input_ids, next_token_logits)
1154
+ next_token_scores = logits_warper(input_ids, next_token_scores)
1155
+
1156
+ # sample
1157
+ probs = nn.functional.softmax(next_token_scores, dim=-1)
1158
+ if generation_config.do_sample:
1159
+ next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
1160
+ else:
1161
+ next_tokens = torch.argmax(probs, dim=-1)
1162
+
1163
+ # update generated ids, model inputs, and length for next step
1164
+ input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
1165
+ model_kwargs = self._update_model_kwargs_for_generation(
1166
+ outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder
1167
+ )
1168
+ unfinished_sequences = unfinished_sequences.mul((sum(next_tokens != i for i in eos_token_id)).long())
1169
+ if return_past_key_values:
1170
+ yield input_ids, outputs.past_key_values
1171
+ else:
1172
+ yield input_ids
1173
+ # stop when each sentence is finished, or if we exceed the maximum length
1174
+ if unfinished_sequences.max() == 0 or stopping_criteria(input_ids, scores):
1175
+ break
1176
+
1177
+ def quantize(self, bits: int, empty_init=False, device=None, **kwargs):
1178
+ if bits == 0:
1179
+ return
1180
+
1181
+ from .quantization import quantize
1182
+
1183
+ if self.quantized:
1184
+ logger.info("Already quantized.")
1185
+ return self
1186
+
1187
+ self.quantized = True
1188
+
1189
+ self.config.quantization_bit = bits
1190
+
1191
+ self.transformer.encoder = quantize(self.transformer.encoder, bits, empty_init=empty_init, device=device,
1192
+ **kwargs)
1193
+ return self
checkpoint-2000/optimizer.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6a36a1229d3e1a94161dffd08ad13fb4cb4fd73c9d8f541a7a661e7adc4707cd
3
+ size 1346374055
checkpoint-2000/pytorch_model.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e5bb05c5aa0234c123bd782c7b8672f9b922b5be7fa1205074361b22197af545
3
+ size 673186943
checkpoint-2000/quantization.py ADDED
@@ -0,0 +1,188 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from torch.nn import Linear
2
+ from torch.nn.parameter import Parameter
3
+
4
+ import bz2
5
+ import torch
6
+ import base64
7
+ import ctypes
8
+ from transformers.utils import logging
9
+
10
+ from typing import List
11
+ from functools import partial
12
+
13
+ logger = logging.get_logger(__name__)
14
+
15
+ try:
16
+ from cpm_kernels.kernels.base import LazyKernelCModule, KernelFunction, round_up
17
+
18
+ class Kernel:
19
+ def __init__(self, code: bytes, function_names: List[str]):
20
+ self.code = code
21
+ self._function_names = function_names
22
+ self._cmodule = LazyKernelCModule(self.code)
23
+
24
+ for name in self._function_names:
25
+ setattr(self, name, KernelFunction(self._cmodule, name))
26
+
27
+ quantization_code = "$QlpoOTFBWSZTWU9yuJUAQHN//////////f/n/8/n///n//bt4dTidcVx8X3V9FV/92/v4B7/AD5FBQFAAAChSgKpFCFAFVSigUAAAEKhSgUUqgFBKigqVREQAABQBQIANDTTIGI00BkZBkNGE0A0BkBkGQGRkaNAaAGQNBoGgDIAAYIGTI0DQAQAaGmmQMRpoDIyDIaMJoBoDIDIMgMjI0aA0AMgaDQNAGQAAwQMmRoGgAgA0NNMgYjTQGRkGQ0YTQDQGQGQZAZGRo0BoAZA0GgaAMgABggZMjQNABABoaaZAxGmgMjIMhowmgGgMgMgyAyMjRoDQAyBoNA0AZAADBAyZGgaAAmqU1NEgJqnptU/Sn4jRR6J6epk2pqb1Q/SgAPUGgyNNGjQ2SBpoAZAAGg0NB6mgDIAAAAA2oaApSREBNAARhGiYEaEwU8pvImlP0k2aam1GaGqbFNM1MHpTwmkepmyU9R6nqPKekHqNNPUxNGhp6n6p6QaZ6o9TG1GMqcoV9ly6nRanHlq6zPNbnGZNi6HSug+2nPiZ13XcnFYZW+45W11CumhzYhchOJ2GLLV1OBjBjGf4TptOddTSOcVxhqYZMYwZXZZY00zI1paX5X9J+b+f4e+x43RXSxXPOdquiGpduatGyXneN696M9t4HU2eR5XX/kPhP261NTx3JO1Ow7LyuDmeo9a7d351T1ZxnvnrvYnrXv/hXxPCeuYx2XsNmO003eg9J3Z6U7b23meJ4ri01OdzTk9BNO96brz+qT5nuvvH3ds/G+m/JcG/F2XYuhXlvO+jP7U3XgrzPN/lr8Sf1n6j4j7jZs+s/T0tNaNNYzTs12rxjwztHlnire3Nzc3N1wuBwOBwXBvZfoHpD7rFmR99V5vj3aXza3xdBbXMalubTg/jIv5dfAi54Pdc75j4z412n3Npj3Ld/ENm7a3b/Cod6h/ret1/5vn/C+l+gdslMvgPSLJ8d8q+U66fevYn/tW1chleEtNTGlcHCbLRlq0tHzF5tsbbZZfHjjLgZu42XCuC3NrdjTasZGNzgxPIrGqp7r3p7L2p5XjnpPSmTd5XtzqnB6U87zzg1Ol0zd0zsLszxR6lkxp35u6/teL0L0W922cR7Lu1lpL9CsHirzuM2T+BgsyViT6LHcm0/Vr6U/7LGGyJeqTEjt0PHWhF5mCT7R9mtlDwriYv0Tyr/OxYt6qp5r0mPVT0608TqnqMZaarU2nFwrTzzlrs1ed7z1ux60wyr4ydCaTi3enW8x68x0zU7tXSlcmPSW1mGpWJMg4zmPC2lK96tp0OE80y4MfEvnZj8zGluR6b22ki1Ou9V2nCd9xovcPvcYMZYy0lvN60ScZ45vN6yeCeeXFb1lVjnnCar5fwXwE2bzJ4HI1XVPXfXZMm44GUsMpYsmLB65TuVdm0cl0b+i/wGNN66XjeV7zuPpHcnK/juhhjdfId5jMdE5nN0dGmmm2zZs2cexD5n9p/dY352XsvXHaZNWWsmmS1atjR452nYudzvqv2HMRyvNNnlMcDl3R2+yx2uVrBubTW9icHDVtbNXlZm7jma1rM4VurZZd2y6nUau7ZXZ7bVU+mnoOVxZGMrVmvX60605JwmzGZhhhjTWtaaaMaaGTGmNMZasY0iX8VMUl8eepaIrzGSpemWOQyZORk2bNpjUybMmxqYmknCGCFynutfksaZpjTNMaaatM0xsxcGR0sociNqxNSmhhR1ZJPbsn8qyF0t2qH6iYBclclalbtTTcHTDsPaX6rlnElph2Jyumumtynv2Kk8GI7rsvXbIcJgHJOSaSXnnGaI3m87RtVXJOZ/YtgdTE6Wpha6ZlE8ayXkef1fh602r2WwvfMXtMdLlkfnLFdYYwYso+bWqm7yJqHXZGw2nrS5ZanSYnWlxBxMF1V940K2wdrI7R6OYf7DGGamMmTSbRhlS45xmVOumF1EyPCmHrrN8wwZOOrdNtLeMtzFzDlWnfTBxMk2NaXIZHBYxYLD4w8yju0ao65Vz1OIXoS9dLanwCe1PWrYuWMqf1if1z2k2yYfKJ741PDgno1ZQ8DRqvUny3mNoWTzGO6m1DkrJI8JiR5cSd+vZdGOO8nrMoc5+NDUFsMSXaZJeNlMmGLtJsovOsUp7I9S5VojKxF6bTVEelXqlfJobQr3LozSh2Jk7VcrVMfhXqszGWMzNqGhqZY0OadxkyyMssKugZR0KNFXBHlqwmJgTE/BNVMk6ItJXZMR0H47GpXv/DMOvNkmVuaV1PRfEdxuqc7Hcd+ZV/zTLaRxWk0nl9CdCeM6mn5rstHIBcpiuwmUZXeq81DacHI2rmrZ5SuE5mOZd6LQrZg9mx32TprA8BMo5jKN6yLTCi3WzQaZSuhzTtM1fUTGVpG8Tw+KXI0tjEpiWxtLYynOlktSbVlaI5kxP8TDH8kx50xoxi5KcA4pcja8KWLRlO/Ks6q06ergnvm1ca3Tq8Uw7LTUsmWyctXPWmpitl/uvGcWTGXGuAXDfhqazGmjkxcJW5hMMMMpYsXl2TZYtVOddG3XCarUt6Ptq9CZXSNzyuRzqRZOjsxdBbFVz6OA5HI43r1jityVlVpVkxmOsyaYWE1NTGq1sOVh36mHMcxtSvcy70edG0ZGR3I1Go1GRlV7mWWo1G0ZGRqlvH40l7o4m5xMWLLLYyNjnqc8556mdPqLJ31n/1nWOncxzG1tizrHs/Z+d2vP/B/l8wdJ6rHUn2nbbDq4p6htFtYzMMMTaZis1K5GKzGNmxhmUx2DDlZ/qNnIx41xnaMfCZWYaZWtNLTNW8ND4Fw1MyZOCdM428suKG1ehW8TesOydg7J+YYcD4cYR+8dFK6M4E3HM9ZfRNNL+Sn6rsl4DsrDl2HpPCnfxjGXtbZtYys1ttlyJ4T+BvexjGWRjMszK4Jpc77D3GyuVD7q0+G8m9G+2+rGm7cOR2y7FdtY2XUYx/oNlfRYxhMYyYZkyyg55enna9Kt/FFi6GMMwYwdwxWgxGMLKYmUyGExTKMZkMFhkymKuh0NOBNnBu+23LdwDoZYYzGGMxtORaTU1pjTGWTTGGtMrNWUsyyTTLLG1qy2ZjbK2DBllWqxMtBMaYZQmcE7zvvRcTkclUwdkxTaSdyySt/7fpL+T1v516Ji97fwr5JbLu305zMn5+GMTTZ9F+y7ExwmGVfG44yxn3dLv6l5i+Wth1jCrDq21nW9LqvvDzz3Vf3LLH/O/32TJ/erx3bXftO4eF+G956D952K/An4NfvOpjFjExjevP/UmE0fIoZXx6/w6lX/no3D0bLt+ixjieBM6ksRd0yB4Lt2SwYNE+gd1detlZWUnpiZfGfFaK+4PyCa/v18V8X75pe9fLXzp7l3VjF76vWZmHwGz1IZNWT7b8yddJ4q5kyrVdfru6atWc7bVYztL9Jf4GXvT+Y8m9/YsXP6H018a8D4XVOqvfzqeR+6yZOD8dPv0+U7/q5Pl+2dNb0MjzGVH5p6MNQ7cOWvw62U9aHE8DprDek+McLyvDz+te+9Zhq5+YTruufMcWMabqysTmZVWjKPfnK0wyVcrsuhjZRdLkHNvD72b9abriOSGIxiLixMOoalNPXzy+wT/tf+U6HHONfsz+xe8ufHBdQWWGWLA9if0rsnmrxK5LvRZQeWsTCsrmOYy8VteVfuRfcVTtDLItLIsMYxZLdU/DbtSemxF6Z6Zo5WBXE4tFdCyVMMXMTEMZXVlS6Xec2T4e0tHsRcEuWshcJ2YsNF5rUx1E8ifCq6Z+ZP7qdCeu/aTwFd53l16/o0NOw6O3dLavP4Hbi4RdmuDk6DoYaninC0+o4uZjbJ7Rxeu0/FbuFg+q7DVS6fQe0rZ6NDGUNNU6DEqOaLTicKnYZMnBWruljQxoaS3dZhocDge0bSTyOvdAbG5hxe2xji7E/L55xX13wWNDi6HCekcFxfCPGxY0MXC+s7afWaMdDyjyr+o8Rudm/NabOZvdl274zH4f5XK9z6On1Pe/K5TdPAslg77BjuO6Y3eO7GqvOPG/stknp1leyvLL0Z7bl9I4noMvLkzytLhWYzrOZzLXCORe028rORzOg4N/L0HlMOQ3Pgmnbb6KczlabORpu980q37TBqRu0/p3PO6234Bl03Ynuz+9W7gnsEcmvYaYY3aMYY0wx3pYd+ujsXauWdaY5Xkbtl23fPzFHiDB/QMo0yFjBllYxTQYYyxkrwn7JufwJ/PfgJ+C83X69ni6zvXcnyXabv0ncbLwsceS+RNlyN2mnneJtX0ngYO0+e+0+UnA+Wch3ji8hj5an4h+i6XBySU4n+R0roVcbw5yvHrmr4Yw8Y7x6c+9POPYHI5HI5HI5HI5HGXGww4nE4nrVyOR8XeqPEO7PLOiukYa3Novk5hV4cdtYZLI93e+uxff2jRo0aNGjRo0aNG1bVtW1dy3m83m8+tQ5ZzHw3nObwOu8La9Rc1dtkdS8A3eTk823tnktXWlxN6Oixe06zrN70Isd9jiOgZFq9yfkPqP/SLhN2Myl8jDM43bl1nbcb4cO57jlh8Jow6pzXZdL4dyODTuuhu77FyO27DdwdRxmvO+O+3N2+BdqyTwLHVczDVY4UPE4O66/ZO2cx1LFzVdSXtF7G4HMbrauOHRw6c8FdZ5m9fHZHYZXfTlZquyynSyTTKke6vcffSD9pzPA/G7n7jxPmuhc1DHMynPMrGL6AdewYmwu5ko+UUyTwrMv27rPH1v1nGqd87+p6N6LU8k3NEng53xXyHS97+44OSg/sy/hn+Se6yfYNjW0/uTgP+PvWYzLMmjhcLB/gGpri6H83/84eUXWT6T9Hsv7785z/7z4icpW+zfXypuR7rx/gMdZb1/wC678pcs8/2a3mDitGHxl9mfPlll5MafWWqxk/eYuTDgcNMzDGWLWvsuglNxs53GtN6uWpktlW1tZZYcuinMMWmnNnJydze3b2Y1McBxrBkXw799izLMZZYyy0TkbsGM4p03S2uVu5s/XXUdSdec6smVxZYYGpVmT8A+8ajuEyV5FatkvVru2x6uxGXXbH4A+jvgP4GMYy3iPLXzq/6z65+E005ey+cwMZD3fZcqc6xpjTFjQ0P3U+e++cPYmTIwj0nrK5NPTfl3WvpfLtXDcb2HQMudYOxFXQBor4L4T6vrOauFctYXJQ++NUWmJe5bmx1jDiZS1dTqWxo4GR8jm3fttpmPHppk9PEyv4/y8/sO07XacOmcqc0x2Vi9BvNJvN5oW8x4mOsydpidRxMYJPx06m1bqPzq9KtK8sxXNXFodD/+MYYaJTLwOhc9brCsV18oOR1i4tXChyTkq4lf4y1Ke+9axjDHqs1mfBbMXuP4Hzi+X7t8vzv7bHerrUPgPCxhjre4fXdfLNtNM+Jd+Zdh8xd8wP87uNPoPgv4W7/5P2BuxfsMabNnMnza+54Pdi5U671GPZY8CehX8Voeoo7FHpkeEc6715FwHZrIrUrHaviPUbPZHND+IhczrP6FcYvhOZ0Di/ETt0OI+YwNWR9r7tpf6WDeZKZDB1+z2IthOl1mPyb5FluvEx9h9d0NnM0Y1XPFkWIsk1WotJ0PBMmkvjvQTd0e71tfeV+8r8lQ/tpzpsmxJ+InrI/dj2UajUajVTUajatRqNRtGo1Go1Go4wjeMpZFMVV9CHbofPraLsJ3JpWV2XOoanCuFky4y3PPNxucK2uKC1Lbdb1eo+m5XomN6HfeZsabHLHRX/K+offtNGGmHWctcVcG44MdSqsOLY9VzX+Zxfxn2HPdWTpzWvkrtJ8M5zorrKcquRytJ5N5DZmcaW02l76nWO+BqPXm1A2Ry/0q71dH/mqrqeFjkYxjEXtsX8qubTk67rGycyqsdm4tZx5D6D5hhi0waaWmiaMP81Yjii5qxPlPuU/GfTL1Y5E6Jyfiq63qTa39A4J0sOGDgO9WF9bOXl0XfPRbsY2bPNKPy1YrFYrFYmRhhlTIyMjJWJYZHXuCXI8OoXsvfljGLFicNifpp2XunoPiG1wtx3p1Tah+/DD66OnVtVXP9rKbVxOnL0tR/rHtqB5UDErUVcl11D4qqvjpOcxX7armUNJB3LpW6bxVvD08e8h3odKKvyCFZBdSh2FVcST9xV3n3T8t1j7Kr9qgrqXg+13Pt5U7JCvFXVIV1YG5lRhkVYZJYYDDD4KOIMoHCp26WS8GB7uBh2zIdgq/PKyInjV2STShuoapUdCpX1yTwqq/z1VvET7Kh5nVPkO8YyxjLt2MaaMmWTLQvx3qnzltnXW0p2jxgbEtSny/Osv8Y9pLMXYoHVPAhkVdWVeODhR6q9/Sxe2liwwZWMVvFXfRkeIDxAePUPIrdJ4ey6yquzH+PD/bUOWAu05qVHtFd8rrKHSoeNIOUqrYr3FXyToqfYJgwmJdKpXXOwYYegNNGMzfZPp/t3t/DVs4zjNTN61rRqaWaa4NYbRjTa0tWwy2Y2tGN8ZO8ofNKq4j9SL7I+cSm4/6ovLV5HNXLI0jJidwrtk6ynCaP6Z++GjRlWS3tLeW129Mi9evxU9mtz6s5J3Z7M2ngTgnKvmpomxpaLCzPfmx0JWE+m3NLDDGOX47RctdYYNK5jakdqLkRlI39n590T5zctGSwwZZDJj6kW8XSi6ot2MmWWJ0DUT3nuvebBudScjZ79g8cWJ8av0k+/bE5WKd5MdbFpbDVMxu1DVMmtNZGJvq1mtRbn6M+g/kP0FwDwr7quZs7xosNGpbscyxhhd9TyJyFwbLcxlTasg75vW7TsV5K7ji44XPMMrdoj+Y3rT0Hie62nlYV/pwczzOmdLqLhYkzGMzCZWGMQzGMSsZYY6Di1t4nlJ+Em63mJxrVLxPbYxNEdgc1dU2iOKyoYYWjNrEeHTYybVk0atSa7ehuwsWMWTqn1TrnS6hYsi71d1+s+k+ic70e20fzE/VaTdxT9ZtU4GIXdeNx3X77guYYfpHeTQjaMX6brOu4OY4K7Y2d9mbHarI5ox3p4GpJ2Vd/Tst60f7j999pppjR+Q/Qf8J/VaORs3cji7FfFuN61+ui9s8hix1OCh5KGVV23BPXvZfz3CLyHpix+exi8z/KnCnosY2eunor+cxyPO/xJ0vKey9OvE9VjqaYu0x3Z3jd6o2b1T12D+F8l232lwaaacD5LE8LBxu7WTlbWraWpew8Xexjel3E+wWD4APITdNqR8F3R3T0lunCQ4GaE9R37DxeCYfcHi4xci5ovKfxVs55y2hf+65E/Xdp6jR5nrebTmi5incpkyOjs50JvrZwstbbW6kfuuQw+2mykf/EXNFzxfKTrxew929TR6bWnGL//F3JFOFCQT3K4lQ"
28
+
29
+ kernels = Kernel(
30
+ bz2.decompress(base64.b64decode(quantization_code)),
31
+ [
32
+ "int4WeightCompression",
33
+ "int4WeightExtractionFloat",
34
+ "int4WeightExtractionHalf",
35
+ "int8WeightExtractionFloat",
36
+ "int8WeightExtractionHalf",
37
+ ],
38
+ )
39
+ except Exception as exception:
40
+ kernels = None
41
+ logger.warning("Failed to load cpm_kernels:" + str(exception))
42
+
43
+
44
+ class W8A16Linear(torch.autograd.Function):
45
+ @staticmethod
46
+ def forward(ctx, inp: torch.Tensor, quant_w: torch.Tensor, scale_w: torch.Tensor, weight_bit_width):
47
+ ctx.inp_shape = inp.size()
48
+ ctx.weight_bit_width = weight_bit_width
49
+ out_features = quant_w.size(0)
50
+ inp = inp.contiguous().view(-1, inp.size(-1))
51
+ weight = extract_weight_to_half(quant_w, scale_w, weight_bit_width)
52
+ ctx.weight_shape = weight.size()
53
+ output = inp.mm(weight.t())
54
+ ctx.save_for_backward(inp, quant_w, scale_w)
55
+ return output.view(*(ctx.inp_shape[:-1] + (out_features,)))
56
+
57
+ @staticmethod
58
+ def backward(ctx, grad_output: torch.Tensor):
59
+ inp, quant_w, scale_w = ctx.saved_tensors
60
+ weight = extract_weight_to_half(quant_w, scale_w, ctx.weight_bit_width)
61
+ grad_output = grad_output.contiguous().view(-1, weight.size(0))
62
+ grad_input = grad_output.mm(weight)
63
+ grad_weight = grad_output.t().mm(inp)
64
+ return grad_input.view(ctx.inp_shape), grad_weight.view(ctx.weight_shape), None, None
65
+
66
+
67
+ def compress_int4_weight(weight: torch.Tensor): # (n, m)
68
+ with torch.cuda.device(weight.device):
69
+ n, m = weight.size(0), weight.size(1)
70
+ assert m % 2 == 0
71
+ m = m // 2
72
+ out = torch.empty(n, m, dtype=torch.int8, device="cuda")
73
+ stream = torch.cuda.current_stream()
74
+
75
+ gridDim = (n, 1, 1)
76
+ blockDim = (min(round_up(m, 32), 1024), 1, 1)
77
+
78
+ kernels.int4WeightCompression(
79
+ gridDim,
80
+ blockDim,
81
+ 0,
82
+ stream,
83
+ [ctypes.c_void_p(weight.data_ptr()), ctypes.c_void_p(out.data_ptr()), ctypes.c_int32(n), ctypes.c_int32(m)],
84
+ )
85
+ return out
86
+
87
+
88
+ def extract_weight_to_half(weight: torch.Tensor, scale_list: torch.Tensor, source_bit_width: int):
89
+ assert scale_list.dtype in [torch.half, torch.bfloat16]
90
+ assert weight.dtype in [torch.int8]
91
+ if source_bit_width == 8:
92
+ return weight.to(scale_list.dtype) * scale_list[:, None]
93
+ elif source_bit_width == 4:
94
+ func = (
95
+ kernels.int4WeightExtractionHalf if scale_list.dtype == torch.half else kernels.int4WeightExtractionBFloat16
96
+ )
97
+ else:
98
+ assert False, "Unsupported bit-width"
99
+
100
+ with torch.cuda.device(weight.device):
101
+ n, m = weight.size(0), weight.size(1)
102
+ out = torch.empty(n, m * (8 // source_bit_width), dtype=scale_list.dtype, device="cuda")
103
+ stream = torch.cuda.current_stream()
104
+
105
+ gridDim = (n, 1, 1)
106
+ blockDim = (min(round_up(m, 32), 1024), 1, 1)
107
+
108
+ func(
109
+ gridDim,
110
+ blockDim,
111
+ 0,
112
+ stream,
113
+ [
114
+ ctypes.c_void_p(weight.data_ptr()),
115
+ ctypes.c_void_p(scale_list.data_ptr()),
116
+ ctypes.c_void_p(out.data_ptr()),
117
+ ctypes.c_int32(n),
118
+ ctypes.c_int32(m),
119
+ ],
120
+ )
121
+ return out
122
+
123
+
124
+ class QuantizedLinear(torch.nn.Module):
125
+ def __init__(self, weight_bit_width: int, weight, bias=None, device="cpu", dtype=None, empty_init=False, *args,
126
+ **kwargs):
127
+ super().__init__()
128
+ self.weight_bit_width = weight_bit_width
129
+
130
+ shape = weight.shape
131
+
132
+ if weight is None or empty_init:
133
+ self.weight = torch.empty(shape[0], shape[1] * weight_bit_width // 8, dtype=torch.int8, device=device)
134
+ self.weight_scale = torch.empty(shape[0], dtype=dtype, device=device)
135
+ else:
136
+ self.weight_scale = weight.abs().max(dim=-1).values / ((2 ** (weight_bit_width - 1)) - 1)
137
+ self.weight = torch.round(weight / self.weight_scale[:, None]).to(torch.int8)
138
+ if weight_bit_width == 4:
139
+ self.weight = compress_int4_weight(self.weight)
140
+
141
+ self.weight = Parameter(self.weight.to(device), requires_grad=False)
142
+ self.weight_scale = Parameter(self.weight_scale.to(device), requires_grad=False)
143
+ self.bias = Parameter(bias.to(device), requires_grad=False) if bias is not None else None
144
+
145
+ def forward(self, input):
146
+ output = W8A16Linear.apply(input, self.weight, self.weight_scale, self.weight_bit_width)
147
+ if self.bias is not None:
148
+ output = output + self.bias
149
+ return output
150
+
151
+
152
+ def quantize(model, weight_bit_width, empty_init=False, device=None):
153
+ """Replace fp16 linear with quantized linear"""
154
+ for layer in model.layers:
155
+ layer.self_attention.query_key_value = QuantizedLinear(
156
+ weight_bit_width=weight_bit_width,
157
+ weight=layer.self_attention.query_key_value.weight.to(torch.cuda.current_device()),
158
+ bias=layer.self_attention.query_key_value.bias,
159
+ dtype=layer.self_attention.query_key_value.weight.dtype,
160
+ device=layer.self_attention.query_key_value.weight.device if device is None else device,
161
+ empty_init=empty_init
162
+ )
163
+ layer.self_attention.dense = QuantizedLinear(
164
+ weight_bit_width=weight_bit_width,
165
+ weight=layer.self_attention.dense.weight.to(torch.cuda.current_device()),
166
+ bias=layer.self_attention.dense.bias,
167
+ dtype=layer.self_attention.dense.weight.dtype,
168
+ device=layer.self_attention.dense.weight.device if device is None else device,
169
+ empty_init=empty_init
170
+ )
171
+ layer.mlp.dense_h_to_4h = QuantizedLinear(
172
+ weight_bit_width=weight_bit_width,
173
+ weight=layer.mlp.dense_h_to_4h.weight.to(torch.cuda.current_device()),
174
+ bias=layer.mlp.dense_h_to_4h.bias,
175
+ dtype=layer.mlp.dense_h_to_4h.weight.dtype,
176
+ device=layer.mlp.dense_h_to_4h.weight.device if device is None else device,
177
+ empty_init=empty_init
178
+ )
179
+ layer.mlp.dense_4h_to_h = QuantizedLinear(
180
+ weight_bit_width=weight_bit_width,
181
+ weight=layer.mlp.dense_4h_to_h.weight.to(torch.cuda.current_device()),
182
+ bias=layer.mlp.dense_4h_to_h.bias,
183
+ dtype=layer.mlp.dense_4h_to_h.weight.dtype,
184
+ device=layer.mlp.dense_4h_to_h.weight.device if device is None else device,
185
+ empty_init=empty_init
186
+ )
187
+
188
+ return model
checkpoint-2000/rng_state_0.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c5775c11ec2fe6a5d58d1857fa436bb0da77027386858b2a4e7ea3eff9c2ab66
3
+ size 17655
checkpoint-2000/rng_state_1.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ca1a0e79990b2325280b0f88b8b2e25a4c5d61927bfa5e5c178ffe1321125ca9
3
+ size 17655
checkpoint-2000/rng_state_2.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0aebaa297ea8d8fd38ae9b49e17f26565783e27c56a220dc278f70f6593bccda
3
+ size 17655
checkpoint-2000/rng_state_3.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4f537c78a9ca12294baa835f33188b772f2cbf238559c0f36ab3505e19304755
3
+ size 17655
checkpoint-2000/scheduler.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:453e1733981451ca80e2fbedefd8bde3424d9b57a85cba2e0cc72dc2057863e3
3
+ size 627
checkpoint-2000/special_tokens_map.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {}
checkpoint-2000/tokenization_chatglm.py ADDED
@@ -0,0 +1,253 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import torch
3
+ from typing import List, Optional, Union, Dict
4
+ from sentencepiece import SentencePieceProcessor
5
+ from transformers import PreTrainedTokenizer
6
+ from transformers.utils import logging, PaddingStrategy
7
+ from transformers.tokenization_utils_base import EncodedInput, BatchEncoding
8
+
9
+
10
+ class SPTokenizer:
11
+ def __init__(self, model_path: str):
12
+ # reload tokenizer
13
+ assert os.path.isfile(model_path), model_path
14
+ self.sp_model = SentencePieceProcessor(model_file=model_path)
15
+
16
+ # BOS / EOS token IDs
17
+ self.n_words: int = self.sp_model.vocab_size()
18
+ self.bos_id: int = self.sp_model.bos_id()
19
+ self.eos_id: int = self.sp_model.eos_id()
20
+ self.pad_id: int = self.sp_model.unk_id()
21
+ assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()
22
+
23
+ special_tokens = ["[MASK]", "[gMASK]", "[sMASK]", "sop", "eop"]
24
+ self.special_tokens = {}
25
+ self.index_special_tokens = {}
26
+ for token in special_tokens:
27
+ self.special_tokens[token] = self.n_words
28
+ self.index_special_tokens[self.n_words] = token
29
+ self.n_words += 1
30
+
31
+ def tokenize(self, s: str):
32
+ return self.sp_model.EncodeAsPieces(s)
33
+
34
+ def encode(self, s: str, bos: bool = False, eos: bool = False) -> List[int]:
35
+ assert type(s) is str
36
+ t = self.sp_model.encode(s)
37
+ if bos:
38
+ t = [self.bos_id] + t
39
+ if eos:
40
+ t = t + [self.eos_id]
41
+ return t
42
+
43
+ def decode(self, t: List[int]) -> str:
44
+ return self.sp_model.decode(t)
45
+
46
+ def decode_tokens(self, tokens: List[str]) -> str:
47
+ text = self.sp_model.DecodePieces(tokens)
48
+ return text
49
+
50
+ def convert_token_to_id(self, token):
51
+ """ Converts a token (str) in an id using the vocab. """
52
+ if token in self.special_tokens:
53
+ return self.special_tokens[token]
54
+ return self.sp_model.PieceToId(token)
55
+
56
+ def convert_id_to_token(self, index):
57
+ """Converts an index (integer) in a token (str) using the vocab."""
58
+ if index in self.index_special_tokens or index in [self.eos_id, self.bos_id, self.pad_id] or index < 0:
59
+ return ""
60
+ return self.sp_model.IdToPiece(index)
61
+
62
+
63
+ class ChatGLMTokenizer(PreTrainedTokenizer):
64
+ vocab_files_names = {"vocab_file": "tokenizer.model"}
65
+
66
+ model_input_names = ["input_ids", "attention_mask", "position_ids"]
67
+
68
+ def __init__(self, vocab_file, padding_side="left", **kwargs):
69
+ super().__init__(padding_side=padding_side, **kwargs)
70
+ self.name = "GLMTokenizer"
71
+
72
+ self.vocab_file = vocab_file
73
+ self.tokenizer = SPTokenizer(vocab_file)
74
+ self.special_tokens = {
75
+ "<bos>": self.tokenizer.bos_id,
76
+ "<eos>": self.tokenizer.eos_id,
77
+ "<pad>": self.tokenizer.pad_id
78
+ }
79
+
80
+ def get_command(self, token):
81
+ if token in self.special_tokens:
82
+ return self.special_tokens[token]
83
+ assert token in self.tokenizer.special_tokens, f"{token} is not a special token for {self.name}"
84
+ return self.tokenizer.special_tokens[token]
85
+
86
+ @property
87
+ def pad_token(self) -> str:
88
+ return "<unk>"
89
+
90
+ @property
91
+ def pad_token_id(self):
92
+ return self.get_command("<pad>")
93
+
94
+ @property
95
+ def eos_token(self) -> str:
96
+ return "</s>"
97
+
98
+ @property
99
+ def eos_token_id(self):
100
+ return self.get_command("<eos>")
101
+
102
+ @property
103
+ def vocab_size(self):
104
+ return self.tokenizer.n_words
105
+
106
+ def get_vocab(self):
107
+ """ Returns vocab as a dict """
108
+ vocab = {self._convert_id_to_token(i): i for i in range(self.vocab_size)}
109
+ vocab.update(self.added_tokens_encoder)
110
+ return vocab
111
+
112
+ def _tokenize(self, text, **kwargs):
113
+ return self.tokenizer.tokenize(text)
114
+
115
+ def _convert_token_to_id(self, token):
116
+ """ Converts a token (str) in an id using the vocab. """
117
+ return self.tokenizer.convert_token_to_id(token)
118
+
119
+ def _convert_id_to_token(self, index):
120
+ """Converts an index (integer) in a token (str) using the vocab."""
121
+ return self.tokenizer.convert_id_to_token(index)
122
+
123
+ def convert_tokens_to_string(self, tokens: List[str]) -> str:
124
+ return self.tokenizer.decode_tokens(tokens)
125
+
126
+ def save_vocabulary(self, save_directory, filename_prefix=None):
127
+ """
128
+ Save the vocabulary and special tokens file to a directory.
129
+
130
+ Args:
131
+ save_directory (`str`):
132
+ The directory in which to save the vocabulary.
133
+ filename_prefix (`str`, *optional*):
134
+ An optional prefix to add to the named of the saved files.
135
+
136
+ Returns:
137
+ `Tuple(str)`: Paths to the files saved.
138
+ """
139
+ if os.path.isdir(save_directory):
140
+ vocab_file = os.path.join(
141
+ save_directory, self.vocab_files_names["vocab_file"]
142
+ )
143
+ else:
144
+ vocab_file = save_directory
145
+
146
+ with open(self.vocab_file, 'rb') as fin:
147
+ proto_str = fin.read()
148
+
149
+ with open(vocab_file, "wb") as writer:
150
+ writer.write(proto_str)
151
+
152
+ return (vocab_file,)
153
+
154
+ def get_prefix_tokens(self):
155
+ prefix_tokens = [self.get_command("[gMASK]"), self.get_command("sop")]
156
+ return prefix_tokens
157
+
158
+ def build_prompt(self, query, history=None):
159
+ if history is None:
160
+ history = []
161
+ prompt = ""
162
+ for i, (old_query, response) in enumerate(history):
163
+ prompt += "[Round {}]\n\n问:{}\n\n答:{}\n\n".format(i + 1, old_query, response)
164
+ prompt += "[Round {}]\n\n问:{}\n\n答:".format(len(history) + 1, query)
165
+ return prompt
166
+
167
+ def build_inputs_with_special_tokens(
168
+ self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
169
+ ) -> List[int]:
170
+ """
171
+ Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
172
+ adding special tokens. A BERT sequence has the following format:
173
+
174
+ - single sequence: `[CLS] X [SEP]`
175
+ - pair of sequences: `[CLS] A [SEP] B [SEP]`
176
+
177
+ Args:
178
+ token_ids_0 (`List[int]`):
179
+ List of IDs to which the special tokens will be added.
180
+ token_ids_1 (`List[int]`, *optional*):
181
+ Optional second list of IDs for sequence pairs.
182
+
183
+ Returns:
184
+ `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
185
+ """
186
+ prefix_tokens = self.get_prefix_tokens()
187
+ token_ids_0 = prefix_tokens + token_ids_0
188
+ if token_ids_1 is not None:
189
+ token_ids_0 = token_ids_0 + token_ids_1 + [self.get_command("<eos>")]
190
+ return token_ids_0
191
+
192
+ def _pad(
193
+ self,
194
+ encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding],
195
+ max_length: Optional[int] = None,
196
+ padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
197
+ pad_to_multiple_of: Optional[int] = None,
198
+ return_attention_mask: Optional[bool] = None,
199
+ ) -> dict:
200
+ """
201
+ Pad encoded inputs (on left/right and up to predefined length or max length in the batch)
202
+
203
+ Args:
204
+ encoded_inputs:
205
+ Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`).
206
+ max_length: maximum length of the returned list and optionally padding length (see below).
207
+ Will truncate by taking into account the special tokens.
208
+ padding_strategy: PaddingStrategy to use for padding.
209
+
210
+ - PaddingStrategy.LONGEST Pad to the longest sequence in the batch
211
+ - PaddingStrategy.MAX_LENGTH: Pad to the max length (default)
212
+ - PaddingStrategy.DO_NOT_PAD: Do not pad
213
+ The tokenizer padding sides are defined in self.padding_side:
214
+
215
+ - 'left': pads on the left of the sequences
216
+ - 'right': pads on the right of the sequences
217
+ pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value.
218
+ This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability
219
+ `>= 7.5` (Volta).
220
+ return_attention_mask:
221
+ (optional) Set to False to avoid returning attention mask (default: set to model specifics)
222
+ """
223
+ # Load from model defaults
224
+ assert self.padding_side == "left"
225
+
226
+ required_input = encoded_inputs[self.model_input_names[0]]
227
+ seq_length = len(required_input)
228
+
229
+ if padding_strategy == PaddingStrategy.LONGEST:
230
+ max_length = len(required_input)
231
+
232
+ if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
233
+ max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of
234
+
235
+ needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(required_input) != max_length
236
+
237
+ # Initialize attention mask if not present.
238
+ if "attention_mask" not in encoded_inputs:
239
+ encoded_inputs["attention_mask"] = [1] * seq_length
240
+
241
+ if "position_ids" not in encoded_inputs:
242
+ encoded_inputs["position_ids"] = list(range(seq_length))
243
+
244
+ if needs_to_be_padded:
245
+ difference = max_length - len(required_input)
246
+
247
+ if "attention_mask" in encoded_inputs:
248
+ encoded_inputs["attention_mask"] = [0] * difference + encoded_inputs["attention_mask"]
249
+ if "position_ids" in encoded_inputs:
250
+ encoded_inputs["position_ids"] = [0] * difference + encoded_inputs["position_ids"]
251
+ encoded_inputs[self.model_input_names[0]] = [self.pad_token_id] * difference + required_input
252
+
253
+ return encoded_inputs
checkpoint-2000/tokenizer.model ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e7dc4c393423b76e4373e5157ddc34803a0189ba96b21ddbb40269d31468a6f2
3
+ size 1018370
checkpoint-2000/tokenizer_config.json ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "auto_map": {
3
+ "AutoTokenizer": [
4
+ "tokenization_chatglm.ChatGLMTokenizer",
5
+ null
6
+ ]
7
+ },
8
+ "clean_up_tokenization_spaces": true,
9
+ "do_lower_case": false,
10
+ "model_max_length": 1000000000000000019884624838656,
11
+ "padding_side": "left",
12
+ "remove_space": false,
13
+ "tokenizer_class": "ChatGLMTokenizer"
14
+ }
checkpoint-2000/trainer_state.json ADDED
@@ -0,0 +1,1216 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "best_metric": null,
3
+ "best_model_checkpoint": null,
4
+ "epoch": 2.7535566773749425,
5
+ "global_step": 2000,
6
+ "is_hyper_param_search": false,
7
+ "is_local_process_zero": true,
8
+ "is_world_process_zero": true,
9
+ "log_history": [
10
+ {
11
+ "epoch": 0.01,
12
+ "learning_rate": 0.00019999739928881725,
13
+ "loss": 2.1906,
14
+ "step": 10
15
+ },
16
+ {
17
+ "epoch": 0.03,
18
+ "learning_rate": 0.00019998959729054295,
19
+ "loss": 1.8919,
20
+ "step": 20
21
+ },
22
+ {
23
+ "epoch": 0.04,
24
+ "learning_rate": 0.00019997659441099206,
25
+ "loss": 1.7855,
26
+ "step": 30
27
+ },
28
+ {
29
+ "epoch": 0.06,
30
+ "learning_rate": 0.00019995839132649917,
31
+ "loss": 1.7235,
32
+ "step": 40
33
+ },
34
+ {
35
+ "epoch": 0.07,
36
+ "learning_rate": 0.0001999349889838836,
37
+ "loss": 1.7027,
38
+ "step": 50
39
+ },
40
+ {
41
+ "epoch": 0.08,
42
+ "learning_rate": 0.00019990638860040006,
43
+ "loss": 1.7118,
44
+ "step": 60
45
+ },
46
+ {
47
+ "epoch": 0.1,
48
+ "learning_rate": 0.00019987259166367533,
49
+ "loss": 1.689,
50
+ "step": 70
51
+ },
52
+ {
53
+ "epoch": 0.11,
54
+ "learning_rate": 0.00019983359993163078,
55
+ "loss": 1.6845,
56
+ "step": 80
57
+ },
58
+ {
59
+ "epoch": 0.12,
60
+ "learning_rate": 0.0001997894154323911,
61
+ "loss": 1.6655,
62
+ "step": 90
63
+ },
64
+ {
65
+ "epoch": 0.14,
66
+ "learning_rate": 0.0001997400404641787,
67
+ "loss": 1.6518,
68
+ "step": 100
69
+ },
70
+ {
71
+ "epoch": 0.15,
72
+ "learning_rate": 0.00019968547759519425,
73
+ "loss": 1.6571,
74
+ "step": 110
75
+ },
76
+ {
77
+ "epoch": 0.17,
78
+ "learning_rate": 0.000199625729663483,
79
+ "loss": 1.6559,
80
+ "step": 120
81
+ },
82
+ {
83
+ "epoch": 0.18,
84
+ "learning_rate": 0.00019956079977678722,
85
+ "loss": 1.6417,
86
+ "step": 130
87
+ },
88
+ {
89
+ "epoch": 0.19,
90
+ "learning_rate": 0.0001994906913123846,
91
+ "loss": 1.6313,
92
+ "step": 140
93
+ },
94
+ {
95
+ "epoch": 0.21,
96
+ "learning_rate": 0.00019941540791691245,
97
+ "loss": 1.6243,
98
+ "step": 150
99
+ },
100
+ {
101
+ "epoch": 0.22,
102
+ "learning_rate": 0.00019933495350617813,
103
+ "loss": 1.6189,
104
+ "step": 160
105
+ },
106
+ {
107
+ "epoch": 0.23,
108
+ "learning_rate": 0.0001992493322649554,
109
+ "loss": 1.628,
110
+ "step": 170
111
+ },
112
+ {
113
+ "epoch": 0.25,
114
+ "learning_rate": 0.00019915854864676664,
115
+ "loss": 1.6149,
116
+ "step": 180
117
+ },
118
+ {
119
+ "epoch": 0.26,
120
+ "learning_rate": 0.00019906260737365122,
121
+ "loss": 1.6032,
122
+ "step": 190
123
+ },
124
+ {
125
+ "epoch": 0.28,
126
+ "learning_rate": 0.00019896151343592008,
127
+ "loss": 1.626,
128
+ "step": 200
129
+ },
130
+ {
131
+ "epoch": 0.29,
132
+ "learning_rate": 0.0001988552720918958,
133
+ "loss": 1.6369,
134
+ "step": 210
135
+ },
136
+ {
137
+ "epoch": 0.3,
138
+ "learning_rate": 0.00019874388886763944,
139
+ "loss": 1.5757,
140
+ "step": 220
141
+ },
142
+ {
143
+ "epoch": 0.32,
144
+ "learning_rate": 0.00019862736955666296,
145
+ "loss": 1.5873,
146
+ "step": 230
147
+ },
148
+ {
149
+ "epoch": 0.33,
150
+ "learning_rate": 0.00019850572021962788,
151
+ "loss": 1.6041,
152
+ "step": 240
153
+ },
154
+ {
155
+ "epoch": 0.34,
156
+ "learning_rate": 0.00019837894718402997,
157
+ "loss": 1.5905,
158
+ "step": 250
159
+ },
160
+ {
161
+ "epoch": 0.36,
162
+ "learning_rate": 0.00019824705704387028,
163
+ "loss": 1.6131,
164
+ "step": 260
165
+ },
166
+ {
167
+ "epoch": 0.37,
168
+ "learning_rate": 0.00019811005665931205,
169
+ "loss": 1.6045,
170
+ "step": 270
171
+ },
172
+ {
173
+ "epoch": 0.39,
174
+ "learning_rate": 0.00019796795315632395,
175
+ "loss": 1.5942,
176
+ "step": 280
177
+ },
178
+ {
179
+ "epoch": 0.4,
180
+ "learning_rate": 0.00019782075392630935,
181
+ "loss": 1.5768,
182
+ "step": 290
183
+ },
184
+ {
185
+ "epoch": 0.41,
186
+ "learning_rate": 0.00019766846662572191,
187
+ "loss": 1.5967,
188
+ "step": 300
189
+ },
190
+ {
191
+ "epoch": 0.43,
192
+ "learning_rate": 0.00019751109917566737,
193
+ "loss": 1.5606,
194
+ "step": 310
195
+ },
196
+ {
197
+ "epoch": 0.44,
198
+ "learning_rate": 0.00019734865976149145,
199
+ "loss": 1.58,
200
+ "step": 320
201
+ },
202
+ {
203
+ "epoch": 0.45,
204
+ "learning_rate": 0.00019718115683235417,
205
+ "loss": 1.5809,
206
+ "step": 330
207
+ },
208
+ {
209
+ "epoch": 0.47,
210
+ "learning_rate": 0.00019700859910079036,
211
+ "loss": 1.5758,
212
+ "step": 340
213
+ },
214
+ {
215
+ "epoch": 0.48,
216
+ "learning_rate": 0.00019683099554225649,
217
+ "loss": 1.57,
218
+ "step": 350
219
+ },
220
+ {
221
+ "epoch": 0.5,
222
+ "learning_rate": 0.0001966483553946637,
223
+ "loss": 1.5783,
224
+ "step": 360
225
+ },
226
+ {
227
+ "epoch": 0.51,
228
+ "learning_rate": 0.00019646068815789755,
229
+ "loss": 1.5699,
230
+ "step": 370
231
+ },
232
+ {
233
+ "epoch": 0.52,
234
+ "learning_rate": 0.00019626800359332362,
235
+ "loss": 1.5767,
236
+ "step": 380
237
+ },
238
+ {
239
+ "epoch": 0.54,
240
+ "learning_rate": 0.00019607031172327996,
241
+ "loss": 1.5635,
242
+ "step": 390
243
+ },
244
+ {
245
+ "epoch": 0.55,
246
+ "learning_rate": 0.00019586762283055573,
247
+ "loss": 1.5763,
248
+ "step": 400
249
+ },
250
+ {
251
+ "epoch": 0.56,
252
+ "learning_rate": 0.0001956599474578563,
253
+ "loss": 1.5719,
254
+ "step": 410
255
+ },
256
+ {
257
+ "epoch": 0.58,
258
+ "learning_rate": 0.00019544729640725498,
259
+ "loss": 1.5712,
260
+ "step": 420
261
+ },
262
+ {
263
+ "epoch": 0.59,
264
+ "learning_rate": 0.00019522968073963106,
265
+ "loss": 1.5699,
266
+ "step": 430
267
+ },
268
+ {
269
+ "epoch": 0.61,
270
+ "learning_rate": 0.00019500711177409454,
271
+ "loss": 1.539,
272
+ "step": 440
273
+ },
274
+ {
275
+ "epoch": 0.62,
276
+ "learning_rate": 0.0001947796010873974,
277
+ "loss": 1.5515,
278
+ "step": 450
279
+ },
280
+ {
281
+ "epoch": 0.63,
282
+ "learning_rate": 0.00019454716051333135,
283
+ "loss": 1.5691,
284
+ "step": 460
285
+ },
286
+ {
287
+ "epoch": 0.65,
288
+ "learning_rate": 0.0001943098021421124,
289
+ "loss": 1.573,
290
+ "step": 470
291
+ },
292
+ {
293
+ "epoch": 0.66,
294
+ "learning_rate": 0.00019406753831975203,
295
+ "loss": 1.5532,
296
+ "step": 480
297
+ },
298
+ {
299
+ "epoch": 0.67,
300
+ "learning_rate": 0.00019382038164741477,
301
+ "loss": 1.5778,
302
+ "step": 490
303
+ },
304
+ {
305
+ "epoch": 0.69,
306
+ "learning_rate": 0.0001935683449807631,
307
+ "loss": 1.5514,
308
+ "step": 500
309
+ },
310
+ {
311
+ "epoch": 0.7,
312
+ "learning_rate": 0.00019331144142928854,
313
+ "loss": 1.5493,
314
+ "step": 510
315
+ },
316
+ {
317
+ "epoch": 0.72,
318
+ "learning_rate": 0.00019304968435562993,
319
+ "loss": 1.5698,
320
+ "step": 520
321
+ },
322
+ {
323
+ "epoch": 0.73,
324
+ "learning_rate": 0.00019278308737487822,
325
+ "loss": 1.5699,
326
+ "step": 530
327
+ },
328
+ {
329
+ "epoch": 0.74,
330
+ "learning_rate": 0.0001925116643538684,
331
+ "loss": 1.554,
332
+ "step": 540
333
+ },
334
+ {
335
+ "epoch": 0.76,
336
+ "learning_rate": 0.00019223542941045817,
337
+ "loss": 1.5527,
338
+ "step": 550
339
+ },
340
+ {
341
+ "epoch": 0.77,
342
+ "learning_rate": 0.00019195439691279363,
343
+ "loss": 1.5517,
344
+ "step": 560
345
+ },
346
+ {
347
+ "epoch": 0.78,
348
+ "learning_rate": 0.00019166858147856203,
349
+ "loss": 1.5673,
350
+ "step": 570
351
+ },
352
+ {
353
+ "epoch": 0.8,
354
+ "learning_rate": 0.00019137799797423126,
355
+ "loss": 1.5469,
356
+ "step": 580
357
+ },
358
+ {
359
+ "epoch": 0.81,
360
+ "learning_rate": 0.00019108266151427673,
361
+ "loss": 1.5499,
362
+ "step": 590
363
+ },
364
+ {
365
+ "epoch": 0.83,
366
+ "learning_rate": 0.00019078258746039507,
367
+ "loss": 1.5502,
368
+ "step": 600
369
+ },
370
+ {
371
+ "epoch": 0.84,
372
+ "learning_rate": 0.00019047779142070527,
373
+ "loss": 1.5402,
374
+ "step": 610
375
+ },
376
+ {
377
+ "epoch": 0.85,
378
+ "learning_rate": 0.0001901682892489367,
379
+ "loss": 1.5444,
380
+ "step": 620
381
+ },
382
+ {
383
+ "epoch": 0.87,
384
+ "learning_rate": 0.00018985409704360456,
385
+ "loss": 1.5524,
386
+ "step": 630
387
+ },
388
+ {
389
+ "epoch": 0.88,
390
+ "learning_rate": 0.00018953523114717245,
391
+ "loss": 1.5338,
392
+ "step": 640
393
+ },
394
+ {
395
+ "epoch": 0.89,
396
+ "learning_rate": 0.00018921170814520247,
397
+ "loss": 1.5416,
398
+ "step": 650
399
+ },
400
+ {
401
+ "epoch": 0.91,
402
+ "learning_rate": 0.00018888354486549237,
403
+ "loss": 1.5524,
404
+ "step": 660
405
+ },
406
+ {
407
+ "epoch": 0.92,
408
+ "learning_rate": 0.00018855075837720034,
409
+ "loss": 1.544,
410
+ "step": 670
411
+ },
412
+ {
413
+ "epoch": 0.94,
414
+ "learning_rate": 0.0001882133659899573,
415
+ "loss": 1.5317,
416
+ "step": 680
417
+ },
418
+ {
419
+ "epoch": 0.95,
420
+ "learning_rate": 0.0001878713852529663,
421
+ "loss": 1.5516,
422
+ "step": 690
423
+ },
424
+ {
425
+ "epoch": 0.96,
426
+ "learning_rate": 0.00018752483395408987,
427
+ "loss": 1.5479,
428
+ "step": 700
429
+ },
430
+ {
431
+ "epoch": 0.98,
432
+ "learning_rate": 0.00018717373011892474,
433
+ "loss": 1.5328,
434
+ "step": 710
435
+ },
436
+ {
437
+ "epoch": 0.99,
438
+ "learning_rate": 0.0001868180920098644,
439
+ "loss": 1.524,
440
+ "step": 720
441
+ },
442
+ {
443
+ "epoch": 1.01,
444
+ "learning_rate": 0.00018645793812514894,
445
+ "loss": 1.4712,
446
+ "step": 730
447
+ },
448
+ {
449
+ "epoch": 1.02,
450
+ "learning_rate": 0.0001860932871979031,
451
+ "loss": 1.3509,
452
+ "step": 740
453
+ },
454
+ {
455
+ "epoch": 1.03,
456
+ "learning_rate": 0.00018572415819516174,
457
+ "loss": 1.3227,
458
+ "step": 750
459
+ },
460
+ {
461
+ "epoch": 1.05,
462
+ "learning_rate": 0.00018535057031688335,
463
+ "loss": 1.3216,
464
+ "step": 760
465
+ },
466
+ {
467
+ "epoch": 1.06,
468
+ "learning_rate": 0.00018497254299495146,
469
+ "loss": 1.3022,
470
+ "step": 770
471
+ },
472
+ {
473
+ "epoch": 1.07,
474
+ "learning_rate": 0.00018459009589216364,
475
+ "loss": 1.308,
476
+ "step": 780
477
+ },
478
+ {
479
+ "epoch": 1.09,
480
+ "learning_rate": 0.00018420324890120916,
481
+ "loss": 1.3295,
482
+ "step": 790
483
+ },
484
+ {
485
+ "epoch": 1.1,
486
+ "learning_rate": 0.0001838120221436338,
487
+ "loss": 1.3249,
488
+ "step": 800
489
+ },
490
+ {
491
+ "epoch": 1.12,
492
+ "learning_rate": 0.00018341643596879367,
493
+ "loss": 1.3109,
494
+ "step": 810
495
+ },
496
+ {
497
+ "epoch": 1.13,
498
+ "learning_rate": 0.00018301651095279655,
499
+ "loss": 1.328,
500
+ "step": 820
501
+ },
502
+ {
503
+ "epoch": 1.14,
504
+ "learning_rate": 0.00018261226789743172,
505
+ "loss": 1.3215,
506
+ "step": 830
507
+ },
508
+ {
509
+ "epoch": 1.16,
510
+ "learning_rate": 0.00018220372782908777,
511
+ "loss": 1.3378,
512
+ "step": 840
513
+ },
514
+ {
515
+ "epoch": 1.17,
516
+ "learning_rate": 0.00018179091199765926,
517
+ "loss": 1.3289,
518
+ "step": 850
519
+ },
520
+ {
521
+ "epoch": 1.18,
522
+ "learning_rate": 0.00018137384187544116,
523
+ "loss": 1.3362,
524
+ "step": 860
525
+ },
526
+ {
527
+ "epoch": 1.2,
528
+ "learning_rate": 0.00018095253915601206,
529
+ "loss": 1.3263,
530
+ "step": 870
531
+ },
532
+ {
533
+ "epoch": 1.21,
534
+ "learning_rate": 0.00018052702575310588,
535
+ "loss": 1.3238,
536
+ "step": 880
537
+ },
538
+ {
539
+ "epoch": 1.23,
540
+ "learning_rate": 0.00018009732379947188,
541
+ "loss": 1.3278,
542
+ "step": 890
543
+ },
544
+ {
545
+ "epoch": 1.24,
546
+ "learning_rate": 0.0001796634556457236,
547
+ "loss": 1.3277,
548
+ "step": 900
549
+ },
550
+ {
551
+ "epoch": 1.25,
552
+ "learning_rate": 0.00017922544385917628,
553
+ "loss": 1.3332,
554
+ "step": 910
555
+ },
556
+ {
557
+ "epoch": 1.27,
558
+ "learning_rate": 0.00017878331122267284,
559
+ "loss": 1.3268,
560
+ "step": 920
561
+ },
562
+ {
563
+ "epoch": 1.28,
564
+ "learning_rate": 0.00017833708073339922,
565
+ "loss": 1.3316,
566
+ "step": 930
567
+ },
568
+ {
569
+ "epoch": 1.29,
570
+ "learning_rate": 0.00017788677560168784,
571
+ "loss": 1.337,
572
+ "step": 940
573
+ },
574
+ {
575
+ "epoch": 1.31,
576
+ "learning_rate": 0.0001774324192498105,
577
+ "loss": 1.3307,
578
+ "step": 950
579
+ },
580
+ {
581
+ "epoch": 1.32,
582
+ "learning_rate": 0.0001769740353107602,
583
+ "loss": 1.3295,
584
+ "step": 960
585
+ },
586
+ {
587
+ "epoch": 1.34,
588
+ "learning_rate": 0.0001765116476270216,
589
+ "loss": 1.3148,
590
+ "step": 970
591
+ },
592
+ {
593
+ "epoch": 1.35,
594
+ "learning_rate": 0.00017604528024933115,
595
+ "loss": 1.3269,
596
+ "step": 980
597
+ },
598
+ {
599
+ "epoch": 1.36,
600
+ "learning_rate": 0.00017557495743542585,
601
+ "loss": 1.3488,
602
+ "step": 990
603
+ },
604
+ {
605
+ "epoch": 1.38,
606
+ "learning_rate": 0.00017510070364878177,
607
+ "loss": 1.3377,
608
+ "step": 1000
609
+ },
610
+ {
611
+ "epoch": 1.39,
612
+ "learning_rate": 0.0001746225435573415,
613
+ "loss": 1.3396,
614
+ "step": 1010
615
+ },
616
+ {
617
+ "epoch": 1.4,
618
+ "learning_rate": 0.0001741405020322309,
619
+ "loss": 1.3375,
620
+ "step": 1020
621
+ },
622
+ {
623
+ "epoch": 1.42,
624
+ "learning_rate": 0.00017365460414646574,
625
+ "loss": 1.326,
626
+ "step": 1030
627
+ },
628
+ {
629
+ "epoch": 1.43,
630
+ "learning_rate": 0.00017316487517364721,
631
+ "loss": 1.3297,
632
+ "step": 1040
633
+ },
634
+ {
635
+ "epoch": 1.45,
636
+ "learning_rate": 0.00017267134058664775,
637
+ "loss": 1.3429,
638
+ "step": 1050
639
+ },
640
+ {
641
+ "epoch": 1.46,
642
+ "learning_rate": 0.00017217402605628572,
643
+ "loss": 1.3339,
644
+ "step": 1060
645
+ },
646
+ {
647
+ "epoch": 1.47,
648
+ "learning_rate": 0.00017167295744999027,
649
+ "loss": 1.3375,
650
+ "step": 1070
651
+ },
652
+ {
653
+ "epoch": 1.49,
654
+ "learning_rate": 0.00017116816083045602,
655
+ "loss": 1.3161,
656
+ "step": 1080
657
+ },
658
+ {
659
+ "epoch": 1.5,
660
+ "learning_rate": 0.00017065966245428723,
661
+ "loss": 1.3281,
662
+ "step": 1090
663
+ },
664
+ {
665
+ "epoch": 1.51,
666
+ "learning_rate": 0.00017014748877063214,
667
+ "loss": 1.322,
668
+ "step": 1100
669
+ },
670
+ {
671
+ "epoch": 1.53,
672
+ "learning_rate": 0.0001696316664198073,
673
+ "loss": 1.3344,
674
+ "step": 1110
675
+ },
676
+ {
677
+ "epoch": 1.54,
678
+ "learning_rate": 0.00016911222223191182,
679
+ "loss": 1.3607,
680
+ "step": 1120
681
+ },
682
+ {
683
+ "epoch": 1.56,
684
+ "learning_rate": 0.00016858918322543186,
685
+ "loss": 1.34,
686
+ "step": 1130
687
+ },
688
+ {
689
+ "epoch": 1.57,
690
+ "learning_rate": 0.00016806257660583534,
691
+ "loss": 1.3324,
692
+ "step": 1140
693
+ },
694
+ {
695
+ "epoch": 1.58,
696
+ "learning_rate": 0.00016753242976415666,
697
+ "loss": 1.3249,
698
+ "step": 1150
699
+ },
700
+ {
701
+ "epoch": 1.6,
702
+ "learning_rate": 0.00016699877027557226,
703
+ "loss": 1.3363,
704
+ "step": 1160
705
+ },
706
+ {
707
+ "epoch": 1.61,
708
+ "learning_rate": 0.00016646162589796615,
709
+ "loss": 1.3403,
710
+ "step": 1170
711
+ },
712
+ {
713
+ "epoch": 1.62,
714
+ "learning_rate": 0.0001659210245704861,
715
+ "loss": 1.3283,
716
+ "step": 1180
717
+ },
718
+ {
719
+ "epoch": 1.64,
720
+ "learning_rate": 0.00016537699441209047,
721
+ "loss": 1.3258,
722
+ "step": 1190
723
+ },
724
+ {
725
+ "epoch": 1.65,
726
+ "learning_rate": 0.0001648295637200856,
727
+ "loss": 1.3258,
728
+ "step": 1200
729
+ },
730
+ {
731
+ "epoch": 1.67,
732
+ "learning_rate": 0.00016427876096865394,
733
+ "loss": 1.3581,
734
+ "step": 1210
735
+ },
736
+ {
737
+ "epoch": 1.68,
738
+ "learning_rate": 0.00016372461480737297,
739
+ "loss": 1.3337,
740
+ "step": 1220
741
+ },
742
+ {
743
+ "epoch": 1.69,
744
+ "learning_rate": 0.00016316715405972508,
745
+ "loss": 1.3432,
746
+ "step": 1230
747
+ },
748
+ {
749
+ "epoch": 1.71,
750
+ "learning_rate": 0.0001626064077215983,
751
+ "loss": 1.3385,
752
+ "step": 1240
753
+ },
754
+ {
755
+ "epoch": 1.72,
756
+ "learning_rate": 0.00016204240495977805,
757
+ "loss": 1.3474,
758
+ "step": 1250
759
+ },
760
+ {
761
+ "epoch": 1.73,
762
+ "learning_rate": 0.0001614751751104301,
763
+ "loss": 1.3198,
764
+ "step": 1260
765
+ },
766
+ {
767
+ "epoch": 1.75,
768
+ "learning_rate": 0.00016090474767757474,
769
+ "loss": 1.3407,
770
+ "step": 1270
771
+ },
772
+ {
773
+ "epoch": 1.76,
774
+ "learning_rate": 0.00016033115233155202,
775
+ "loss": 1.3203,
776
+ "step": 1280
777
+ },
778
+ {
779
+ "epoch": 1.78,
780
+ "learning_rate": 0.00015975441890747855,
781
+ "loss": 1.3311,
782
+ "step": 1290
783
+ },
784
+ {
785
+ "epoch": 1.79,
786
+ "learning_rate": 0.00015917457740369565,
787
+ "loss": 1.3271,
788
+ "step": 1300
789
+ },
790
+ {
791
+ "epoch": 1.8,
792
+ "learning_rate": 0.000158591657980209,
793
+ "loss": 1.3432,
794
+ "step": 1310
795
+ },
796
+ {
797
+ "epoch": 1.82,
798
+ "learning_rate": 0.00015800569095711982,
799
+ "loss": 1.3302,
800
+ "step": 1320
801
+ },
802
+ {
803
+ "epoch": 1.83,
804
+ "learning_rate": 0.00015741670681304796,
805
+ "loss": 1.3326,
806
+ "step": 1330
807
+ },
808
+ {
809
+ "epoch": 1.84,
810
+ "learning_rate": 0.00015682473618354635,
811
+ "loss": 1.3358,
812
+ "step": 1340
813
+ },
814
+ {
815
+ "epoch": 1.86,
816
+ "learning_rate": 0.0001562298098595078,
817
+ "loss": 1.3343,
818
+ "step": 1350
819
+ },
820
+ {
821
+ "epoch": 1.87,
822
+ "learning_rate": 0.0001556319587855631,
823
+ "loss": 1.3214,
824
+ "step": 1360
825
+ },
826
+ {
827
+ "epoch": 1.89,
828
+ "learning_rate": 0.0001550312140584718,
829
+ "loss": 1.3342,
830
+ "step": 1370
831
+ },
832
+ {
833
+ "epoch": 1.9,
834
+ "learning_rate": 0.00015442760692550443,
835
+ "loss": 1.3253,
836
+ "step": 1380
837
+ },
838
+ {
839
+ "epoch": 1.91,
840
+ "learning_rate": 0.0001538211687828174,
841
+ "loss": 1.3494,
842
+ "step": 1390
843
+ },
844
+ {
845
+ "epoch": 1.93,
846
+ "learning_rate": 0.00015321193117381996,
847
+ "loss": 1.324,
848
+ "step": 1400
849
+ },
850
+ {
851
+ "epoch": 1.94,
852
+ "learning_rate": 0.00015259992578753334,
853
+ "loss": 1.3267,
854
+ "step": 1410
855
+ },
856
+ {
857
+ "epoch": 1.96,
858
+ "learning_rate": 0.00015198518445694255,
859
+ "loss": 1.343,
860
+ "step": 1420
861
+ },
862
+ {
863
+ "epoch": 1.97,
864
+ "learning_rate": 0.00015136773915734066,
865
+ "loss": 1.3209,
866
+ "step": 1430
867
+ },
868
+ {
869
+ "epoch": 1.98,
870
+ "learning_rate": 0.00015074762200466556,
871
+ "loss": 1.3224,
872
+ "step": 1440
873
+ },
874
+ {
875
+ "epoch": 2.0,
876
+ "learning_rate": 0.0001501248652538296,
877
+ "loss": 1.3481,
878
+ "step": 1450
879
+ },
880
+ {
881
+ "epoch": 2.01,
882
+ "learning_rate": 0.00014949950129704162,
883
+ "loss": 1.1061,
884
+ "step": 1460
885
+ },
886
+ {
887
+ "epoch": 2.02,
888
+ "learning_rate": 0.00014887156266212237,
889
+ "loss": 0.9927,
890
+ "step": 1470
891
+ },
892
+ {
893
+ "epoch": 2.04,
894
+ "learning_rate": 0.00014824108201081247,
895
+ "loss": 0.9946,
896
+ "step": 1480
897
+ },
898
+ {
899
+ "epoch": 2.05,
900
+ "learning_rate": 0.00014760809213707344,
901
+ "loss": 0.9764,
902
+ "step": 1490
903
+ },
904
+ {
905
+ "epoch": 2.07,
906
+ "learning_rate": 0.00014697262596538227,
907
+ "loss": 0.9851,
908
+ "step": 1500
909
+ },
910
+ {
911
+ "epoch": 2.08,
912
+ "learning_rate": 0.00014633471654901842,
913
+ "loss": 0.9941,
914
+ "step": 1510
915
+ },
916
+ {
917
+ "epoch": 2.09,
918
+ "learning_rate": 0.000145694397068345,
919
+ "loss": 0.9875,
920
+ "step": 1520
921
+ },
922
+ {
923
+ "epoch": 2.11,
924
+ "learning_rate": 0.0001450517008290827,
925
+ "loss": 0.9828,
926
+ "step": 1530
927
+ },
928
+ {
929
+ "epoch": 2.12,
930
+ "learning_rate": 0.00014440666126057744,
931
+ "loss": 0.9913,
932
+ "step": 1540
933
+ },
934
+ {
935
+ "epoch": 2.13,
936
+ "learning_rate": 0.00014375931191406159,
937
+ "loss": 0.9881,
938
+ "step": 1550
939
+ },
940
+ {
941
+ "epoch": 2.15,
942
+ "learning_rate": 0.00014310968646090883,
943
+ "loss": 0.9866,
944
+ "step": 1560
945
+ },
946
+ {
947
+ "epoch": 2.16,
948
+ "learning_rate": 0.0001424578186908828,
949
+ "loss": 0.9934,
950
+ "step": 1570
951
+ },
952
+ {
953
+ "epoch": 2.18,
954
+ "learning_rate": 0.0001418037425103795,
955
+ "loss": 0.9963,
956
+ "step": 1580
957
+ },
958
+ {
959
+ "epoch": 2.19,
960
+ "learning_rate": 0.00014114749194066363,
961
+ "loss": 0.9791,
962
+ "step": 1590
963
+ },
964
+ {
965
+ "epoch": 2.2,
966
+ "learning_rate": 0.00014048910111609915,
967
+ "loss": 0.9851,
968
+ "step": 1600
969
+ },
970
+ {
971
+ "epoch": 2.22,
972
+ "learning_rate": 0.0001398286042823736,
973
+ "loss": 0.9829,
974
+ "step": 1610
975
+ },
976
+ {
977
+ "epoch": 2.23,
978
+ "learning_rate": 0.00013916603579471705,
979
+ "loss": 0.9995,
980
+ "step": 1620
981
+ },
982
+ {
983
+ "epoch": 2.24,
984
+ "learning_rate": 0.00013850143011611497,
985
+ "loss": 0.9843,
986
+ "step": 1630
987
+ },
988
+ {
989
+ "epoch": 2.26,
990
+ "learning_rate": 0.0001378348218155158,
991
+ "loss": 1.0025,
992
+ "step": 1640
993
+ },
994
+ {
995
+ "epoch": 2.27,
996
+ "learning_rate": 0.00013716624556603274,
997
+ "loss": 1.0135,
998
+ "step": 1650
999
+ },
1000
+ {
1001
+ "epoch": 2.29,
1002
+ "learning_rate": 0.00013649573614314044,
1003
+ "loss": 1.0007,
1004
+ "step": 1660
1005
+ },
1006
+ {
1007
+ "epoch": 2.3,
1008
+ "learning_rate": 0.00013582332842286592,
1009
+ "loss": 1.0117,
1010
+ "step": 1670
1011
+ },
1012
+ {
1013
+ "epoch": 2.31,
1014
+ "learning_rate": 0.00013514905737997473,
1015
+ "loss": 1.0068,
1016
+ "step": 1680
1017
+ },
1018
+ {
1019
+ "epoch": 2.33,
1020
+ "learning_rate": 0.0001344729580861517,
1021
+ "loss": 0.998,
1022
+ "step": 1690
1023
+ },
1024
+ {
1025
+ "epoch": 2.34,
1026
+ "learning_rate": 0.0001337950657081768,
1027
+ "loss": 0.9995,
1028
+ "step": 1700
1029
+ },
1030
+ {
1031
+ "epoch": 2.35,
1032
+ "learning_rate": 0.00013311541550609565,
1033
+ "loss": 0.9975,
1034
+ "step": 1710
1035
+ },
1036
+ {
1037
+ "epoch": 2.37,
1038
+ "learning_rate": 0.00013243404283138597,
1039
+ "loss": 1.014,
1040
+ "step": 1720
1041
+ },
1042
+ {
1043
+ "epoch": 2.38,
1044
+ "learning_rate": 0.0001317509831251184,
1045
+ "loss": 0.9954,
1046
+ "step": 1730
1047
+ },
1048
+ {
1049
+ "epoch": 2.4,
1050
+ "learning_rate": 0.00013106627191611332,
1051
+ "loss": 1.0051,
1052
+ "step": 1740
1053
+ },
1054
+ {
1055
+ "epoch": 2.41,
1056
+ "learning_rate": 0.00013037994481909264,
1057
+ "loss": 1.0156,
1058
+ "step": 1750
1059
+ },
1060
+ {
1061
+ "epoch": 2.42,
1062
+ "learning_rate": 0.0001296920375328275,
1063
+ "loss": 1.0155,
1064
+ "step": 1760
1065
+ },
1066
+ {
1067
+ "epoch": 2.44,
1068
+ "learning_rate": 0.00012900258583828137,
1069
+ "loss": 1.0223,
1070
+ "step": 1770
1071
+ },
1072
+ {
1073
+ "epoch": 2.45,
1074
+ "learning_rate": 0.00012831162559674887,
1075
+ "loss": 1.0181,
1076
+ "step": 1780
1077
+ },
1078
+ {
1079
+ "epoch": 2.46,
1080
+ "learning_rate": 0.00012761919274799054,
1081
+ "loss": 1.0023,
1082
+ "step": 1790
1083
+ },
1084
+ {
1085
+ "epoch": 2.48,
1086
+ "learning_rate": 0.00012692532330836346,
1087
+ "loss": 0.9917,
1088
+ "step": 1800
1089
+ },
1090
+ {
1091
+ "epoch": 2.49,
1092
+ "learning_rate": 0.0001262300533689478,
1093
+ "loss": 1.0132,
1094
+ "step": 1810
1095
+ },
1096
+ {
1097
+ "epoch": 2.51,
1098
+ "learning_rate": 0.00012553341909366978,
1099
+ "loss": 1.0136,
1100
+ "step": 1820
1101
+ },
1102
+ {
1103
+ "epoch": 2.52,
1104
+ "learning_rate": 0.0001248354567174203,
1105
+ "loss": 0.9961,
1106
+ "step": 1830
1107
+ },
1108
+ {
1109
+ "epoch": 2.53,
1110
+ "learning_rate": 0.00012413620254417057,
1111
+ "loss": 1.0244,
1112
+ "step": 1840
1113
+ },
1114
+ {
1115
+ "epoch": 2.55,
1116
+ "learning_rate": 0.0001234356929450835,
1117
+ "loss": 1.0126,
1118
+ "step": 1850
1119
+ },
1120
+ {
1121
+ "epoch": 2.56,
1122
+ "learning_rate": 0.00012273396435662212,
1123
+ "loss": 1.0121,
1124
+ "step": 1860
1125
+ },
1126
+ {
1127
+ "epoch": 2.57,
1128
+ "learning_rate": 0.00012203105327865407,
1129
+ "loss": 1.0293,
1130
+ "step": 1870
1131
+ },
1132
+ {
1133
+ "epoch": 2.59,
1134
+ "learning_rate": 0.00012132699627255347,
1135
+ "loss": 1.0104,
1136
+ "step": 1880
1137
+ },
1138
+ {
1139
+ "epoch": 2.6,
1140
+ "learning_rate": 0.00012062182995929882,
1141
+ "loss": 1.0119,
1142
+ "step": 1890
1143
+ },
1144
+ {
1145
+ "epoch": 2.62,
1146
+ "learning_rate": 0.00011991559101756852,
1147
+ "loss": 1.0047,
1148
+ "step": 1900
1149
+ },
1150
+ {
1151
+ "epoch": 2.63,
1152
+ "learning_rate": 0.00011920831618183282,
1153
+ "loss": 1.0095,
1154
+ "step": 1910
1155
+ },
1156
+ {
1157
+ "epoch": 2.64,
1158
+ "learning_rate": 0.00011850004224044315,
1159
+ "loss": 1.0112,
1160
+ "step": 1920
1161
+ },
1162
+ {
1163
+ "epoch": 2.66,
1164
+ "learning_rate": 0.0001177908060337188,
1165
+ "loss": 1.0131,
1166
+ "step": 1930
1167
+ },
1168
+ {
1169
+ "epoch": 2.67,
1170
+ "learning_rate": 0.00011708064445203042,
1171
+ "loss": 1.0082,
1172
+ "step": 1940
1173
+ },
1174
+ {
1175
+ "epoch": 2.68,
1176
+ "learning_rate": 0.00011636959443388132,
1177
+ "loss": 1.0246,
1178
+ "step": 1950
1179
+ },
1180
+ {
1181
+ "epoch": 2.7,
1182
+ "learning_rate": 0.00011565769296398618,
1183
+ "loss": 1.0102,
1184
+ "step": 1960
1185
+ },
1186
+ {
1187
+ "epoch": 2.71,
1188
+ "learning_rate": 0.00011494497707134731,
1189
+ "loss": 1.0103,
1190
+ "step": 1970
1191
+ },
1192
+ {
1193
+ "epoch": 2.73,
1194
+ "learning_rate": 0.00011423148382732853,
1195
+ "loss": 1.0081,
1196
+ "step": 1980
1197
+ },
1198
+ {
1199
+ "epoch": 2.74,
1200
+ "learning_rate": 0.000113517250343727,
1201
+ "loss": 1.0092,
1202
+ "step": 1990
1203
+ },
1204
+ {
1205
+ "epoch": 2.75,
1206
+ "learning_rate": 0.0001128023137708429,
1207
+ "loss": 0.9983,
1208
+ "step": 2000
1209
+ }
1210
+ ],
1211
+ "max_steps": 4356,
1212
+ "num_train_epochs": 6,
1213
+ "total_flos": 2.4195737582441595e+18,
1214
+ "trial_name": null,
1215
+ "trial_params": null
1216
+ }
checkpoint-2000/training_args.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:326a38cf0f9916b8ede07877300460617478f67a798b1bed6b7d9bf33cb7066b
3
+ size 3289
checkpoint-3000/config.json ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "/home/ubuntu/chatglm2-6b",
3
+ "add_bias_linear": false,
4
+ "add_qkv_bias": true,
5
+ "apply_query_key_layer_scaling": true,
6
+ "apply_residual_connection_post_layernorm": false,
7
+ "architectures": [
8
+ "ChatGLMForConditionalGeneration"
9
+ ],
10
+ "attention_dropout": 0.0,
11
+ "attention_softmax_in_fp32": true,
12
+ "auto_map": {
13
+ "AutoConfig": "configuration_chatglm.ChatGLMConfig",
14
+ "AutoModel": "modeling_chatglm.ChatGLMForConditionalGeneration",
15
+ "AutoModelForCausalLM": "modeling_chatglm.ChatGLMForConditionalGeneration",
16
+ "AutoModelForSeq2SeqLM": "modeling_chatglm.ChatGLMForConditionalGeneration"
17
+ },
18
+ "bias_dropout_fusion": true,
19
+ "eos_token_id": 2,
20
+ "ffn_hidden_size": 13696,
21
+ "fp32_residual_connection": false,
22
+ "hidden_dropout": 0.0,
23
+ "hidden_size": 4096,
24
+ "kv_channels": 128,
25
+ "layernorm_epsilon": 1e-05,
26
+ "model_type": "chatglm",
27
+ "multi_query_attention": true,
28
+ "multi_query_group_num": 2,
29
+ "num_attention_heads": 32,
30
+ "num_layers": 28,
31
+ "original_rope": true,
32
+ "pad_token_id": 0,
33
+ "padded_vocab_size": 65024,
34
+ "post_layer_norm": true,
35
+ "pre_seq_len": null,
36
+ "prefix_projection": false,
37
+ "quantization_bit": 0,
38
+ "rmsnorm": true,
39
+ "seq_length": 32768,
40
+ "tie_word_embeddings": false,
41
+ "torch_dtype": "float16",
42
+ "transformers_version": "4.30.2",
43
+ "use_cache": true,
44
+ "vocab_size": 65024
45
+ }
checkpoint-3000/configuration_chatglm.py ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import PretrainedConfig
2
+
3
+
4
+ class ChatGLMConfig(PretrainedConfig):
5
+ model_type = "chatglm"
6
+ def __init__(
7
+ self,
8
+ num_layers=28,
9
+ padded_vocab_size=65024,
10
+ hidden_size=4096,
11
+ ffn_hidden_size=13696,
12
+ kv_channels=128,
13
+ num_attention_heads=32,
14
+ seq_length=2048,
15
+ hidden_dropout=0.0,
16
+ attention_dropout=0.0,
17
+ layernorm_epsilon=1e-5,
18
+ rmsnorm=True,
19
+ apply_residual_connection_post_layernorm=False,
20
+ post_layer_norm=True,
21
+ add_bias_linear=False,
22
+ add_qkv_bias=False,
23
+ bias_dropout_fusion=True,
24
+ multi_query_attention=False,
25
+ multi_query_group_num=1,
26
+ apply_query_key_layer_scaling=True,
27
+ attention_softmax_in_fp32=True,
28
+ fp32_residual_connection=False,
29
+ quantization_bit=0,
30
+ pre_seq_len=None,
31
+ prefix_projection=False,
32
+ **kwargs
33
+ ):
34
+ self.num_layers = num_layers
35
+ self.vocab_size = padded_vocab_size
36
+ self.padded_vocab_size = padded_vocab_size
37
+ self.hidden_size = hidden_size
38
+ self.ffn_hidden_size = ffn_hidden_size
39
+ self.kv_channels = kv_channels
40
+ self.num_attention_heads = num_attention_heads
41
+ self.seq_length = seq_length
42
+ self.hidden_dropout = hidden_dropout
43
+ self.attention_dropout = attention_dropout
44
+ self.layernorm_epsilon = layernorm_epsilon
45
+ self.rmsnorm = rmsnorm
46
+ self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm
47
+ self.post_layer_norm = post_layer_norm
48
+ self.add_bias_linear = add_bias_linear
49
+ self.add_qkv_bias = add_qkv_bias
50
+ self.bias_dropout_fusion = bias_dropout_fusion
51
+ self.multi_query_attention = multi_query_attention
52
+ self.multi_query_group_num = multi_query_group_num
53
+ self.apply_query_key_layer_scaling = apply_query_key_layer_scaling
54
+ self.attention_softmax_in_fp32 = attention_softmax_in_fp32
55
+ self.fp32_residual_connection = fp32_residual_connection
56
+ self.quantization_bit = quantization_bit
57
+ self.pre_seq_len = pre_seq_len
58
+ self.prefix_projection = prefix_projection
59
+ super().__init__(**kwargs)
checkpoint-3000/finetuning_args.json ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "finetuning_type": "freeze",
3
+ "lora_alpha": 9.0,
4
+ "lora_dropout": 0.05,
5
+ "lora_rank": 100,
6
+ "lora_target": [
7
+ "query_key_value"
8
+ ],
9
+ "name_module_trainable": "mlp",
10
+ "num_layer_trainable": 1,
11
+ "pre_seq_len": 64,
12
+ "prefix_projection": false
13
+ }
checkpoint-3000/generation_config.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "eos_token_id": 2,
4
+ "pad_token_id": 0,
5
+ "transformers_version": "4.30.2"
6
+ }
checkpoint-3000/modeling_chatglm.py ADDED
@@ -0,0 +1,1193 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """ PyTorch ChatGLM model. """
2
+
3
+ import math
4
+ import copy
5
+ import warnings
6
+ import re
7
+ import sys
8
+
9
+ import torch
10
+ import torch.utils.checkpoint
11
+ import torch.nn.functional as F
12
+ from torch import nn
13
+ from torch.nn import CrossEntropyLoss, LayerNorm
14
+ from torch.nn.utils import skip_init
15
+ from typing import Optional, Tuple, Union, List, Callable, Dict, Any
16
+
17
+ from transformers.modeling_outputs import (
18
+ BaseModelOutputWithPast,
19
+ CausalLMOutputWithPast,
20
+ )
21
+ from transformers.modeling_utils import PreTrainedModel
22
+ from transformers.utils import logging
23
+ from transformers.generation.logits_process import LogitsProcessor
24
+ from transformers.generation.utils import LogitsProcessorList, StoppingCriteriaList, GenerationConfig, ModelOutput
25
+
26
+ from .configuration_chatglm import ChatGLMConfig
27
+
28
+ # flags required to enable jit fusion kernels
29
+
30
+ if sys.platform != 'darwin':
31
+ torch._C._jit_set_profiling_mode(False)
32
+ torch._C._jit_set_profiling_executor(False)
33
+ torch._C._jit_override_can_fuse_on_cpu(True)
34
+ torch._C._jit_override_can_fuse_on_gpu(True)
35
+
36
+ logger = logging.get_logger(__name__)
37
+
38
+ _CHECKPOINT_FOR_DOC = "THUDM/ChatGLM2-6B"
39
+ _CONFIG_FOR_DOC = "ChatGLM6BConfig"
40
+
41
+ CHATGLM_6B_PRETRAINED_MODEL_ARCHIVE_LIST = [
42
+ "THUDM/chatglm2-6b",
43
+ # See all ChatGLM models at https://huggingface.co/models?filter=chatglm
44
+ ]
45
+
46
+
47
+ def default_init(cls, *args, **kwargs):
48
+ return cls(*args, **kwargs)
49
+
50
+
51
+ class InvalidScoreLogitsProcessor(LogitsProcessor):
52
+ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
53
+ if torch.isnan(scores).any() or torch.isinf(scores).any():
54
+ scores.zero_()
55
+ scores[..., 5] = 5e4
56
+ return scores
57
+
58
+
59
+ class PrefixEncoder(torch.nn.Module):
60
+ """
61
+ The torch.nn model to encode the prefix
62
+ Input shape: (batch-size, prefix-length)
63
+ Output shape: (batch-size, prefix-length, 2*layers*hidden)
64
+ """
65
+
66
+ def __init__(self, config: ChatGLMConfig):
67
+ super().__init__()
68
+ self.prefix_projection = config.prefix_projection
69
+ if self.prefix_projection:
70
+ # Use a two-layer MLP to encode the prefix
71
+ kv_size = config.num_layers * config.kv_channels * config.multi_query_group_num * 2
72
+ self.embedding = torch.nn.Embedding(config.pre_seq_len, kv_size)
73
+ self.trans = torch.nn.Sequential(
74
+ torch.nn.Linear(kv_size, config.hidden_size),
75
+ torch.nn.Tanh(),
76
+ torch.nn.Linear(config.hidden_size, kv_size)
77
+ )
78
+ else:
79
+ self.embedding = torch.nn.Embedding(config.pre_seq_len,
80
+ config.num_layers * config.kv_channels * config.multi_query_group_num * 2)
81
+
82
+ def forward(self, prefix: torch.Tensor):
83
+ if self.prefix_projection:
84
+ prefix_tokens = self.embedding(prefix)
85
+ past_key_values = self.trans(prefix_tokens)
86
+ else:
87
+ past_key_values = self.embedding(prefix)
88
+ return past_key_values
89
+
90
+
91
+ def split_tensor_along_last_dim(
92
+ tensor: torch.Tensor,
93
+ num_partitions: int,
94
+ contiguous_split_chunks: bool = False,
95
+ ) -> List[torch.Tensor]:
96
+ """Split a tensor along its last dimension.
97
+
98
+ Arguments:
99
+ tensor: input tensor.
100
+ num_partitions: number of partitions to split the tensor
101
+ contiguous_split_chunks: If True, make each chunk contiguous
102
+ in memory.
103
+
104
+ Returns:
105
+ A list of Tensors
106
+ """
107
+ # Get the size and dimension.
108
+ last_dim = tensor.dim() - 1
109
+ last_dim_size = tensor.size()[last_dim] // num_partitions
110
+ # Split.
111
+ tensor_list = torch.split(tensor, last_dim_size, dim=last_dim)
112
+ # Note: torch.split does not create contiguous tensors by default.
113
+ if contiguous_split_chunks:
114
+ return tuple(chunk.contiguous() for chunk in tensor_list)
115
+
116
+ return tensor_list
117
+
118
+
119
+ class RotaryEmbedding(nn.Module):
120
+ def __init__(self, dim, original_impl=False, device=None, dtype=None):
121
+ super().__init__()
122
+ inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, device=device).to(dtype=dtype) / dim))
123
+ self.register_buffer("inv_freq", inv_freq)
124
+ self.dim = dim
125
+ self.original_impl = original_impl
126
+
127
+ def forward_impl(
128
+ self, seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000
129
+ ):
130
+ """Enhanced Transformer with Rotary Position Embedding.
131
+
132
+ Derived from: https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/
133
+ transformers/rope/__init__.py. MIT License:
134
+ https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license.
135
+ """
136
+ # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
137
+ theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, dtype=dtype, device=device) / n_elem))
138
+
139
+ # Create position indexes `[0, 1, ..., seq_len - 1]`
140
+ seq_idx = torch.arange(seq_len, dtype=dtype, device=device)
141
+
142
+ # Calculate the product of position index and $\theta_i$
143
+ idx_theta = torch.outer(seq_idx, theta).float()
144
+
145
+ cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1)
146
+
147
+ # this is to mimic the behaviour of complex32, else we will get different results
148
+ if dtype in (torch.float16, torch.bfloat16, torch.int8):
149
+ cache = cache.bfloat16() if dtype == torch.bfloat16 else cache.half()
150
+ return cache
151
+
152
+ def forward(self, max_seq_len, offset=0):
153
+ return self.forward_impl(
154
+ max_seq_len, self.dim, dtype=self.inv_freq.dtype, device=self.inv_freq.device
155
+ )
156
+
157
+
158
+ @torch.jit.script
159
+ def apply_rotary_pos_emb(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor:
160
+ # x: [sq, b, np, hn]
161
+ sq, b, np, hn = x.size(0), x.size(1), x.size(2), x.size(3)
162
+ rot_dim = rope_cache.shape[-2] * 2
163
+ x, x_pass = x[..., :rot_dim], x[..., rot_dim:]
164
+ # truncate to support variable sizes
165
+ rope_cache = rope_cache[:sq]
166
+ xshaped = x.reshape(sq, -1, np, rot_dim // 2, 2)
167
+ rope_cache = rope_cache.view(sq, -1, 1, xshaped.size(3), 2)
168
+ x_out2 = torch.stack(
169
+ [
170
+ xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1],
171
+ xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1],
172
+ ],
173
+ -1,
174
+ )
175
+ x_out2 = x_out2.flatten(3)
176
+ return torch.cat((x_out2, x_pass), dim=-1)
177
+
178
+
179
+ class RMSNorm(torch.nn.Module):
180
+ def __init__(self, normalized_shape, eps=1e-5, device=None, dtype=None, **kwargs):
181
+ super().__init__()
182
+ self.weight = torch.nn.Parameter(torch.empty(normalized_shape, device=device, dtype=dtype))
183
+ self.eps = eps
184
+
185
+ def forward(self, hidden_states: torch.Tensor):
186
+ input_dtype = hidden_states.dtype
187
+ variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
188
+ hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
189
+
190
+ return (self.weight * hidden_states).to(input_dtype)
191
+
192
+
193
+ class CoreAttention(torch.nn.Module):
194
+ def __init__(self, config: ChatGLMConfig, layer_number):
195
+ super(CoreAttention, self).__init__()
196
+
197
+ self.apply_query_key_layer_scaling = config.apply_query_key_layer_scaling
198
+ self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32
199
+ if self.apply_query_key_layer_scaling:
200
+ self.attention_softmax_in_fp32 = True
201
+ self.layer_number = max(1, layer_number)
202
+
203
+ projection_size = config.kv_channels * config.num_attention_heads
204
+
205
+ # Per attention head and per partition values.
206
+ self.hidden_size_per_partition = projection_size
207
+ self.hidden_size_per_attention_head = projection_size // config.num_attention_heads
208
+ self.num_attention_heads_per_partition = config.num_attention_heads
209
+
210
+ coeff = None
211
+ self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)
212
+ if self.apply_query_key_layer_scaling:
213
+ coeff = self.layer_number
214
+ self.norm_factor *= coeff
215
+ self.coeff = coeff
216
+
217
+ self.attention_dropout = torch.nn.Dropout(config.attention_dropout)
218
+
219
+ def forward(self, query_layer, key_layer, value_layer, attention_mask):
220
+ pytorch_major_version = int(torch.__version__.split('.')[0])
221
+ if pytorch_major_version >= 2:
222
+ query_layer, key_layer, value_layer = [k.permute(1, 2, 0, 3) for k in [query_layer, key_layer, value_layer]]
223
+ if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]:
224
+ context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
225
+ is_causal=True)
226
+ else:
227
+ if attention_mask is not None:
228
+ attention_mask = ~attention_mask
229
+ context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
230
+ attention_mask)
231
+ context_layer = context_layer.permute(2, 0, 1, 3)
232
+ new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
233
+ context_layer = context_layer.reshape(*new_context_layer_shape)
234
+ else:
235
+ # Raw attention scores
236
+
237
+ # [b, np, sq, sk]
238
+ output_size = (query_layer.size(1), query_layer.size(2), query_layer.size(0), key_layer.size(0))
239
+
240
+ # [sq, b, np, hn] -> [sq, b * np, hn]
241
+ query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1)
242
+ # [sk, b, np, hn] -> [sk, b * np, hn]
243
+ key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1)
244
+
245
+ # preallocting input tensor: [b * np, sq, sk]
246
+ matmul_input_buffer = torch.empty(
247
+ output_size[0] * output_size[1], output_size[2], output_size[3], dtype=query_layer.dtype,
248
+ device=query_layer.device
249
+ )
250
+
251
+ # Raw attention scores. [b * np, sq, sk]
252
+ matmul_result = torch.baddbmm(
253
+ matmul_input_buffer,
254
+ query_layer.transpose(0, 1), # [b * np, sq, hn]
255
+ key_layer.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk]
256
+ beta=0.0,
257
+ alpha=(1.0 / self.norm_factor),
258
+ )
259
+
260
+ # change view to [b, np, sq, sk]
261
+ attention_scores = matmul_result.view(*output_size)
262
+
263
+ # ===========================
264
+ # Attention probs and dropout
265
+ # ===========================
266
+
267
+ # attention scores and attention mask [b, np, sq, sk]
268
+ if self.attention_softmax_in_fp32:
269
+ attention_scores = attention_scores.float()
270
+ if self.coeff is not None:
271
+ attention_scores = attention_scores * self.coeff
272
+ if attention_mask is None and attention_scores.shape[2] == attention_scores.shape[3]:
273
+ attention_mask = torch.ones(output_size[0], 1, output_size[2], output_size[3],
274
+ device=attention_scores.device, dtype=torch.bool)
275
+ attention_mask.tril_()
276
+ attention_mask = ~attention_mask
277
+ if attention_mask is not None:
278
+ attention_scores = attention_scores.masked_fill(attention_mask, float("-inf"))
279
+ attention_probs = F.softmax(attention_scores, dim=-1)
280
+ attention_probs = attention_probs.type_as(value_layer)
281
+
282
+ # This is actually dropping out entire tokens to attend to, which might
283
+ # seem a bit unusual, but is taken from the original Transformer paper.
284
+ attention_probs = self.attention_dropout(attention_probs)
285
+ # =========================
286
+ # Context layer. [sq, b, hp]
287
+ # =========================
288
+
289
+ # value_layer -> context layer.
290
+ # [sk, b, np, hn] --> [b, np, sq, hn]
291
+
292
+ # context layer shape: [b, np, sq, hn]
293
+ output_size = (value_layer.size(1), value_layer.size(2), query_layer.size(0), value_layer.size(3))
294
+ # change view [sk, b * np, hn]
295
+ value_layer = value_layer.view(value_layer.size(0), output_size[0] * output_size[1], -1)
296
+ # change view [b * np, sq, sk]
297
+ attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
298
+ # matmul: [b * np, sq, hn]
299
+ context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1))
300
+ # change view [b, np, sq, hn]
301
+ context_layer = context_layer.view(*output_size)
302
+ # [b, np, sq, hn] --> [sq, b, np, hn]
303
+ context_layer = context_layer.permute(2, 0, 1, 3).contiguous()
304
+ # [sq, b, np, hn] --> [sq, b, hp]
305
+ new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
306
+ context_layer = context_layer.view(*new_context_layer_shape)
307
+
308
+ return context_layer
309
+
310
+
311
+ class SelfAttention(torch.nn.Module):
312
+ """Parallel self-attention layer abstract class.
313
+
314
+ Self-attention layer takes input with size [s, b, h]
315
+ and returns output of the same size.
316
+ """
317
+
318
+ def __init__(self, config: ChatGLMConfig, layer_number, device=None):
319
+ super(SelfAttention, self).__init__()
320
+ self.layer_number = max(1, layer_number)
321
+
322
+ self.projection_size = config.kv_channels * config.num_attention_heads
323
+
324
+ # Per attention head and per partition values.
325
+ self.hidden_size_per_attention_head = self.projection_size // config.num_attention_heads
326
+ self.num_attention_heads_per_partition = config.num_attention_heads
327
+
328
+ self.multi_query_attention = config.multi_query_attention
329
+ self.qkv_hidden_size = 3 * self.projection_size
330
+ if self.multi_query_attention:
331
+ self.num_multi_query_groups_per_partition = config.multi_query_group_num
332
+ self.qkv_hidden_size = (
333
+ self.projection_size + 2 * self.hidden_size_per_attention_head * config.multi_query_group_num
334
+ )
335
+ self.query_key_value = nn.Linear(config.hidden_size, self.qkv_hidden_size,
336
+ bias=config.add_bias_linear or config.add_qkv_bias,
337
+ device=device, **_config_to_kwargs(config)
338
+ )
339
+
340
+ self.core_attention = CoreAttention(config, self.layer_number)
341
+
342
+ # Output.
343
+ self.dense = nn.Linear(self.projection_size, config.hidden_size, bias=config.add_bias_linear,
344
+ device=device, **_config_to_kwargs(config)
345
+ )
346
+
347
+ def _allocate_memory(self, inference_max_sequence_len, batch_size, device=None, dtype=None):
348
+ if self.multi_query_attention:
349
+ num_attention_heads = self.num_multi_query_groups_per_partition
350
+ else:
351
+ num_attention_heads = self.num_attention_heads_per_partition
352
+ return torch.empty(
353
+ inference_max_sequence_len,
354
+ batch_size,
355
+ num_attention_heads,
356
+ self.hidden_size_per_attention_head,
357
+ dtype=dtype,
358
+ device=device,
359
+ )
360
+
361
+ def forward(
362
+ self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True
363
+ ):
364
+ # hidden_states: [sq, b, h]
365
+
366
+ # =================================================
367
+ # Pre-allocate memory for key-values for inference.
368
+ # =================================================
369
+ # =====================
370
+ # Query, Key, and Value
371
+ # =====================
372
+
373
+ # Attention heads [sq, b, h] --> [sq, b, (np * 3 * hn)]
374
+ mixed_x_layer = self.query_key_value(hidden_states)
375
+
376
+ if self.multi_query_attention:
377
+ (query_layer, key_layer, value_layer) = mixed_x_layer.split(
378
+ [
379
+ self.num_attention_heads_per_partition * self.hidden_size_per_attention_head,
380
+ self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
381
+ self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
382
+ ],
383
+ dim=-1,
384
+ )
385
+ query_layer = query_layer.view(
386
+ query_layer.size()[:-1] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)
387
+ )
388
+ key_layer = key_layer.view(
389
+ key_layer.size()[:-1] + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)
390
+ )
391
+ value_layer = value_layer.view(
392
+ value_layer.size()[:-1]
393
+ + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)
394
+ )
395
+ else:
396
+ new_tensor_shape = mixed_x_layer.size()[:-1] + \
397
+ (self.num_attention_heads_per_partition,
398
+ 3 * self.hidden_size_per_attention_head)
399
+ mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)
400
+
401
+ # [sq, b, np, 3 * hn] --> 3 [sq, b, np, hn]
402
+ (query_layer, key_layer, value_layer) = split_tensor_along_last_dim(mixed_x_layer, 3)
403
+
404
+ # apply relative positional encoding (rotary embedding)
405
+ if rotary_pos_emb is not None:
406
+ query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb)
407
+ key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb)
408
+
409
+ # adjust key and value for inference
410
+ if kv_cache is not None:
411
+ cache_k, cache_v = kv_cache
412
+ key_layer = torch.cat((cache_k, key_layer), dim=0)
413
+ value_layer = torch.cat((cache_v, value_layer), dim=0)
414
+ if use_cache:
415
+ kv_cache = (key_layer, value_layer)
416
+ else:
417
+ kv_cache = None
418
+
419
+ if self.multi_query_attention:
420
+ key_layer = key_layer.unsqueeze(-2)
421
+ key_layer = key_layer.expand(
422
+ -1, -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1
423
+ )
424
+ key_layer = key_layer.contiguous().view(
425
+ key_layer.size()[:2] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)
426
+ )
427
+ value_layer = value_layer.unsqueeze(-2)
428
+ value_layer = value_layer.expand(
429
+ -1, -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1
430
+ )
431
+ value_layer = value_layer.contiguous().view(
432
+ value_layer.size()[:2] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)
433
+ )
434
+
435
+ # ==================================
436
+ # core attention computation
437
+ # ==================================
438
+
439
+ context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask)
440
+
441
+ # =================
442
+ # Output. [sq, b, h]
443
+ # =================
444
+
445
+ output = self.dense(context_layer)
446
+
447
+ return output, kv_cache
448
+
449
+
450
+ def _config_to_kwargs(args):
451
+ common_kwargs = {
452
+ "dtype": args.torch_dtype,
453
+ }
454
+ return common_kwargs
455
+
456
+
457
+ class MLP(torch.nn.Module):
458
+ """MLP.
459
+
460
+ MLP will take the input with h hidden state, project it to 4*h
461
+ hidden dimension, perform nonlinear transformation, and project the
462
+ state back into h hidden dimension.
463
+ """
464
+
465
+ def __init__(self, config: ChatGLMConfig, device=None):
466
+ super(MLP, self).__init__()
467
+
468
+ self.add_bias = config.add_bias_linear
469
+
470
+ # Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf
471
+ self.dense_h_to_4h = nn.Linear(
472
+ config.hidden_size,
473
+ config.ffn_hidden_size * 2,
474
+ bias=self.add_bias,
475
+ device=device,
476
+ **_config_to_kwargs(config)
477
+ )
478
+
479
+ def swiglu(x):
480
+ x = torch.chunk(x, 2, dim=-1)
481
+ return F.silu(x[0]) * x[1]
482
+
483
+ self.activation_func = swiglu
484
+
485
+ # Project back to h.
486
+ self.dense_4h_to_h = nn.Linear(
487
+ config.ffn_hidden_size,
488
+ config.hidden_size,
489
+ bias=self.add_bias,
490
+ device=device,
491
+ **_config_to_kwargs(config)
492
+ )
493
+
494
+ def forward(self, hidden_states):
495
+ # [s, b, 4hp]
496
+ intermediate_parallel = self.dense_h_to_4h(hidden_states)
497
+ intermediate_parallel = self.activation_func(intermediate_parallel)
498
+ # [s, b, h]
499
+ output = self.dense_4h_to_h(intermediate_parallel)
500
+ return output
501
+
502
+
503
+ class GLMBlock(torch.nn.Module):
504
+ """A single transformer layer.
505
+
506
+ Transformer layer takes input with size [s, b, h] and returns an
507
+ output of the same size.
508
+ """
509
+
510
+ def __init__(self, config: ChatGLMConfig, layer_number, device=None):
511
+ super(GLMBlock, self).__init__()
512
+ self.layer_number = layer_number
513
+
514
+ self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm
515
+
516
+ self.fp32_residual_connection = config.fp32_residual_connection
517
+
518
+ LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm
519
+ # Layernorm on the input data.
520
+ self.input_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
521
+ dtype=config.torch_dtype)
522
+
523
+ # Self attention.
524
+ self.self_attention = SelfAttention(config, layer_number, device=device)
525
+ self.hidden_dropout = config.hidden_dropout
526
+
527
+ # Layernorm on the attention output
528
+ self.post_attention_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
529
+ dtype=config.torch_dtype)
530
+
531
+ # MLP
532
+ self.mlp = MLP(config, device=device)
533
+
534
+ def forward(
535
+ self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True,
536
+ ):
537
+ # hidden_states: [s, b, h]
538
+
539
+ # Layer norm at the beginning of the transformer layer.
540
+ layernorm_output = self.input_layernorm(hidden_states)
541
+ # Self attention.
542
+ attention_output, kv_cache = self.self_attention(
543
+ layernorm_output,
544
+ attention_mask,
545
+ rotary_pos_emb,
546
+ kv_cache=kv_cache,
547
+ use_cache=use_cache
548
+ )
549
+
550
+ # Residual connection.
551
+ if self.apply_residual_connection_post_layernorm:
552
+ residual = layernorm_output
553
+ else:
554
+ residual = hidden_states
555
+
556
+ layernorm_input = torch.nn.functional.dropout(attention_output, p=self.hidden_dropout, training=self.training)
557
+ layernorm_input = residual + layernorm_input
558
+
559
+ # Layer norm post the self attention.
560
+ layernorm_output = self.post_attention_layernorm(layernorm_input)
561
+
562
+ # MLP.
563
+ mlp_output = self.mlp(layernorm_output)
564
+
565
+ # Second residual connection.
566
+ if self.apply_residual_connection_post_layernorm:
567
+ residual = layernorm_output
568
+ else:
569
+ residual = layernorm_input
570
+
571
+ output = torch.nn.functional.dropout(mlp_output, p=self.hidden_dropout, training=self.training)
572
+ output = residual + output
573
+
574
+ return output, kv_cache
575
+
576
+
577
+ class GLMTransformer(torch.nn.Module):
578
+ """Transformer class."""
579
+
580
+ def __init__(self, config: ChatGLMConfig, device=None):
581
+ super(GLMTransformer, self).__init__()
582
+
583
+ self.fp32_residual_connection = config.fp32_residual_connection
584
+ self.post_layer_norm = config.post_layer_norm
585
+
586
+ # Number of layers.
587
+ self.num_layers = config.num_layers
588
+
589
+ # Transformer layers.
590
+ def build_layer(layer_number):
591
+ return GLMBlock(config, layer_number, device=device)
592
+
593
+ self.layers = torch.nn.ModuleList([build_layer(i + 1) for i in range(self.num_layers)])
594
+
595
+ if self.post_layer_norm:
596
+ LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm
597
+ # Final layer norm before output.
598
+ self.final_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
599
+ dtype=config.torch_dtype)
600
+
601
+ self.gradient_checkpointing = False
602
+
603
+ def _get_layer(self, layer_number):
604
+ return self.layers[layer_number]
605
+
606
+ def forward(
607
+ self, hidden_states, attention_mask, rotary_pos_emb, kv_caches=None,
608
+ use_cache: Optional[bool] = True,
609
+ output_hidden_states: Optional[bool] = False,
610
+ ):
611
+ if not kv_caches:
612
+ kv_caches = [None for _ in range(self.num_layers)]
613
+ presents = () if use_cache else None
614
+ if self.gradient_checkpointing and self.training:
615
+ if use_cache:
616
+ logger.warning_once(
617
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
618
+ )
619
+ use_cache = False
620
+
621
+ all_self_attentions = None
622
+ all_hidden_states = () if output_hidden_states else None
623
+ for index in range(self.num_layers):
624
+ if output_hidden_states:
625
+ all_hidden_states = all_hidden_states + (hidden_states,)
626
+
627
+ layer = self._get_layer(index)
628
+ if self.gradient_checkpointing and self.training:
629
+ layer_ret = torch.utils.checkpoint.checkpoint(
630
+ layer,
631
+ hidden_states,
632
+ attention_mask,
633
+ rotary_pos_emb,
634
+ kv_caches[index],
635
+ use_cache
636
+ )
637
+ else:
638
+ layer_ret = layer(
639
+ hidden_states,
640
+ attention_mask,
641
+ rotary_pos_emb,
642
+ kv_cache=kv_caches[index],
643
+ use_cache=use_cache
644
+ )
645
+ hidden_states, kv_cache = layer_ret
646
+ if use_cache:
647
+ presents = presents + (kv_cache,)
648
+
649
+ if output_hidden_states:
650
+ all_hidden_states = all_hidden_states + (hidden_states,)
651
+
652
+ # Final layer norm.
653
+ if self.post_layer_norm:
654
+ hidden_states = self.final_layernorm(hidden_states)
655
+
656
+ return hidden_states, presents, all_hidden_states, all_self_attentions
657
+
658
+
659
+ class ChatGLMPreTrainedModel(PreTrainedModel):
660
+ """
661
+ An abstract class to handle weights initialization and
662
+ a simple interface for downloading and loading pretrained models.
663
+ """
664
+
665
+ is_parallelizable = False
666
+ supports_gradient_checkpointing = True
667
+ config_class = ChatGLMConfig
668
+ base_model_prefix = "transformer"
669
+ _no_split_modules = ["GLMBlock"]
670
+
671
+ def _init_weights(self, module: nn.Module):
672
+ """Initialize the weights."""
673
+ return
674
+
675
+ def get_masks(self, input_ids, past_key_values, padding_mask=None):
676
+ batch_size, seq_length = input_ids.shape
677
+ full_attention_mask = torch.ones(batch_size, seq_length, seq_length, device=input_ids.device)
678
+ full_attention_mask.tril_()
679
+ past_length = 0
680
+ if past_key_values:
681
+ past_length = past_key_values[0][0].shape[0]
682
+ if past_length:
683
+ full_attention_mask = torch.cat((torch.ones(batch_size, seq_length, past_length,
684
+ device=input_ids.device), full_attention_mask), dim=-1)
685
+ if padding_mask is not None:
686
+ full_attention_mask = full_attention_mask * padding_mask.unsqueeze(1)
687
+ if not past_length and padding_mask is not None:
688
+ full_attention_mask -= padding_mask.unsqueeze(-1) - 1
689
+ full_attention_mask = (full_attention_mask < 0.5).bool()
690
+ full_attention_mask.unsqueeze_(1)
691
+ return full_attention_mask
692
+
693
+ def get_position_ids(self, input_ids, device):
694
+ batch_size, seq_length = input_ids.shape
695
+ position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1)
696
+ return position_ids
697
+
698
+ def _set_gradient_checkpointing(self, module, value=False):
699
+ if isinstance(module, GLMTransformer):
700
+ module.gradient_checkpointing = value
701
+
702
+
703
+ class Embedding(torch.nn.Module):
704
+ """Language model embeddings."""
705
+
706
+ def __init__(self, config: ChatGLMConfig, device=None):
707
+ super(Embedding, self).__init__()
708
+
709
+ self.hidden_size = config.hidden_size
710
+ # Word embeddings (parallel).
711
+ self.word_embeddings = nn.Embedding(
712
+ config.padded_vocab_size,
713
+ self.hidden_size,
714
+ dtype=config.torch_dtype,
715
+ device=device
716
+ )
717
+ self.fp32_residual_connection = config.fp32_residual_connection
718
+
719
+ def forward(self, input_ids):
720
+ # Embeddings.
721
+ words_embeddings = self.word_embeddings(input_ids)
722
+ embeddings = words_embeddings
723
+ # Data format change to avoid explicit tranposes : [b s h] --> [s b h].
724
+ embeddings = embeddings.transpose(0, 1).contiguous()
725
+ # If the input flag for fp32 residual connection is set, convert for float.
726
+ if self.fp32_residual_connection:
727
+ embeddings = embeddings.float()
728
+ return embeddings
729
+
730
+
731
+ class ChatGLMModel(ChatGLMPreTrainedModel):
732
+ def __init__(self, config: ChatGLMConfig, device=None, empty_init=True):
733
+ super().__init__(config)
734
+ if empty_init:
735
+ init_method = skip_init
736
+ else:
737
+ init_method = default_init
738
+ init_kwargs = {}
739
+ if device is not None:
740
+ init_kwargs["device"] = device
741
+ self.embedding = init_method(Embedding, config, **init_kwargs)
742
+ self.num_layers = config.num_layers
743
+ self.multi_query_group_num = config.multi_query_group_num
744
+ self.kv_channels = config.kv_channels
745
+
746
+ # Rotary positional embeddings
747
+ self.seq_length = config.seq_length
748
+ rotary_dim = (
749
+ config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels
750
+ )
751
+
752
+ self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, original_impl=config.original_rope, device=device,
753
+ dtype=config.torch_dtype)
754
+ self.encoder = init_method(GLMTransformer, config, **init_kwargs)
755
+ self.output_layer = init_method(nn.Linear, config.hidden_size, config.padded_vocab_size, bias=False,
756
+ dtype=config.torch_dtype, **init_kwargs)
757
+ self.pre_seq_len = config.pre_seq_len
758
+ self.prefix_projection = config.prefix_projection
759
+ if self.pre_seq_len is not None:
760
+ for param in self.parameters():
761
+ param.requires_grad = False
762
+ self.prefix_tokens = torch.arange(self.pre_seq_len).long()
763
+ self.prefix_encoder = PrefixEncoder(config)
764
+ self.dropout = torch.nn.Dropout(0.1)
765
+
766
+ def get_input_embeddings(self):
767
+ return self.embedding.word_embeddings
768
+
769
+ def get_prompt(self, batch_size, device, dtype=torch.half):
770
+ prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(device)
771
+ past_key_values = self.prefix_encoder(prefix_tokens).type(dtype)
772
+ past_key_values = past_key_values.view(
773
+ batch_size,
774
+ self.pre_seq_len,
775
+ self.num_layers * 2,
776
+ self.multi_query_group_num,
777
+ self.kv_channels
778
+ )
779
+ # seq_len, b, nh, hidden_size
780
+ past_key_values = self.dropout(past_key_values)
781
+ past_key_values = past_key_values.permute([2, 1, 0, 3, 4]).split(2)
782
+ return past_key_values
783
+
784
+ def forward(
785
+ self,
786
+ input_ids,
787
+ position_ids: Optional[torch.Tensor] = None,
788
+ attention_mask: Optional[torch.BoolTensor] = None,
789
+ full_attention_mask: Optional[torch.BoolTensor] = None,
790
+ past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
791
+ inputs_embeds: Optional[torch.Tensor] = None,
792
+ use_cache: Optional[bool] = None,
793
+ output_hidden_states: Optional[bool] = None,
794
+ return_dict: Optional[bool] = None,
795
+ ):
796
+ output_hidden_states = (
797
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
798
+ )
799
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
800
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
801
+
802
+ batch_size, seq_length = input_ids.shape
803
+
804
+ if inputs_embeds is None:
805
+ inputs_embeds = self.embedding(input_ids)
806
+
807
+ if self.pre_seq_len is not None:
808
+ if past_key_values is None:
809
+ past_key_values = self.get_prompt(batch_size=batch_size, device=input_ids.device,
810
+ dtype=inputs_embeds.dtype)
811
+ if attention_mask is not None:
812
+ attention_mask = torch.cat([attention_mask.new_ones((batch_size, self.pre_seq_len)),
813
+ attention_mask], dim=-1)
814
+
815
+ if full_attention_mask is None:
816
+ if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1):
817
+ full_attention_mask = self.get_masks(input_ids, past_key_values, padding_mask=attention_mask)
818
+
819
+ # Rotary positional embeddings
820
+ rotary_pos_emb = self.rotary_pos_emb(self.seq_length)
821
+ if position_ids is not None:
822
+ rotary_pos_emb = rotary_pos_emb[position_ids]
823
+ else:
824
+ rotary_pos_emb = rotary_pos_emb[None, :seq_length]
825
+ rotary_pos_emb = rotary_pos_emb.transpose(0, 1).contiguous()
826
+
827
+ # Run encoder.
828
+ hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder(
829
+ inputs_embeds, full_attention_mask, rotary_pos_emb=rotary_pos_emb,
830
+ kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states
831
+ )
832
+
833
+ if not return_dict:
834
+ return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)
835
+
836
+ return BaseModelOutputWithPast(
837
+ last_hidden_state=hidden_states,
838
+ past_key_values=presents,
839
+ hidden_states=all_hidden_states,
840
+ attentions=all_self_attentions,
841
+ )
842
+
843
+ def quantize(self, weight_bit_width: int):
844
+ from .quantization import quantize
845
+ quantize(self.encoder, weight_bit_width)
846
+ return self
847
+
848
+
849
+ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
850
+ def __init__(self, config: ChatGLMConfig, empty_init=True, device=None):
851
+ super().__init__(config)
852
+
853
+ self.max_sequence_length = config.max_length
854
+ self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device)
855
+ self.config = config
856
+ self.quantized = False
857
+
858
+ if self.config.quantization_bit:
859
+ self.quantize(self.config.quantization_bit, empty_init=True)
860
+
861
+ def _update_model_kwargs_for_generation(
862
+ self,
863
+ outputs: ModelOutput,
864
+ model_kwargs: Dict[str, Any],
865
+ is_encoder_decoder: bool = False,
866
+ standardize_cache_format: bool = False,
867
+ ) -> Dict[str, Any]:
868
+ # update past_key_values
869
+ model_kwargs["past_key_values"] = self._extract_past_from_model_output(
870
+ outputs, standardize_cache_format=standardize_cache_format
871
+ )
872
+
873
+ # update attention mask
874
+ if "attention_mask" in model_kwargs:
875
+ attention_mask = model_kwargs["attention_mask"]
876
+ model_kwargs["attention_mask"] = torch.cat(
877
+ [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
878
+ )
879
+
880
+ # update position ids
881
+ if "position_ids" in model_kwargs:
882
+ position_ids = model_kwargs["position_ids"]
883
+ new_position_id = position_ids[..., -1:].clone()
884
+ new_position_id += 1
885
+ model_kwargs["position_ids"] = torch.cat(
886
+ [position_ids, new_position_id], dim=-1
887
+ )
888
+
889
+ model_kwargs["is_first_forward"] = False
890
+ return model_kwargs
891
+
892
+ def prepare_inputs_for_generation(
893
+ self,
894
+ input_ids: torch.LongTensor,
895
+ past_key_values: Optional[torch.Tensor] = None,
896
+ attention_mask: Optional[torch.Tensor] = None,
897
+ position_ids: Optional[torch.Tensor] = None,
898
+ is_first_forward: bool = True,
899
+ **kwargs
900
+ ) -> dict:
901
+ # only last token for input_ids if past is not None
902
+ if position_ids is None:
903
+ position_ids = self.get_position_ids(input_ids, device=input_ids.device)
904
+ if not is_first_forward:
905
+ position_ids = position_ids[..., -1:]
906
+ input_ids = input_ids[:, -1:]
907
+ return {
908
+ "input_ids": input_ids,
909
+ "past_key_values": past_key_values,
910
+ "position_ids": position_ids,
911
+ "attention_mask": attention_mask,
912
+ "return_last_logit": True
913
+ }
914
+
915
+ def forward(
916
+ self,
917
+ input_ids: Optional[torch.Tensor] = None,
918
+ position_ids: Optional[torch.Tensor] = None,
919
+ attention_mask: Optional[torch.Tensor] = None,
920
+ past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
921
+ inputs_embeds: Optional[torch.Tensor] = None,
922
+ labels: Optional[torch.Tensor] = None,
923
+ use_cache: Optional[bool] = None,
924
+ output_attentions: Optional[bool] = None,
925
+ output_hidden_states: Optional[bool] = None,
926
+ return_dict: Optional[bool] = None,
927
+ return_last_logit: Optional[bool] = False,
928
+ ):
929
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
930
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
931
+
932
+ transformer_outputs = self.transformer(
933
+ input_ids=input_ids,
934
+ position_ids=position_ids,
935
+ attention_mask=attention_mask,
936
+ past_key_values=past_key_values,
937
+ inputs_embeds=inputs_embeds,
938
+ use_cache=use_cache,
939
+ output_hidden_states=output_hidden_states,
940
+ return_dict=return_dict,
941
+ )
942
+
943
+ hidden_states = transformer_outputs[0]
944
+ if return_last_logit:
945
+ hidden_states = hidden_states[-1:]
946
+ lm_logits = self.transformer.output_layer(hidden_states)
947
+ lm_logits = lm_logits.transpose(0, 1).contiguous()
948
+
949
+ loss = None
950
+ if labels is not None:
951
+ lm_logits = lm_logits.to(torch.float32)
952
+
953
+ # Shift so that tokens < n predict n
954
+ shift_logits = lm_logits[..., :-1, :].contiguous()
955
+ shift_labels = labels[..., 1:].contiguous()
956
+ # Flatten the tokens
957
+ loss_fct = CrossEntropyLoss(ignore_index=-100)
958
+ loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
959
+
960
+ lm_logits = lm_logits.to(hidden_states.dtype)
961
+ loss = loss.to(hidden_states.dtype)
962
+
963
+ if not return_dict:
964
+ output = (lm_logits,) + transformer_outputs[1:]
965
+ return ((loss,) + output) if loss is not None else output
966
+
967
+ return CausalLMOutputWithPast(
968
+ loss=loss,
969
+ logits=lm_logits,
970
+ past_key_values=transformer_outputs.past_key_values,
971
+ hidden_states=transformer_outputs.hidden_states,
972
+ attentions=transformer_outputs.attentions,
973
+ )
974
+
975
+ @staticmethod
976
+ def _reorder_cache(
977
+ past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor
978
+ ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:
979
+ """
980
+ This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
981
+ [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
982
+ beam_idx at every generation step.
983
+
984
+ Output shares the same memory storage as `past`.
985
+ """
986
+ return tuple(
987
+ (
988
+ layer_past[0].index_select(1, beam_idx.to(layer_past[0].device)),
989
+ layer_past[1].index_select(1, beam_idx.to(layer_past[1].device)),
990
+ )
991
+ for layer_past in past
992
+ )
993
+
994
+ def process_response(self, response):
995
+ response = response.strip()
996
+ response = response.replace("[[训练时间]]", "2023年")
997
+ return response
998
+
999
+ def build_inputs(self, tokenizer, query: str, history: List[Tuple[str, str]] = None):
1000
+ prompt = tokenizer.build_prompt(query, history=history)
1001
+ inputs = tokenizer([prompt], return_tensors="pt")
1002
+ inputs = inputs.to(self.device)
1003
+ return inputs
1004
+
1005
+ def build_stream_inputs(self, tokenizer, query: str, history: List[Tuple[str, str]] = None):
1006
+ if history:
1007
+ prompt = "\n\n[Round {}]\n\n问:{}\n\n答:".format(len(history) + 1, query)
1008
+ input_ids = tokenizer.encode(prompt, add_special_tokens=False)
1009
+ input_ids = input_ids[1:]
1010
+ inputs = tokenizer.batch_encode_plus([(input_ids, None)], return_tensors="pt", add_special_tokens=False)
1011
+ else:
1012
+ prompt = "[Round {}]\n\n问:{}\n\n答:".format(len(history) + 1, query)
1013
+ inputs = tokenizer([prompt], return_tensors="pt")
1014
+ inputs = inputs.to(self.device)
1015
+ return inputs
1016
+
1017
+ @torch.no_grad()
1018
+ def chat(self, tokenizer, query: str, history: List[Tuple[str, str]] = None, max_length: int = 8192, num_beams=1,
1019
+ do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None, **kwargs):
1020
+ if history is None:
1021
+ history = []
1022
+ if logits_processor is None:
1023
+ logits_processor = LogitsProcessorList()
1024
+ logits_processor.append(InvalidScoreLogitsProcessor())
1025
+ gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "do_sample": do_sample, "top_p": top_p,
1026
+ "temperature": temperature, "logits_processor": logits_processor, **kwargs}
1027
+ inputs = self.build_inputs(tokenizer, query, history=history)
1028
+ outputs = self.generate(**inputs, **gen_kwargs)
1029
+ outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):]
1030
+ response = tokenizer.decode(outputs)
1031
+ response = self.process_response(response)
1032
+ history = history + [(query, response)]
1033
+ return response, history
1034
+
1035
+ @torch.no_grad()
1036
+ def stream_chat(self, tokenizer, query: str, history: List[Tuple[str, str]] = None, past_key_values=None,
1037
+ max_length: int = 8192, do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None,
1038
+ return_past_key_values=False, **kwargs):
1039
+ if history is None:
1040
+ history = []
1041
+ if logits_processor is None:
1042
+ logits_processor = LogitsProcessorList()
1043
+ logits_processor.append(InvalidScoreLogitsProcessor())
1044
+ gen_kwargs = {"max_length": max_length, "do_sample": do_sample, "top_p": top_p,
1045
+ "temperature": temperature, "logits_processor": logits_processor, **kwargs}
1046
+ if past_key_values is None and not return_past_key_values:
1047
+ inputs = self.build_inputs(tokenizer, query, history=history)
1048
+ else:
1049
+ inputs = self.build_stream_inputs(tokenizer, query, history=history)
1050
+ if past_key_values is not None:
1051
+ past_length = past_key_values[0][0].shape[0]
1052
+ if self.transformer.pre_seq_len is not None:
1053
+ past_length -= self.transformer.pre_seq_len
1054
+ inputs.position_ids += past_length
1055
+ attention_mask = inputs.attention_mask
1056
+ attention_mask = torch.cat((attention_mask.new_ones(1, past_length), attention_mask), dim=1)
1057
+ inputs['attention_mask'] = attention_mask
1058
+ for outputs in self.stream_generate(**inputs, past_key_values=past_key_values,
1059
+ return_past_key_values=return_past_key_values, **gen_kwargs):
1060
+ if return_past_key_values:
1061
+ outputs, past_key_values = outputs
1062
+ outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):]
1063
+ response = tokenizer.decode(outputs)
1064
+ if response and response[-1] != "�":
1065
+ response = self.process_response(response)
1066
+ new_history = history + [(query, response)]
1067
+ if return_past_key_values:
1068
+ yield response, new_history, past_key_values
1069
+ else:
1070
+ yield response, new_history
1071
+
1072
+ @torch.no_grad()
1073
+ def stream_generate(
1074
+ self,
1075
+ input_ids,
1076
+ generation_config: Optional[GenerationConfig] = None,
1077
+ logits_processor: Optional[LogitsProcessorList] = None,
1078
+ stopping_criteria: Optional[StoppingCriteriaList] = None,
1079
+ prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
1080
+ return_past_key_values=False,
1081
+ **kwargs,
1082
+ ):
1083
+ batch_size, input_ids_seq_length = input_ids.shape[0], input_ids.shape[-1]
1084
+
1085
+ if generation_config is None:
1086
+ generation_config = self.generation_config
1087
+ generation_config = copy.deepcopy(generation_config)
1088
+ model_kwargs = generation_config.update(**kwargs)
1089
+ bos_token_id, eos_token_id = generation_config.bos_token_id, generation_config.eos_token_id
1090
+
1091
+ if isinstance(eos_token_id, int):
1092
+ eos_token_id = [eos_token_id]
1093
+
1094
+ has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
1095
+ if has_default_max_length and generation_config.max_new_tokens is None:
1096
+ warnings.warn(
1097
+ f"Using `max_length`'s default ({generation_config.max_length}) to control the generation length. "
1098
+ "This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we"
1099
+ " recommend using `max_new_tokens` to control the maximum length of the generation.",
1100
+ UserWarning,
1101
+ )
1102
+ elif generation_config.max_new_tokens is not None:
1103
+ generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length
1104
+ if not has_default_max_length:
1105
+ logger.warn(
1106
+ f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="
1107
+ f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "
1108
+ "Please refer to the documentation for more information. "
1109
+ "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)",
1110
+ UserWarning,
1111
+ )
1112
+
1113
+ if input_ids_seq_length >= generation_config.max_length:
1114
+ input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"
1115
+ logger.warning(
1116
+ f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to"
1117
+ f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
1118
+ " increasing `max_new_tokens`."
1119
+ )
1120
+
1121
+ # 2. Set generation parameters if not already defined
1122
+ logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
1123
+ stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
1124
+
1125
+ logits_processor = self._get_logits_processor(
1126
+ generation_config=generation_config,
1127
+ input_ids_seq_length=input_ids_seq_length,
1128
+ encoder_input_ids=input_ids,
1129
+ prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
1130
+ logits_processor=logits_processor,
1131
+ )
1132
+
1133
+ stopping_criteria = self._get_stopping_criteria(
1134
+ generation_config=generation_config, stopping_criteria=stopping_criteria
1135
+ )
1136
+ logits_warper = self._get_logits_warper(generation_config)
1137
+
1138
+ unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1)
1139
+ scores = None
1140
+ while True:
1141
+ model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
1142
+ # forward pass to get next token
1143
+ outputs = self(
1144
+ **model_inputs,
1145
+ return_dict=True,
1146
+ output_attentions=False,
1147
+ output_hidden_states=False,
1148
+ )
1149
+
1150
+ next_token_logits = outputs.logits[:, -1, :]
1151
+
1152
+ # pre-process distribution
1153
+ next_token_scores = logits_processor(input_ids, next_token_logits)
1154
+ next_token_scores = logits_warper(input_ids, next_token_scores)
1155
+
1156
+ # sample
1157
+ probs = nn.functional.softmax(next_token_scores, dim=-1)
1158
+ if generation_config.do_sample:
1159
+ next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
1160
+ else:
1161
+ next_tokens = torch.argmax(probs, dim=-1)
1162
+
1163
+ # update generated ids, model inputs, and length for next step
1164
+ input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
1165
+ model_kwargs = self._update_model_kwargs_for_generation(
1166
+ outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder
1167
+ )
1168
+ unfinished_sequences = unfinished_sequences.mul((sum(next_tokens != i for i in eos_token_id)).long())
1169
+ if return_past_key_values:
1170
+ yield input_ids, outputs.past_key_values
1171
+ else:
1172
+ yield input_ids
1173
+ # stop when each sentence is finished, or if we exceed the maximum length
1174
+ if unfinished_sequences.max() == 0 or stopping_criteria(input_ids, scores):
1175
+ break
1176
+
1177
+ def quantize(self, bits: int, empty_init=False, device=None, **kwargs):
1178
+ if bits == 0:
1179
+ return
1180
+
1181
+ from .quantization import quantize
1182
+
1183
+ if self.quantized:
1184
+ logger.info("Already quantized.")
1185
+ return self
1186
+
1187
+ self.quantized = True
1188
+
1189
+ self.config.quantization_bit = bits
1190
+
1191
+ self.transformer.encoder = quantize(self.transformer.encoder, bits, empty_init=empty_init, device=device,
1192
+ **kwargs)
1193
+ return self
checkpoint-3000/optimizer.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3ead273acff3493dd766f61114a57f2bf7f2e6094047cb9a67630d425d8eda3e
3
+ size 1346374055
checkpoint-3000/pytorch_model.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:276ad90420f3c49a67c1a26373b7feecb75e928b4b12e1da43a70298d8e45cb6
3
+ size 673186943
checkpoint-3000/quantization.py ADDED
@@ -0,0 +1,188 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from torch.nn import Linear
2
+ from torch.nn.parameter import Parameter
3
+
4
+ import bz2
5
+ import torch
6
+ import base64
7
+ import ctypes
8
+ from transformers.utils import logging
9
+
10
+ from typing import List
11
+ from functools import partial
12
+
13
+ logger = logging.get_logger(__name__)
14
+
15
+ try:
16
+ from cpm_kernels.kernels.base import LazyKernelCModule, KernelFunction, round_up
17
+
18
+ class Kernel:
19
+ def __init__(self, code: bytes, function_names: List[str]):
20
+ self.code = code
21
+ self._function_names = function_names
22
+ self._cmodule = LazyKernelCModule(self.code)
23
+
24
+ for name in self._function_names:
25
+ setattr(self, name, KernelFunction(self._cmodule, name))
26
+
27
+ quantization_code = "$QlpoOTFBWSZTWU9yuJUAQHN//////////f/n/8/n///n//bt4dTidcVx8X3V9FV/92/v4B7/AD5FBQFAAAChSgKpFCFAFVSigUAAAEKhSgUUqgFBKigqVREQAABQBQIANDTTIGI00BkZBkNGE0A0BkBkGQGRkaNAaAGQNBoGgDIAAYIGTI0DQAQAaGmmQMRpoDIyDIaMJoBoDIDIMgMjI0aA0AMgaDQNAGQAAwQMmRoGgAgA0NNMgYjTQGRkGQ0YTQDQGQGQZAZGRo0BoAZA0GgaAMgABggZMjQNABABoaaZAxGmgMjIMhowmgGgMgMgyAyMjRoDQAyBoNA0AZAADBAyZGgaAAmqU1NEgJqnptU/Sn4jRR6J6epk2pqb1Q/SgAPUGgyNNGjQ2SBpoAZAAGg0NB6mgDIAAAAA2oaApSREBNAARhGiYEaEwU8pvImlP0k2aam1GaGqbFNM1MHpTwmkepmyU9R6nqPKekHqNNPUxNGhp6n6p6QaZ6o9TG1GMqcoV9ly6nRanHlq6zPNbnGZNi6HSug+2nPiZ13XcnFYZW+45W11CumhzYhchOJ2GLLV1OBjBjGf4TptOddTSOcVxhqYZMYwZXZZY00zI1paX5X9J+b+f4e+x43RXSxXPOdquiGpduatGyXneN696M9t4HU2eR5XX/kPhP261NTx3JO1Ow7LyuDmeo9a7d351T1ZxnvnrvYnrXv/hXxPCeuYx2XsNmO003eg9J3Z6U7b23meJ4ri01OdzTk9BNO96brz+qT5nuvvH3ds/G+m/JcG/F2XYuhXlvO+jP7U3XgrzPN/lr8Sf1n6j4j7jZs+s/T0tNaNNYzTs12rxjwztHlnire3Nzc3N1wuBwOBwXBvZfoHpD7rFmR99V5vj3aXza3xdBbXMalubTg/jIv5dfAi54Pdc75j4z412n3Npj3Ld/ENm7a3b/Cod6h/ret1/5vn/C+l+gdslMvgPSLJ8d8q+U66fevYn/tW1chleEtNTGlcHCbLRlq0tHzF5tsbbZZfHjjLgZu42XCuC3NrdjTasZGNzgxPIrGqp7r3p7L2p5XjnpPSmTd5XtzqnB6U87zzg1Ol0zd0zsLszxR6lkxp35u6/teL0L0W922cR7Lu1lpL9CsHirzuM2T+BgsyViT6LHcm0/Vr6U/7LGGyJeqTEjt0PHWhF5mCT7R9mtlDwriYv0Tyr/OxYt6qp5r0mPVT0608TqnqMZaarU2nFwrTzzlrs1ed7z1ux60wyr4ydCaTi3enW8x68x0zU7tXSlcmPSW1mGpWJMg4zmPC2lK96tp0OE80y4MfEvnZj8zGluR6b22ki1Ou9V2nCd9xovcPvcYMZYy0lvN60ScZ45vN6yeCeeXFb1lVjnnCar5fwXwE2bzJ4HI1XVPXfXZMm44GUsMpYsmLB65TuVdm0cl0b+i/wGNN66XjeV7zuPpHcnK/juhhjdfId5jMdE5nN0dGmmm2zZs2cexD5n9p/dY352XsvXHaZNWWsmmS1atjR452nYudzvqv2HMRyvNNnlMcDl3R2+yx2uVrBubTW9icHDVtbNXlZm7jma1rM4VurZZd2y6nUau7ZXZ7bVU+mnoOVxZGMrVmvX60605JwmzGZhhhjTWtaaaMaaGTGmNMZasY0iX8VMUl8eepaIrzGSpemWOQyZORk2bNpjUybMmxqYmknCGCFynutfksaZpjTNMaaatM0xsxcGR0sociNqxNSmhhR1ZJPbsn8qyF0t2qH6iYBclclalbtTTcHTDsPaX6rlnElph2Jyumumtynv2Kk8GI7rsvXbIcJgHJOSaSXnnGaI3m87RtVXJOZ/YtgdTE6Wpha6ZlE8ayXkef1fh602r2WwvfMXtMdLlkfnLFdYYwYso+bWqm7yJqHXZGw2nrS5ZanSYnWlxBxMF1V940K2wdrI7R6OYf7DGGamMmTSbRhlS45xmVOumF1EyPCmHrrN8wwZOOrdNtLeMtzFzDlWnfTBxMk2NaXIZHBYxYLD4w8yju0ao65Vz1OIXoS9dLanwCe1PWrYuWMqf1if1z2k2yYfKJ741PDgno1ZQ8DRqvUny3mNoWTzGO6m1DkrJI8JiR5cSd+vZdGOO8nrMoc5+NDUFsMSXaZJeNlMmGLtJsovOsUp7I9S5VojKxF6bTVEelXqlfJobQr3LozSh2Jk7VcrVMfhXqszGWMzNqGhqZY0OadxkyyMssKugZR0KNFXBHlqwmJgTE/BNVMk6ItJXZMR0H47GpXv/DMOvNkmVuaV1PRfEdxuqc7Hcd+ZV/zTLaRxWk0nl9CdCeM6mn5rstHIBcpiuwmUZXeq81DacHI2rmrZ5SuE5mOZd6LQrZg9mx32TprA8BMo5jKN6yLTCi3WzQaZSuhzTtM1fUTGVpG8Tw+KXI0tjEpiWxtLYynOlktSbVlaI5kxP8TDH8kx50xoxi5KcA4pcja8KWLRlO/Ks6q06ergnvm1ca3Tq8Uw7LTUsmWyctXPWmpitl/uvGcWTGXGuAXDfhqazGmjkxcJW5hMMMMpYsXl2TZYtVOddG3XCarUt6Ptq9CZXSNzyuRzqRZOjsxdBbFVz6OA5HI43r1jityVlVpVkxmOsyaYWE1NTGq1sOVh36mHMcxtSvcy70edG0ZGR3I1Go1GRlV7mWWo1G0ZGRqlvH40l7o4m5xMWLLLYyNjnqc8556mdPqLJ31n/1nWOncxzG1tizrHs/Z+d2vP/B/l8wdJ6rHUn2nbbDq4p6htFtYzMMMTaZis1K5GKzGNmxhmUx2DDlZ/qNnIx41xnaMfCZWYaZWtNLTNW8ND4Fw1MyZOCdM428suKG1ehW8TesOydg7J+YYcD4cYR+8dFK6M4E3HM9ZfRNNL+Sn6rsl4DsrDl2HpPCnfxjGXtbZtYys1ttlyJ4T+BvexjGWRjMszK4Jpc77D3GyuVD7q0+G8m9G+2+rGm7cOR2y7FdtY2XUYx/oNlfRYxhMYyYZkyyg55enna9Kt/FFi6GMMwYwdwxWgxGMLKYmUyGExTKMZkMFhkymKuh0NOBNnBu+23LdwDoZYYzGGMxtORaTU1pjTGWTTGGtMrNWUsyyTTLLG1qy2ZjbK2DBllWqxMtBMaYZQmcE7zvvRcTkclUwdkxTaSdyySt/7fpL+T1v516Ji97fwr5JbLu305zMn5+GMTTZ9F+y7ExwmGVfG44yxn3dLv6l5i+Wth1jCrDq21nW9LqvvDzz3Vf3LLH/O/32TJ/erx3bXftO4eF+G956D952K/An4NfvOpjFjExjevP/UmE0fIoZXx6/w6lX/no3D0bLt+ixjieBM6ksRd0yB4Lt2SwYNE+gd1detlZWUnpiZfGfFaK+4PyCa/v18V8X75pe9fLXzp7l3VjF76vWZmHwGz1IZNWT7b8yddJ4q5kyrVdfru6atWc7bVYztL9Jf4GXvT+Y8m9/YsXP6H018a8D4XVOqvfzqeR+6yZOD8dPv0+U7/q5Pl+2dNb0MjzGVH5p6MNQ7cOWvw62U9aHE8DprDek+McLyvDz+te+9Zhq5+YTruufMcWMabqysTmZVWjKPfnK0wyVcrsuhjZRdLkHNvD72b9abriOSGIxiLixMOoalNPXzy+wT/tf+U6HHONfsz+xe8ufHBdQWWGWLA9if0rsnmrxK5LvRZQeWsTCsrmOYy8VteVfuRfcVTtDLItLIsMYxZLdU/DbtSemxF6Z6Zo5WBXE4tFdCyVMMXMTEMZXVlS6Xec2T4e0tHsRcEuWshcJ2YsNF5rUx1E8ifCq6Z+ZP7qdCeu/aTwFd53l16/o0NOw6O3dLavP4Hbi4RdmuDk6DoYaninC0+o4uZjbJ7Rxeu0/FbuFg+q7DVS6fQe0rZ6NDGUNNU6DEqOaLTicKnYZMnBWruljQxoaS3dZhocDge0bSTyOvdAbG5hxe2xji7E/L55xX13wWNDi6HCekcFxfCPGxY0MXC+s7afWaMdDyjyr+o8Rudm/NabOZvdl274zH4f5XK9z6On1Pe/K5TdPAslg77BjuO6Y3eO7GqvOPG/stknp1leyvLL0Z7bl9I4noMvLkzytLhWYzrOZzLXCORe028rORzOg4N/L0HlMOQ3Pgmnbb6KczlabORpu980q37TBqRu0/p3PO6234Bl03Ynuz+9W7gnsEcmvYaYY3aMYY0wx3pYd+ujsXauWdaY5Xkbtl23fPzFHiDB/QMo0yFjBllYxTQYYyxkrwn7JufwJ/PfgJ+C83X69ni6zvXcnyXabv0ncbLwsceS+RNlyN2mnneJtX0ngYO0+e+0+UnA+Wch3ji8hj5an4h+i6XBySU4n+R0roVcbw5yvHrmr4Yw8Y7x6c+9POPYHI5HI5HI5HI5HGXGww4nE4nrVyOR8XeqPEO7PLOiukYa3Novk5hV4cdtYZLI93e+uxff2jRo0aNGjRo0aNG1bVtW1dy3m83m8+tQ5ZzHw3nObwOu8La9Rc1dtkdS8A3eTk823tnktXWlxN6Oixe06zrN70Isd9jiOgZFq9yfkPqP/SLhN2Myl8jDM43bl1nbcb4cO57jlh8Jow6pzXZdL4dyODTuuhu77FyO27DdwdRxmvO+O+3N2+BdqyTwLHVczDVY4UPE4O66/ZO2cx1LFzVdSXtF7G4HMbrauOHRw6c8FdZ5m9fHZHYZXfTlZquyynSyTTKke6vcffSD9pzPA/G7n7jxPmuhc1DHMynPMrGL6AdewYmwu5ko+UUyTwrMv27rPH1v1nGqd87+p6N6LU8k3NEng53xXyHS97+44OSg/sy/hn+Se6yfYNjW0/uTgP+PvWYzLMmjhcLB/gGpri6H83/84eUXWT6T9Hsv7785z/7z4icpW+zfXypuR7rx/gMdZb1/wC678pcs8/2a3mDitGHxl9mfPlll5MafWWqxk/eYuTDgcNMzDGWLWvsuglNxs53GtN6uWpktlW1tZZYcuinMMWmnNnJydze3b2Y1McBxrBkXw799izLMZZYyy0TkbsGM4p03S2uVu5s/XXUdSdec6smVxZYYGpVmT8A+8ajuEyV5FatkvVru2x6uxGXXbH4A+jvgP4GMYy3iPLXzq/6z65+E005ey+cwMZD3fZcqc6xpjTFjQ0P3U+e++cPYmTIwj0nrK5NPTfl3WvpfLtXDcb2HQMudYOxFXQBor4L4T6vrOauFctYXJQ++NUWmJe5bmx1jDiZS1dTqWxo4GR8jm3fttpmPHppk9PEyv4/y8/sO07XacOmcqc0x2Vi9BvNJvN5oW8x4mOsydpidRxMYJPx06m1bqPzq9KtK8sxXNXFodD/+MYYaJTLwOhc9brCsV18oOR1i4tXChyTkq4lf4y1Ke+9axjDHqs1mfBbMXuP4Hzi+X7t8vzv7bHerrUPgPCxhjre4fXdfLNtNM+Jd+Zdh8xd8wP87uNPoPgv4W7/5P2BuxfsMabNnMnza+54Pdi5U671GPZY8CehX8Voeoo7FHpkeEc6715FwHZrIrUrHaviPUbPZHND+IhczrP6FcYvhOZ0Di/ETt0OI+YwNWR9r7tpf6WDeZKZDB1+z2IthOl1mPyb5FluvEx9h9d0NnM0Y1XPFkWIsk1WotJ0PBMmkvjvQTd0e71tfeV+8r8lQ/tpzpsmxJ+InrI/dj2UajUajVTUajatRqNRtGo1Go1Go4wjeMpZFMVV9CHbofPraLsJ3JpWV2XOoanCuFky4y3PPNxucK2uKC1Lbdb1eo+m5XomN6HfeZsabHLHRX/K+offtNGGmHWctcVcG44MdSqsOLY9VzX+Zxfxn2HPdWTpzWvkrtJ8M5zorrKcquRytJ5N5DZmcaW02l76nWO+BqPXm1A2Ry/0q71dH/mqrqeFjkYxjEXtsX8qubTk67rGycyqsdm4tZx5D6D5hhi0waaWmiaMP81Yjii5qxPlPuU/GfTL1Y5E6Jyfiq63qTa39A4J0sOGDgO9WF9bOXl0XfPRbsY2bPNKPy1YrFYrFYmRhhlTIyMjJWJYZHXuCXI8OoXsvfljGLFicNifpp2XunoPiG1wtx3p1Tah+/DD66OnVtVXP9rKbVxOnL0tR/rHtqB5UDErUVcl11D4qqvjpOcxX7armUNJB3LpW6bxVvD08e8h3odKKvyCFZBdSh2FVcST9xV3n3T8t1j7Kr9qgrqXg+13Pt5U7JCvFXVIV1YG5lRhkVYZJYYDDD4KOIMoHCp26WS8GB7uBh2zIdgq/PKyInjV2STShuoapUdCpX1yTwqq/z1VvET7Kh5nVPkO8YyxjLt2MaaMmWTLQvx3qnzltnXW0p2jxgbEtSny/Osv8Y9pLMXYoHVPAhkVdWVeODhR6q9/Sxe2liwwZWMVvFXfRkeIDxAePUPIrdJ4ey6yquzH+PD/bUOWAu05qVHtFd8rrKHSoeNIOUqrYr3FXyToqfYJgwmJdKpXXOwYYegNNGMzfZPp/t3t/DVs4zjNTN61rRqaWaa4NYbRjTa0tWwy2Y2tGN8ZO8ofNKq4j9SL7I+cSm4/6ovLV5HNXLI0jJidwrtk6ynCaP6Z++GjRlWS3tLeW129Mi9evxU9mtz6s5J3Z7M2ngTgnKvmpomxpaLCzPfmx0JWE+m3NLDDGOX47RctdYYNK5jakdqLkRlI39n590T5zctGSwwZZDJj6kW8XSi6ot2MmWWJ0DUT3nuvebBudScjZ79g8cWJ8av0k+/bE5WKd5MdbFpbDVMxu1DVMmtNZGJvq1mtRbn6M+g/kP0FwDwr7quZs7xosNGpbscyxhhd9TyJyFwbLcxlTasg75vW7TsV5K7ji44XPMMrdoj+Y3rT0Hie62nlYV/pwczzOmdLqLhYkzGMzCZWGMQzGMSsZYY6Di1t4nlJ+Em63mJxrVLxPbYxNEdgc1dU2iOKyoYYWjNrEeHTYybVk0atSa7ehuwsWMWTqn1TrnS6hYsi71d1+s+k+ic70e20fzE/VaTdxT9ZtU4GIXdeNx3X77guYYfpHeTQjaMX6brOu4OY4K7Y2d9mbHarI5ox3p4GpJ2Vd/Tst60f7j999pppjR+Q/Qf8J/VaORs3cji7FfFuN61+ui9s8hix1OCh5KGVV23BPXvZfz3CLyHpix+exi8z/KnCnosY2eunor+cxyPO/xJ0vKey9OvE9VjqaYu0x3Z3jd6o2b1T12D+F8l232lwaaacD5LE8LBxu7WTlbWraWpew8Xexjel3E+wWD4APITdNqR8F3R3T0lunCQ4GaE9R37DxeCYfcHi4xci5ovKfxVs55y2hf+65E/Xdp6jR5nrebTmi5incpkyOjs50JvrZwstbbW6kfuuQw+2mykf/EXNFzxfKTrxew929TR6bWnGL//F3JFOFCQT3K4lQ"
28
+
29
+ kernels = Kernel(
30
+ bz2.decompress(base64.b64decode(quantization_code)),
31
+ [
32
+ "int4WeightCompression",
33
+ "int4WeightExtractionFloat",
34
+ "int4WeightExtractionHalf",
35
+ "int8WeightExtractionFloat",
36
+ "int8WeightExtractionHalf",
37
+ ],
38
+ )
39
+ except Exception as exception:
40
+ kernels = None
41
+ logger.warning("Failed to load cpm_kernels:" + str(exception))
42
+
43
+
44
+ class W8A16Linear(torch.autograd.Function):
45
+ @staticmethod
46
+ def forward(ctx, inp: torch.Tensor, quant_w: torch.Tensor, scale_w: torch.Tensor, weight_bit_width):
47
+ ctx.inp_shape = inp.size()
48
+ ctx.weight_bit_width = weight_bit_width
49
+ out_features = quant_w.size(0)
50
+ inp = inp.contiguous().view(-1, inp.size(-1))
51
+ weight = extract_weight_to_half(quant_w, scale_w, weight_bit_width)
52
+ ctx.weight_shape = weight.size()
53
+ output = inp.mm(weight.t())
54
+ ctx.save_for_backward(inp, quant_w, scale_w)
55
+ return output.view(*(ctx.inp_shape[:-1] + (out_features,)))
56
+
57
+ @staticmethod
58
+ def backward(ctx, grad_output: torch.Tensor):
59
+ inp, quant_w, scale_w = ctx.saved_tensors
60
+ weight = extract_weight_to_half(quant_w, scale_w, ctx.weight_bit_width)
61
+ grad_output = grad_output.contiguous().view(-1, weight.size(0))
62
+ grad_input = grad_output.mm(weight)
63
+ grad_weight = grad_output.t().mm(inp)
64
+ return grad_input.view(ctx.inp_shape), grad_weight.view(ctx.weight_shape), None, None
65
+
66
+
67
+ def compress_int4_weight(weight: torch.Tensor): # (n, m)
68
+ with torch.cuda.device(weight.device):
69
+ n, m = weight.size(0), weight.size(1)
70
+ assert m % 2 == 0
71
+ m = m // 2
72
+ out = torch.empty(n, m, dtype=torch.int8, device="cuda")
73
+ stream = torch.cuda.current_stream()
74
+
75
+ gridDim = (n, 1, 1)
76
+ blockDim = (min(round_up(m, 32), 1024), 1, 1)
77
+
78
+ kernels.int4WeightCompression(
79
+ gridDim,
80
+ blockDim,
81
+ 0,
82
+ stream,
83
+ [ctypes.c_void_p(weight.data_ptr()), ctypes.c_void_p(out.data_ptr()), ctypes.c_int32(n), ctypes.c_int32(m)],
84
+ )
85
+ return out
86
+
87
+
88
+ def extract_weight_to_half(weight: torch.Tensor, scale_list: torch.Tensor, source_bit_width: int):
89
+ assert scale_list.dtype in [torch.half, torch.bfloat16]
90
+ assert weight.dtype in [torch.int8]
91
+ if source_bit_width == 8:
92
+ return weight.to(scale_list.dtype) * scale_list[:, None]
93
+ elif source_bit_width == 4:
94
+ func = (
95
+ kernels.int4WeightExtractionHalf if scale_list.dtype == torch.half else kernels.int4WeightExtractionBFloat16
96
+ )
97
+ else:
98
+ assert False, "Unsupported bit-width"
99
+
100
+ with torch.cuda.device(weight.device):
101
+ n, m = weight.size(0), weight.size(1)
102
+ out = torch.empty(n, m * (8 // source_bit_width), dtype=scale_list.dtype, device="cuda")
103
+ stream = torch.cuda.current_stream()
104
+
105
+ gridDim = (n, 1, 1)
106
+ blockDim = (min(round_up(m, 32), 1024), 1, 1)
107
+
108
+ func(
109
+ gridDim,
110
+ blockDim,
111
+ 0,
112
+ stream,
113
+ [
114
+ ctypes.c_void_p(weight.data_ptr()),
115
+ ctypes.c_void_p(scale_list.data_ptr()),
116
+ ctypes.c_void_p(out.data_ptr()),
117
+ ctypes.c_int32(n),
118
+ ctypes.c_int32(m),
119
+ ],
120
+ )
121
+ return out
122
+
123
+
124
+ class QuantizedLinear(torch.nn.Module):
125
+ def __init__(self, weight_bit_width: int, weight, bias=None, device="cpu", dtype=None, empty_init=False, *args,
126
+ **kwargs):
127
+ super().__init__()
128
+ self.weight_bit_width = weight_bit_width
129
+
130
+ shape = weight.shape
131
+
132
+ if weight is None or empty_init:
133
+ self.weight = torch.empty(shape[0], shape[1] * weight_bit_width // 8, dtype=torch.int8, device=device)
134
+ self.weight_scale = torch.empty(shape[0], dtype=dtype, device=device)
135
+ else:
136
+ self.weight_scale = weight.abs().max(dim=-1).values / ((2 ** (weight_bit_width - 1)) - 1)
137
+ self.weight = torch.round(weight / self.weight_scale[:, None]).to(torch.int8)
138
+ if weight_bit_width == 4:
139
+ self.weight = compress_int4_weight(self.weight)
140
+
141
+ self.weight = Parameter(self.weight.to(device), requires_grad=False)
142
+ self.weight_scale = Parameter(self.weight_scale.to(device), requires_grad=False)
143
+ self.bias = Parameter(bias.to(device), requires_grad=False) if bias is not None else None
144
+
145
+ def forward(self, input):
146
+ output = W8A16Linear.apply(input, self.weight, self.weight_scale, self.weight_bit_width)
147
+ if self.bias is not None:
148
+ output = output + self.bias
149
+ return output
150
+
151
+
152
+ def quantize(model, weight_bit_width, empty_init=False, device=None):
153
+ """Replace fp16 linear with quantized linear"""
154
+ for layer in model.layers:
155
+ layer.self_attention.query_key_value = QuantizedLinear(
156
+ weight_bit_width=weight_bit_width,
157
+ weight=layer.self_attention.query_key_value.weight.to(torch.cuda.current_device()),
158
+ bias=layer.self_attention.query_key_value.bias,
159
+ dtype=layer.self_attention.query_key_value.weight.dtype,
160
+ device=layer.self_attention.query_key_value.weight.device if device is None else device,
161
+ empty_init=empty_init
162
+ )
163
+ layer.self_attention.dense = QuantizedLinear(
164
+ weight_bit_width=weight_bit_width,
165
+ weight=layer.self_attention.dense.weight.to(torch.cuda.current_device()),
166
+ bias=layer.self_attention.dense.bias,
167
+ dtype=layer.self_attention.dense.weight.dtype,
168
+ device=layer.self_attention.dense.weight.device if device is None else device,
169
+ empty_init=empty_init
170
+ )
171
+ layer.mlp.dense_h_to_4h = QuantizedLinear(
172
+ weight_bit_width=weight_bit_width,
173
+ weight=layer.mlp.dense_h_to_4h.weight.to(torch.cuda.current_device()),
174
+ bias=layer.mlp.dense_h_to_4h.bias,
175
+ dtype=layer.mlp.dense_h_to_4h.weight.dtype,
176
+ device=layer.mlp.dense_h_to_4h.weight.device if device is None else device,
177
+ empty_init=empty_init
178
+ )
179
+ layer.mlp.dense_4h_to_h = QuantizedLinear(
180
+ weight_bit_width=weight_bit_width,
181
+ weight=layer.mlp.dense_4h_to_h.weight.to(torch.cuda.current_device()),
182
+ bias=layer.mlp.dense_4h_to_h.bias,
183
+ dtype=layer.mlp.dense_4h_to_h.weight.dtype,
184
+ device=layer.mlp.dense_4h_to_h.weight.device if device is None else device,
185
+ empty_init=empty_init
186
+ )
187
+
188
+ return model
checkpoint-3000/rng_state_0.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:98d65493cf73c9e9b5e7761466376515ae12ce70f43ca0eb634ab33477f2a65f
3
+ size 17655
checkpoint-3000/rng_state_1.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b44cdbbd91d15b5951be6d2a1e5d65141cc20b30c1bb91a7b6db904af71da3cf
3
+ size 17655
checkpoint-3000/rng_state_2.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:fd9f3794e660af32407f5f088f5128bab1c509ca125bfd35edbb3ebffc0fcff5
3
+ size 17655
checkpoint-3000/rng_state_3.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ac8a17aba25e6596c3f91822289a72fba5c858852b8146f6613fdd3e9e6cb1dd
3
+ size 17655