Fix version compatibility issue with transformers>4.34.0 for flash-attention2 patch

[Trangle]

Why are these changes needed?

1. the rotary_emb logic changed in transformers==4.35.0, fix the compatibility

Related issue number (if applicable)

#2648

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]

does this work for older versions of transformers?

[Trangle]

does this work for older versions of transformers?

test in 4.34 and 4.35，4.35 is when the logic of rotary_emb changed.

[Trangle]

does this work for older versions of transformers?

Also tested in 4.30.

However, there are some questions. Do we not need to perform a restore operation on the kv head here?

about after line # 72

if getattr(self, "num_key_value_groups", None):
k = repeat_kv(k, self.num_key_value_groups)
v= repeat_kv(v, self.num_key_value_groups)

[Niyx52094]

how long is this feature ok for users? When I use this file for transformers 4.35.0 to fine tune llama2 7b in 8X A100, it gives an error about " out of memory ".

[Trangle]

how long is this feature ok for users? When I use this file for transformers 4.35.0 to fine tune llama2 7b in 8X A100, it gives an error about " out of memory ".

This is an issue with transformers after ver 4.35. After reducing the batch size, try again. Currently, this issue has not been fixed in 4.36 yet. You can try fixed versions below 4.35, such as 4.34.