Fixes for FlashAttention

[tmm1]

Why are these changes needed?

We need to mark LlamaRMSNorm layers as bf16 after peft conversion to fix errors from flash-attn about dtype support

I also updated the flash attention patch

I am able to run training on llama2 with this change.

Related issue number (if applicable)

Closes #1828

Checks

• I've run format.sh to lint the changes in this PR.
• I've included any doc changes needed.
• I've made sure the relevant tests are passing (if applicable).

[merrymercy]

@tmm1 Thanks! This is merged.

[philschmid]

This code looks quite similar to the patch in my blog We know noticed that 70B models with GQA are not supported.

Have you seen the same issue?

[tmm1]

Hi @philschmid, yes you could refer to this approach:

LAION-AI/Open-Assistant@3c8f93e

[philschmid]

Yeah saw the commit as well. Just wanted to share here that you are aware its not working for 70B atm.

[tmm1]

@philschmid I'm looking at this again now to add 70B support. Did you end up doing any more work in this area?

I'm also interested in making the patch work for forward-pass with past_key_value support, which is something that's still not quite working.

[merrymercy]

@tmm1 Hi, I did some minor style cleanup in PR #2212.
I also found the current implementation does not support generative inference or 70B.
Could you take a look and fix it? Thanks!