NVIDIA / NVIDIA/Megatron-LM

[BUG] RoPE ignores position IDs

Open
#453 6 comments 0 reactions 0 assignees View on GitHub
bug
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.