Allow MPT models to return attention weights

mosaicml · Sep 15, 2023 · a4f5679 · a4f5679
1 parent 7ec2fe0
commit a4f5679
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/llmfoundry/models/layers/blocks.py b/llmfoundry/models/layers/blocks.py
@@ -91,6 +91,7 @@ def forward(
         attn_bias: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.ByteTensor] = None,
         is_causal: bool = True,
+        output_attentions: bool = False,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[
             torch.Tensor, torch.Tensor]]]:
         a = self.norm_1(x)
@@ -100,6 +101,7 @@ def forward(
             attn_bias=attn_bias,
             attention_mask=attention_mask,
             is_causal=is_causal,
+            needs_weights=output_attentions,
         )
         x = x + self.resid_attn_dropout(b)
         m = x

diff --git a/llmfoundry/models/mpt/modeling_mpt.py b/llmfoundry/models/mpt/modeling_mpt.py
@@ -434,6 +434,7 @@ def forward(
                 attn_bias=attn_bias,
                 attention_mask=attention_mask,
                 is_causal=self.is_causal,
+                output_attentions=output_attentions == True,
             )
             if past_key_values is not None:
                 past_key_values[b_idx] = past_key_value