Update modeling_llama_no_vcache.py

modeling_llama_no_vcache.py

@@ -414,28 +414,16 @@ class LlamaAttention(nn.Module):
         query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
         key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
         value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-       
+        
+        
         # import pdb; pdb.set_trace()
         
         if value_states.shape[2]>576:
             reuse = True
-            if self.load_ae_v:
-                self.ae_v.load_state_dict(torch.load("weights/"+"autoencoder_epoch_1_L1_1280_nonorm_layer_"+str(self.layer_idx)+".pth", map_location='cuda'))
-                self.load_ae_v = False
-            else:
-                pass
             value_states_ = value_states.clone()
-            value_states_v = value_states[:,:,35:35+576,:]
-            value_states_v = value_states_v.permute(0, 2, 1, 3)
-            value_states_v=value_states_v.reshape(value_states_v.shape[0],value_states_v.shape[1],5120)
-            # import pdb; pdb.set_trace()
-            value_states_v = self.ae_v(value_states_v)
-            value_states_v = value_states_v.reshape(value_states_v.shape[0],value_states_v.shape[1], 40, 128)
-            value_states_v = value_states_v.permute(0, 2, 1, 3)
-            value_states[:,:,35:35+576,:] = value_states_v
         else:
             reuse = False
-            
+        
         kv_seq_len = key_states.shape[-2]
         if past_key_value is not None:
             if self.layer_idx is None:
@@ -475,8 +463,7 @@ class LlamaAttention(nn.Module):
         attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
         attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
         
-        if reuse:
-            value_states = value_states_
+        
         #if self.layer_idx==5:
         #    print(value_states[0,0,256,:])
         
@@ -486,8 +473,9 @@ class LlamaAttention(nn.Module):
         else:
             pass
         
+        #if self.layer_idx==5:
+        #    print(value_states.shape)
         if value_states.shape[2]>576:
-            value_states_ = value_states.clone()
             value_states_v = value_states[:,:,35:35+576,:]
             value_states_v = value_states_v.permute(0, 2, 1, 3)
             value_states_v=value_states_v.reshape(value_states_v.shape[0],value_states_v.shape[1],5120)
@@ -497,6 +485,11 @@ class LlamaAttention(nn.Module):
             value_states_v = value_states_v.permute(0, 2, 1, 3)
             value_states[:,:,35:35+576,:] = value_states_v
         
+        if reuse:
+            value_states = value_states_
+            
+        #if self.layer_idx==5:
+        #    print(value_states[0,0,256,:])
         attn_output = torch.matmul(attn_weights, value_states)
 
         if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):