Support rope cache indexing using positions

[acisseJZhong]

Add support to indexing rope cache using position_ids, this might be needed during

1. inference, where we passed in position_ids into transformer forward
2. CP load balancing where we need to index rope cache given positions ids

Test:
running dpskv3 16b base

also tested in https://github.com/wwwjn/torchtitan/pull/1/files when passing position_ids

[tianyu-l]

Thanks! May I ask: for the dsv3 16B test on dp4tp2 before vs. after, did you explicitly pass positions into the model? We should try both (1, seq_len) and (batch_size, seq_len) inputs (they could be the trivial 0 -> seq_len - 1 ids).

Also had an inline comment.

[tianyu-l]

Note that when positions is not None, this is making implicit assumption that positions has the the expected sharding when it's used, namely
sharded on batch dim by DP, replicate on TP mesh, sharded on seq dim by CP

I don't have a good solution right now -- making it Replicate by default will fail here when positions is None https://github.com/pytorch/pytorch/blob/main/torch/distributed/tensor/parallel/style.py#L521
but clearly this is leaving a footgun. I'd suggest we add a comment for now.