Fix Rope Compatibility with Cos/Sin Position Embedding for Batch Size > 1

[wizyoung]

Summary

Fix Rope Compatibility with Cos/Sin Position Embedding for Batch Size > 1

This PR addresses an issue with the compatibility of the ROPE implementation when using cosine/sine position embeddings with a batch size greater than 1.
In the default behavior of transformers, position_ids is set to None during training, which results in the following computation:

cache_position = torch.arange(seq_len)
position_ids = cache_position.unsqueeze(0)

This leads to the shape of the position embeddings being (1, seq_len, head_dim), which is consistent with the implementation in Liger.
However, if position_ids are pre-calculated for any reason(In my experiment, I implement m-rope in another different way making position_ids pre-calculated), the current implementation fails to handle this scenario correctly. This PR introduces a fix to ensure that the ROPE implementation can accommodate pre-computed position_ids.
In the unit test test_ropy.py, I have added a variable expand_position_ids to simulate this condition. The previous implementation fails under this scenario, while the new patch successfully resolves the issue.

pytest details:

Testing Done

• Hardware Type:
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence