[BUG] RoPE ignores position IDs
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
**Describe the bug**
If using `--reset-position-ids`, the RoPE implementation does not take this into account; it will still use the embeddings from position 0 to sequence length - 1.
**To Reproduce**
As you can see in [`language_model.py`](https://github.com/NVIDIA/Megatron-LM/blob/0609f27fe8376f17ab65c001d3d8f35cd8175950/megatron/model/language_model.py#L491-L496), the generated rotary embedding tensor does not take the `enc_position_ids` into account.
**Expected behavior**
From the rotary embedding tensor, the indices in the sequence dimension for the given `enc_position_ids` should be taken instead. For `enc_position_ids = torch.arange(0, seq_len)`, this would return the same (correct) values as it does now. However, then it would also handle other arbitrary position IDs. It's also important that this handles `dec_position_ids` separately.
**Proposed fix**
Most simply, this would be implemented in the forward pass of `RotaryEmbedding` by taking position IDs instead of a sequence length as its only argument. Again, it's also really important to create a new rotary embedding tensor for the decoder part of a encoder-decoder transformer using `dec_input_ids`.
Contributor guide
Assessment
This issue has not been assessed yet.