[Proposal] Change Mistral's config to reduce context size from 32k to 4k

[collingray]

Proposal

The current TL config for Mistral sets its context at 32k, which results in ~34gb worth of attention masks being allocated (32768*32768*32 layers), putting it out of range for all consumer-grade single-gpu setups. As described in #479, a more memory efficient implementation of attention masks would be ideal, but in the short term it can just be capped at 2048/4096. If the longer context is needed, it could instead be parameterized, but I suspect that 4k is enough, at least for the short term.

I can put up a pr if needed.

Checklist

• I have checked that there is no similar issue in the repo (required)

[neelnanda-io]

This would be great! My ideal would be to change the default to be 2K, and also add a parameter to from_pretrained called context_size_override that allows you to override the default.

@ojh31 This may be why you were having memory issues with Mistral 7B!

I think longer-term we should just get rid of the attention mask attribute, it seems wasteful and can just be made at inference time, it should be pretty cheap.

[collingray]

Ok, I'll get a pr up for this