Slim embedding

config.json

@@ -1,5 +1,5 @@
 {
-  "_name_or_path": "THUDM/chatglm-6b",
+  "_name_or_path": "THUDM/chatglm-6b-int4-qe",
   "architectures": [
     "ChatGLMForConditionalGeneration"
   ],
@@ -8,21 +8,23 @@
     "AutoModel": "modeling_chatglm.ChatGLMForConditionalGeneration",
     "AutoModelForSeq2SeqLM": "modeling_chatglm.ChatGLMForConditionalGeneration"
   },
-  "bos_token_id": 150004,
-  "eos_token_id": 150005,
-  "pad_token_id": 20003,
+  "bos_token_id": 130004,
+  "eos_token_id": 130005,
+  "gmask_token_id": 130001,
   "hidden_size": 4096,
   "inner_hidden_size": 16384,
   "layernorm_epsilon": 1e-05,
+  "mask_token_id": 130000,
   "max_sequence_length": 2048,
   "model_type": "chatglm",
   "num_attention_heads": 32,
   "num_layers": 28,
+  "pad_token_id": 3,
   "position_encoding_2d": true,
   "quantization_bit": 4,
   "quantization_embeddings": true,
   "torch_dtype": "float16",
-  "transformers_version": "4.26.1",
+  "transformers_version": "4.27.1",
   "use_cache": true,
-  "vocab_size": 150528
-}
+  "vocab_size": 130528
+}

configuration_chatglm.py

@@ -66,6 +66,8 @@ class ChatGLMConfig(PretrainedConfig):
             use_cache=False,
             bos_token_id=150004,
             eos_token_id=150005,
+            mask_token_id=150000,
+            gmask_token_id=150001,
             pad_token_id=0,
             max_sequence_length=2048,
             inner_hidden_size=16384,
@@ -87,6 +89,8 @@ class ChatGLMConfig(PretrainedConfig):
         self.bos_token_id = bos_token_id
         self.eos_token_id = eos_token_id
         self.pad_token_id = pad_token_id
+        self.mask_token_id = mask_token_id
+        self.gmask_token_id = gmask_token_id
         self.position_encoding_2d = position_encoding_2d
         self.quantization_bit = quantization_bit
         self.quantization_embeddings = quantization_embeddings

modeling_chatglm.py

@@ -921,9 +921,9 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
 
 
             if position_ids is None:
-                MASK, gMASK = 150000, 150001
-                mask_token = MASK if MASK in input_ids else gMASK
-                use_gmask = False if MASK in input_ids else True
+                MASK, gMASK = self.config.mask_token_id, self.config.gmask_token_id
+                mask_token = gMASK if gMASK in input_ids else MASK
+                use_gmask = True if gMASK in input_ids else False
 
                 mask_positions = [seq.tolist().index(mask_token) for seq in input_ids]
                 position_ids = self.get_position_ids(
@@ -1084,9 +1084,9 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
             **kwargs
     ) -> dict:
         batch_size, seq_length = input_ids.shape
-        MASK, gMASK = 150000, 150001
-        mask_token = MASK if MASK in input_ids else gMASK
-        use_gmask = False if MASK in input_ids else True
+        MASK, gMASK = self.config.mask_token_id, self.config.gmask_token_id
+        mask_token = gMASK if gMASK in input_ids else MASK
+        use_gmask = True if gMASK in input_ids else False
         seqs = input_ids.tolist()
         mask_positions = [seq.index(mask_token) for seq in seqs]
 
@@ -1408,6 +1408,11 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
 
         self.transformer = quantize(self.transformer, bits, use_quantization_cache=use_quantization_cache, empty_init=empty_init, **kwargs)
 
+        if self.device == torch.device("cpu"):
+            dtype = torch.float32
+        else:
+            dtype = torch.half
+
         if quantize_embeddings:
             logger.info("Applying quantization to embeddings")
             self.transformer.word_embeddings = QuantizedEmbedding(
@@ -1415,11 +1420,11 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
                 weight_tensor=self.transformer.word_embeddings.weight.to(self.device),
                 num_embeddings=self.transformer.word_embeddings.num_embeddings,
                 embedding_dim=self.transformer.word_embeddings.embedding_dim,
-                dtype=torch.half,
-                empty_init=True,
+                dtype=dtype,
+                empty_init=empty_init,
                 device=self.transformer.word_embeddings.weight.device,
             )
-            self.lm_head =  QuantizedLinear(
+            self.lm_head = QuantizedLinear(
                 weight_bit_width=bits,
                 weight_tensor=self.lm_head.weight.to(self.device),
                 bias_tensor=None,
@@ -1428,8 +1433,8 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
                 bias=False,
                 quantized_weight=self.transformer.word_embeddings.weight,
                 quantized_weight_scale=self.transformer.word_embeddings.weight_scale,
-                dtype=torch.half,
-                empty_init=True,
+                dtype=dtype,
+                empty_init=empty_init,
                 device=self.lm_head.weight.device,
             )
 

quantization.py

@@ -369,7 +369,7 @@ class QuantizedEmbedding(Embedding):  # TODO: backward, check empty_init
                 )
                 self.weight_scale = torch.empty(shape[0], dtype=kwargs["dtype"], device=kwargs["device"])
             else:
-                self.weight_scale = (weight_tensor.abs().max(dim=-1).values / ((2 ** (weight_bit_width - 1)) - 1)).half()
+                self.weight_scale = (weight_tensor.abs().max(dim=-1).values / ((2 ** (weight_bit_width - 1)) - 1)).to(kwargs["dtype"])
                 self.weight = torch.round(weight_tensor / self.weight_scale[:, None]).to(torch.int8)
                 if weight_bit_width == 4:
                     self.weight = compress_int4_weight(self.weight)

tokenization_chatglm.py

@@ -48,11 +48,13 @@ class SPTokenizer:
     def __init__(
             self,
             vocab_file,
+            num_image_tokens=20000,
             max_blank_length=80,
             byte_fallback=True,
     ):
         assert vocab_file is not None
         self.vocab_file = vocab_file
+        self.num_image_tokens = num_image_tokens
         self.special_tokens = ["[MASK]", "[gMASK]", "[sMASK]", "<unused_0>", "<sop>", "<eop>", "<ENC>", "<dBLOCK>"]
         self.max_blank_length = max_blank_length
         self.byte_fallback = byte_fallback
@@ -70,10 +72,6 @@ class SPTokenizer:
     def get_tab_token():
         return f"<|tab|>"
 
-    @property
-    def num_image_tokens(self):
-        return 20000
-
     @property
     def num_text_tokens(self):
         return self.text_tokenizer.num_tokens
@@ -178,6 +176,7 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
             mask_token='[MASK]',
             gmask_token='[gMASK]',
             padding_side="left",
+            num_image_tokens=20000,
             **kwargs
     ) -> None:
         super().__init__(
@@ -197,10 +196,16 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
         self.mask_token = mask_token
         self.gmask_token = gmask_token
 
-        self.sp_tokenizer = SPTokenizer(vocab_file)
+        self.sp_tokenizer = SPTokenizer(vocab_file, num_image_tokens=num_image_tokens)
 
         """ Initialisation """
 
+    @property
+    def gmask_token_id(self) -> Optional[int]:
+        if self.gmask_token is None:
+            return None
+        return self.convert_tokens_to_ids(self.gmask_token)
+
     @property
     def eop_token_id(self) -> Optional[int]:
         """

tokenizer_config.json

@@ -10,6 +10,7 @@
   "remove_space": false,
   "do_lower_case": false,
   "tokenizer_class": "ChatGLMTokenizer",
+  "num_image_tokens": 0,
   "auto_map": {
     "AutoTokenizer": [
       "tokenization_chatglm.ChatGLMTokenizer",