Allow MPT models to return attention weights

[lorabit110]

Previously, output_attentions is not propagated to the attention layer. Even when it's set to True, no attention weights were returned.

[dakinggg]

Could you please add a simple unit test for this?

[vchiley]

This error should be noted if you plan to use it.

[lorabit110]

Could you please add a simple unit test for this?

Updated an existing unit test to check attention.

[lorabit110]

This error should be noted if you plan to use it.

Yeah. I am aware of that flash attention won't return attention weights.

[dakinggg]

One comment to make sure we're returning the right shape of stuff, but otherwise lgtm! Thanks for the PR!

[dakinggg]

@lorabit110 could you run precommit run --all-files locally? Thanks!