Update modeling_llama.py

modeling_llama.py

@@ -318,7 +318,7 @@ class GroupedAutoEncoder(nn.Module):
     def forward(self, x):
         # Split input into groups
         group_inputs = torch.split(x, self.group_input_dim, dim=2)
-        # import pdb; pdb.set_trace()
+
         # Apply group-wise encoding
         encoded_groups = [encoder(group) for group, encoder in zip(group_inputs, self.encoders)]
 
@@ -370,8 +370,8 @@ class LlamaAttention(nn.Module):
         input_dim = 5120
         hidden_dim = 320
         num_groups = 40
-        # self.ae_v = AutoEncoder(input_dim, hidden_dim)#.cuda()
-        self.ae_v = GroupedAutoEncoder(input_dim=input_dim, hidden_dim=hidden_dim, num_groups=num_groups)# .cuda()
+        self.ae_v = GroupedAutoEncoder(input_dim=input_dim, hidden_dim=hidden_dim, num_groups=num_groups)
+        self.load_ae_v = True
         #self.ae_v.eval()
         
     def _init_rope(self):
@@ -446,7 +446,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
+            value_states_ = value_states.clone()
+        else:
+            reuse = False
+        
         kv_seq_len = key_states.shape[-2]
         if past_key_value is not None:
             if self.layer_idx is None:
@@ -458,14 +467,15 @@ class LlamaAttention(nn.Module):
             kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
         cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
         query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
-
+        
         if past_key_value is not None:
             cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
+            # print(value_states.shape)
             key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
 
         key_states = repeat_kv(key_states, self.num_key_value_groups)
         value_states = repeat_kv(value_states, self.num_key_value_groups)
-
+        
         attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
 
         if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
@@ -485,11 +495,19 @@ 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)
         
-        # import pdb; pdb.set_trace()
-        if attn_weights.shape[2]>576:
-            # print("loading ... ")
-            #print(value_states.shape)
+        
+        #if self.layer_idx==5:
+        #    print(value_states[0,0,256,:])
+        
+        if self.load_ae_v:
             self.ae_v.load_state_dict(torch.load("weights_group_320/"+"autoencoder_epoch_1_L1_nonorm_layer_"+str(self.layer_idx)+".pth", map_location='cuda'))
+            self.load_ae_v = False
+        else:
+            pass
+        
+        #if self.layer_idx==5:
+        #    print(value_states.shape)
+        if value_states.shape[2]>576:
             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)
@@ -498,7 +516,12 @@ class LlamaAttention(nn.Module):
             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
+        
+        # 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):
@@ -1480,4 +1503,4 @@ class LlamaForSequenceClassification(LlamaPreTrainedModel):
             past_key_values=transformer_outputs.past_key_values,
             hidden_states=transformer_outputs.hidden_states,
             attentions=transformer_outputs.attentions,
-        )
+        )