update bugs

modeling_stablelm_epoch.py

@@ -529,23 +529,23 @@ class DecoderLayer(nn.Module):
         )
         hidden_states = residual + hidden_states
 
-        # # Cross Attention
-        # residual = hidden_states
+        # Cross Attention
+        residual = hidden_states
         
-        # bsz, q_len, _ = hidden_states.size()
-        # _, kv_len, _ = cross_states.size()
-
-        # cross_attn_mask = torch.zeros((bsz, 1, kv_len, q_len), device=hidden_states.device)
-        # hidden_states, cross_attn_weights, _ = self.cross_attn(
-        #     hidden_states=hidden_states,
-        #     cross_states=cross_states,
-        #     attention_mask=cross_attn_mask,
-        #     position_ids=position_ids,
-        #     past_key_value=past_key_value,
-        #     output_attentions=output_attentions,
-        #     use_cache=use_cache,
-        # )
-        # hidden_states = residual + hidden_states
+        bsz, q_len, _ = hidden_states.size()
+        _, kv_len, _ = cross_states.size()
+
+        cross_attn_mask = torch.zeros((bsz, 1, kv_len, q_len), device=hidden_states.device)
+        hidden_states, cross_attn_weights, _ = self.cross_attn(
+            hidden_states=hidden_states,
+            cross_states=cross_states,
+            attention_mask=cross_attn_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+        )
+        hidden_states = residual + hidden_states
 
         # Fully Connected
         residual = hidden_states