Remove redundant transposes for rope rotation

[ShashankMosaicML]

Now that we use transformers v4.36.0, we do not need to transpose query and keys back and forth for apply rope rotations to them. This PR fixes this.

We can see that the loss and mfu plots are almost identical

[vchiley]

Should we gate on HF version? Or is that the min version of HF that we install?

[ShashankMosaicML]

Should we gate on HF version? Or is that the min version of HF that we install?

Yes, the minimum version of HF we install now is 4.36 (

llm-foundry/setup.py

Line 52 in 06b9a1f

'transformers>=4.36,<4.37',

). Hence, we do not need to gate since the 'unsqueeze_dim' argument was added in version 4.35.

[dakinggg]

@ShashankMosaicML Lets gate on transformers version please. Since this is in the model code, it'd be good to make it compatible with all (or as much as possible) transformers versions

[ShashankMosaicML]

@ShashankMosaicML Lets gate on transformers version please. Since this is in the model code, it'd be good to make it compatible with all (or as much as possible) transformers versions

Done.