finetune

modeling_chatglm.py

@@ -3,7 +3,6 @@ import json
 import math
 import copy
 import warnings
-import re
 import sys
 
 import torch
@@ -30,6 +29,7 @@ from .configuration_chatglm import ChatGLMConfig
 
 try:
     from transformers.utils import is_flash_attn_greater_or_equal_2_10, is_flash_attn_2_available
+
     if is_flash_attn_2_available():
         from flash_attn import flash_attn_func, flash_attn_varlen_func
         from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
@@ -215,6 +215,7 @@ class RMSNorm(torch.nn.Module):
         return (self.weight * hidden_states).to(input_dtype)
 
 
+
 class CoreAttention(torch.nn.Module):
     def __init__(self, config: ChatGLMConfig, layer_number):
         super(CoreAttention, self).__init__()
@@ -332,130 +333,6 @@ class CoreAttention(torch.nn.Module):
         return context_layer
 
 
-class SdpaAttention(CoreAttention):
-    def forward(self, query_layer, key_layer, value_layer, attention_mask):
-        if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]:
-            context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
-                                                                             is_causal=True,
-                                                                             dropout_p=self.config.attention_dropout if self.training else 0.0)
-        else:
-            if attention_mask is not None:
-                attention_mask = ~attention_mask
-            context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
-                                                                             attention_mask,
-                                                                             dropout_p=self.config.attention_dropout if self.training else 0.0)
-        context_layer = context_layer.transpose(1, 2).contiguous()
-        new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
-        context_layer = context_layer.reshape(*new_context_layer_shape)
-        return context_layer
-
-
-def _get_unpad_data(attention_mask):
-    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
-    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
-    max_seqlen_in_batch = seqlens_in_batch.max().item()
-    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
-    return (
-        indices,
-        cu_seqlens,
-        max_seqlen_in_batch,
-    )
-
-
-# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2
-class FlashAttention2(CoreAttention):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
-
-    def forward(self, query_states, key_states, value_states, attention_mask):
-        query_states = query_states.transpose(1, 2)
-        key_states = key_states.transpose(1, 2)
-        value_states = value_states.transpose(1, 2)
-        batch_size, query_length = query_states.shape[:2]
-        if not self._flash_attn_uses_top_left_mask:
-            causal = self.is_causal
-        else:
-            # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
-            causal = self.is_causal and query_length != 1
-        dropout = self.config.attention_dropout if self.training else 0.0
-        # Contains at least one padding token in the sequence
-        if attention_mask is not None:
-            query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
-                query_states, key_states, value_states, attention_mask, query_length
-            )
-
-            cu_seqlens_q, cu_seqlens_k = cu_seq_lens
-            max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
-
-            attn_output_unpad = flash_attn_varlen_func(
-                query_states,
-                key_states,
-                value_states,
-                cu_seqlens_q=cu_seqlens_q,
-                cu_seqlens_k=cu_seqlens_k,
-                max_seqlen_q=max_seqlen_in_batch_q,
-                max_seqlen_k=max_seqlen_in_batch_k,
-                dropout_p=dropout,
-                softmax_scale=None,
-                causal=causal,
-            )
-
-            attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
-        else:
-            attn_output = flash_attn_func(
-                query_states, key_states, value_states, dropout, softmax_scale=None, causal=causal
-            )
-        attn_output = attn_output.reshape(batch_size, query_length, self.hidden_size_per_partition).contiguous()
-        return attn_output
-
-    def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
-        indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
-        batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
-
-        key_layer = index_first_axis(
-            key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
-        )
-        value_layer = index_first_axis(
-            value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
-        )
-        if query_length == kv_seq_len:
-            query_layer = index_first_axis(
-                query_layer.reshape(batch_size * kv_seq_len, self.num_attention_heads_per_partition, head_dim),
-                indices_k
-            )
-            cu_seqlens_q = cu_seqlens_k
-            max_seqlen_in_batch_q = max_seqlen_in_batch_k
-            indices_q = indices_k
-        elif query_length == 1:
-            max_seqlen_in_batch_q = 1
-            cu_seqlens_q = torch.arange(
-                batch_size + 1, dtype=torch.int32, device=query_layer.device
-            )  # There is a memcpy here, that is very bad.
-            indices_q = cu_seqlens_q[:-1]
-            query_layer = query_layer.squeeze(1)
-        else:
-            # The -q_len: slice assumes left padding.
-            attention_mask = attention_mask[:, -query_length:]
-            query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
-
-        return (
-            query_layer,
-            key_layer,
-            value_layer,
-            indices_q,
-            (cu_seqlens_q, cu_seqlens_k),
-            (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
-        )
-
-
-CORE_ATTENTION_CLASSES = {
-    "eager": CoreAttention,
-    "sdpa": SdpaAttention,
-    "flash_attention_2": FlashAttention2
-}
-
-
 class SelfAttention(torch.nn.Module):
     """Parallel self-attention layer abstract class.
 
@@ -820,18 +697,12 @@ class ChatGLMPreTrainedModel(PreTrainedModel):
     config_class = ChatGLMConfig
     base_model_prefix = "transformer"
     _no_split_modules = ["GLMBlock"]
-    _supports_flash_attn_2 = True
-    _supports_sdpa = True
 
     def _init_weights(self, module: nn.Module):
         """Initialize the weights."""
         return
 
     def get_masks(self, input_embeds, past_key_values, padding_mask=None):
-        if self.config._attn_implementation == "flash_attention_2":
-            if padding_mask is not None and not padding_mask.all():
-                return padding_mask
-            return None
         batch_size, seq_length, embed_size = input_embeds.shape
         full_attention_mask = torch.ones(batch_size, seq_length, seq_length, device=input_embeds.device)
         full_attention_mask.tril_()
@@ -978,7 +849,6 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
             # not allow for inputs_embeds, because we want to process image feature
             assert input_ids is not None and inputs_embeds is None, f"{input_ids} {inputs_embeds}"
             if not is_empty(images):  # multi-modality
-
                 image_size: int = self.config.vision_config['image_size']
                 patch_size: int = self.config.vision_config['patch_size']
                 num_patches = (image_size // patch_size // 2) ** 2
@@ -998,8 +868,7 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
                         self.config.eoi_token_id)
                     assert eoi_token_pos - boi_token_pos == 2
                     new_input_embeds.append(torch.cat(
-                        (inputs_embeds[i, :boi_token_pos], images_features[i].to(inputs_embeds.device),
-                         inputs_embeds[i, eoi_token_pos + 1:])))
+                        (inputs_embeds[i, :boi_token_pos], images_features[i], inputs_embeds[i, eoi_token_pos + 1:])))
                     new_position_ids.append(torch.cat(
                         (position_ids[i, :boi_token_pos + 1], position_ids[i, boi_token_pos + 1].repeat(num_patches),
                          position_ids[i, eoi_token_pos:])
@@ -1015,9 +884,6 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
 
         batch_size, seq_length = input_ids.shape
 
-        if inputs_embeds is None:
-            inputs_embeds = self.embedding(input_ids)
-
         if self.pre_seq_len is not None:
             if past_key_values is None:
                 past_key_values = self.get_prompt(batch_size=batch_size, device=input_ids.device,
@@ -1028,10 +894,32 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
 
         if full_attention_mask is None:
             if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1):
+                if self.training:
+                    # https://github.com/THUDM/GLM-4/issues/264
+                    new_input_ids, new_attention_mask = [], []
+                    for i in range(len(input_ids)):
+                        input_id = input_ids[i].tolist()
+                        boi_token_pos, eoi_token_pos = input_id.index(self.config.boi_token_id), input_id.index(self.config.eoi_token_id)
+                        assert eoi_token_pos - boi_token_pos == 2
+
+                        new_attention_mask.append(torch.cat(
+                            (attention_mask[i, :boi_token_pos + 1], torch.ones(num_patches).to(attention_mask.device),
+                             attention_mask[i, eoi_token_pos:])))
+
+                        new_input_ids.append(torch.cat(
+                            (input_ids[i, :boi_token_pos + 1], input_ids[i, -1].repeat(num_patches),
+                             input_ids[i, eoi_token_pos:])))
+
+                    attention_mask = torch.stack(new_attention_mask, dim=0)
+                    input_ids = torch.stack(new_input_ids, dim=0)
+
+                if inputs_embeds is None:
+                        inputs_embeds = self.embedding(input_ids)
                 full_attention_mask = self.get_masks(inputs_embeds, past_key_values, padding_mask=attention_mask)
 
         # Rotary positional embeddings
         rotary_pos_emb = self.rotary_pos_emb(self.seq_length)
+
         if position_ids is not None:
             rotary_pos_emb = rotary_pos_emb[position_ids]
         else:
@@ -1189,6 +1077,22 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
 
         loss = None
         if labels is not None:
+            # https://github.com/THUDM/GLM-4/issues/264
+            new_labels = []
+            for i in range(len(input_ids)):
+                input_id = input_ids[i].tolist()
+                boi_token_pos, eoi_token_pos = input_id.index(self.config.boi_token_id), input_id.index(
+                    self.config.eoi_token_id)
+                assert eoi_token_pos - boi_token_pos == 2
+
+                new_labels.append(torch.cat(
+                    (
+                        labels[i, :boi_token_pos + 1],
+                        torch.tensor([-100]).to(labels.device).to(labels.dtype).repeat(1600),
+                        labels[i, eoi_token_pos:])))  # 在两个token之间加入
+
+            labels = torch.stack(new_labels, dim=0)
+
             lm_logits = lm_logits.to(torch.float32)
 
             # Shift so that tokens < n predict n