Recover

modeling_cocom.py

@@ -3,45 +3,11 @@ import torch
 import math 
 from peft import get_peft_model, LoraConfig, TaskType
 import os
-from flash_attn.flash_attn_interface import flash_attn_func
-import torch.nn as nn
-import torch
 
 def freeze_model(model):
     for param in model.parameters():
         param.requires_grad = False
 
-class CustomFlashAttention(nn.Module):
-    def __init__(self, embed_dim, num_heads, dropout=0.0):
-        super().__init__()
-        self.embed_dim = embed_dim
-        self.num_heads = num_heads
-        self.dropout = dropout
-        self.head_dim = embed_dim // num_heads
-        assert self.head_dim * num_heads == embed_dim, "Embedding size must be divisible by the number of heads."
-
-        # Define projection layers
-        self.qkv_proj = nn.Linear(embed_dim, 3 * embed_dim)
-        self.out_proj = nn.Linear(embed_dim, embed_dim)
-
-    def forward(self, hidden_states):
-        batch_size, seq_len, embed_dim = hidden_states.size()
-        qkv = self.qkv_proj(hidden_states)  # Project to Q, K, V
-        qkv = qkv.view(batch_size, seq_len, 3, self.num_heads, self.head_dim)
-        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, batch_size, num_heads, seq_len, head_dim)
-        query, key, value = qkv[0], qkv[1], qkv[2]
-
-        # FlashAttention expects contiguous inputs
-        query = query.contiguous()
-        key = key.contiguous()
-        value = value.contiguous()
-
-        # Apply FlashAttention
-        attn_output, _ = flash_attn_func(query, key, value, dropout_p=self.dropout, causal=False)
-
-        # Reshape and project back to the original dimension
-        attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, embed_dim)
-        return self.out_proj(attn_output)
 
 class BERT_Compressor(torch.nn.Module):
     def __init__(self, compr_model_name, compr_rate, compr_linear_type, decoder_hidden_size):
@@ -109,7 +75,7 @@ class COCOMConfig(PretrainedConfig):
                 device_map = "cuda",
                  **kwargs):
         super().__init__(**kwargs)
-        
+
         self.decoder_model_name = decoder_model_name # model name of decoder
         self.quantization = quantization # quantization, could be no, int4, int8
         self.generation_top_k = generation_top_k # top k for each query, for pretraining, set to 1
@@ -226,12 +192,6 @@ class COCOM(PreTrainedModel):
         self.sep = cfg.sep
         self.compr_rate = cfg.compr_rate
         self.local_rank = os.getenv('LOCAL_RANK', '0')
-        for layer in self.decoder.encoder.layer:
-            layer.attention.self = CustomFlashAttention(
-                embed_dim=cfg.hidden_size,
-                num_heads=cfg.num_attention_heads,
-                dropout=cfg.attention_probs_dropout_prob,
-            )
 
     def compress_and_replace_emb(self, enc_input_ids, enc_attention_mask, dec_input_ids):
         indices = range(0, enc_input_ids.size(0) + 1, self.generation_top_k)
@@ -348,4 +308,7 @@ class COCOM(PreTrainedModel):
             'dec_attention_mask': inp_dec['attention_mask'].to(self.decoder.device)
         }
 
-        return self.generate(model_input, max_new_tokens)
+        return self.generate(model_input, max_new_tokens)
+        
+
+