sync with main

modeling_chatglm.py

@@ -40,12 +40,6 @@ logger = logging.get_logger(__name__)
 _CHECKPOINT_FOR_DOC = "THUDM/ChatGLM"
 _CONFIG_FOR_DOC = "ChatGLMConfig"
 
-CHATGLM_6B_PRETRAINED_MODEL_ARCHIVE_LIST = [
-    "THUDM/chatglm3-6b",
-    # See all ChatGLM models at https://huggingface.co/models?filter=chatglm
-]
-
-
 def default_init(cls, *args, **kwargs):
     return cls(*args, **kwargs)
 
@@ -253,15 +247,12 @@ class CoreAttention(torch.nn.Module):
             # This is actually dropping out entire tokens to attend to, which might
             # seem a bit unusual, but is taken from the original Transformer paper.
             attention_probs = self.attention_dropout(attention_probs)
-            # =========================
-            # Context layer. [sq, b, hp]
-            # =========================
-
-            # value_layer -> context layer.
-            # [sk, b, np, hn] --> [b, np, sq, hn]
 
+            # query layer shape: [b * np, sq, hn]
+            # value layer shape: [b, np, sk, hn]
+            # attention shape: [b, np, sq, sk]
             # context layer shape: [b, np, sq, hn]
-            output_size = (value_layer.size(1), value_layer.size(2), query_layer.size(0), value_layer.size(3))
+            output_size = (value_layer.size(0), value_layer.size(1), query_layer.size(1), value_layer.size(3))
             # change view [b * np, sk, hn]
             value_layer = value_layer.view(output_size[0] * output_size[1], value_layer.size(2), -1)
             # change view [b * np, sq, sk]
@@ -386,7 +377,10 @@ class SelfAttention(torch.nn.Module):
             key_layer = torch.cat((cache_k, key_layer), dim=2)
             value_layer = torch.cat((cache_v, value_layer), dim=2)
         if use_cache:
-            kv_cache = (key_layer, value_layer)
+            if kv_cache is None:
+                kv_cache = torch.cat((key_layer.unsqueeze(0).unsqueeze(0), value_layer.unsqueeze(0).unsqueeze(0)), dim=1)
+            else:
+                kv_cache = (key_layer, value_layer)
         else:
             kv_cache = None
 
@@ -605,7 +599,7 @@ class GLMTransformer(torch.nn.Module):
                     hidden_states,
                     attention_mask=attention_mask,
                     rotary_pos_emb=rotary_pos_emb,
-                    kv_cache=kv_caches[index],
+                    kv_caches=kv_caches[index],
                     use_cache=use_cache,
                     use_reentrant=False
                 )
@@ -619,7 +613,15 @@ class GLMTransformer(torch.nn.Module):
                 )
             hidden_states, kv_cache = layer_ret
             if use_cache:
-                presents = presents + (kv_cache,)
+                # token by token decoding, use tuple format
+                if kv_caches[0] is not None:
+                    presents = presents + (kv_cache,)
+                # prefilling in decoding, use tensor format to save cuda memory
+                else:
+                    if len(presents) == 0:
+                        presents = kv_cache
+                    else:
+                        presents = torch.cat((presents, kv_cache.to(presents.device)), dim=0)
 
         if output_hidden_states:
             all_hidden_states = all_hidden_states + (hidden_states,)
@@ -773,6 +775,12 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
             inputs_embeds, full_attention_mask, rotary_pos_emb=rotary_pos_emb,
             kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states
         )
+        if presents is not None and type(presents) is torch.Tensor:
+            presents = presents.split(1, dim=0)
+            presents = list(presents)
+            presents = [list(x.squeeze(0).split(1, dim=0)) for x in presents]
+            presents = [tuple([x.squeeze(0) for x in y]) for y in presents]
+            presents = tuple(presents)
 
         if not return_dict:
             return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)