Skip to content

[PyTorch][Common] Refactor RoPE - #1626

Merged
yaox12 merged 8 commits into
NVIDIA:mainfrom
yaox12:xiny/refactor_rope
Apr 7, 2025
Merged

[PyTorch][Common] Refactor RoPE#1626
yaox12 merged 8 commits into
NVIDIA:mainfrom
yaox12:xiny/refactor_rope

Conversation

@yaox12

@yaox12yaox12 commented Mar 31, 2025

Copy link
Copy Markdown
Member

Description

Refactor RoPE to

  1. Now both fused and non-fused RoPE support CP + SBHD/BSDH/THD.
  2. When CP > 1, the freqs tensor is required to be the full tensor.
    1. Previously, the freqs tensor is supposed to be the full tensor with the THD format while sliced (outside of TE) with other formats (sbhd/bshd). This PR unifies them.
    2. This is not a breaking change, because the previous apply_rotary_pos_emb doesn't accept cp_size and cp_rank for sbhd and bshd, so passing sliced freqs with cp_size = 1 and cp_rank = 0 should still work as before.
  3. Both fused and non-fused RoPE support interleaved mode.

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

Please list the changes introduced in this PR:

  • Change A
  • Change B

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

@yaox12
refactor to add cp support for sbhd/bshd
5129fb6
Signed-off-by: Xin Yao <xiny@nvidia.com>
@yaox12
support interleaved
d727f1c
Signed-off-by: Xin Yao <xiny@nvidia.com>
@yaox12
yaox12force-pushed the xiny/refactor_rope branch from 2525377 to d727f1cCompareMarch 31, 2025 01:17
@yaox12
format
18ca2f3
Signed-off-by: Xin Yao <xiny@nvidia.com>
yaox12
yaox12 commented Mar 31, 2025
View reviewed changes
Comment threadtransformer_engine/pytorch/dot_product_attention/rope.py
@tomlifu

tomlifu commented Mar 31, 2025

Copy link
Copy Markdown
Contributor

Rotary interleaved part looks good to me.

@yaox12

yaox12 commented Apr 1, 2025

Copy link
Copy Markdown
MemberAuthor

/te-ci pytorch

@yaox12yaox12 mentioned this pull request Apr 2, 2025
Merged
6 tasks
sudhakarsingh27
sudhakarsingh27 reviewed Apr 4, 2025
View reviewed changes
Comment threadtests/pytorch/test_fused_rope.py
Comment threadtests/pytorch/test_fused_rope.py
Comment threadtransformer_engine/pytorch/dot_product_attention/rope.py
if tensor_format == "sbhd":
output = tex.fused_rope_forward(t, freqs, False)
elif tensor_format == "bshd":
output = tex.fused_rope_forward(t.transpose(0, 1), freqs, True).transpose(0, 1)

@sudhakarsingh27sudhakarsingh27Apr 3, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

curious where the t.transpose happens now

@yaox12yaox12Apr 7, 2025•
edited
Loading

Copy link
Copy Markdown
MemberAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Previously, the C++ interface fused_rope_forward only accepts sbhd format, so for bshd we hack it by transposing bshd to sbhd and then transposing the result back. There was a flag called transpose_output_memory to make sure the transpose of result doesn't really do a memory copy. This method is not intuitive. Now I add a new argument qkv_format to support all formats natively.

Comment threadtransformer_engine/pytorch/dot_product_attention/rope.py
yaox12 added 3 commits April 6, 2025 18:58
@yaox12
Merge branch 'main' into xiny/refactor_rope
d8d8b39
@yaox12
update
e0ffb65
Signed-off-by: Xin Yao <xiny@nvidia.com>
@yaox12
merge sbhd/bshd and thd functions
757a39b
Signed-off-by: Xin Yao <xiny@nvidia.com>
@yaox12

yaox12 commented Apr 7, 2025

Copy link
Copy Markdown
MemberAuthor

/te-ci pytorch

@yaox12

yaox12 commented Apr 7, 2025

Copy link
Copy Markdown
MemberAuthor

CI passed except for irrelevant failures and Blackwell runners got stuck.

@yaox12
yaox12 merged commit ba605f1 into NVIDIA:mainApr 7, 2025
wdykas pushed a commit to wdykas/TransformerEngine that referenced this pull request Apr 14, 2025
@yaox12@wdykas
[PyTorch][Common] Refactor RoPE (NVIDIA#1626)
b37b2b2
* refactor to add cp support for sbhd/bshd
Signed-off-by: Xin Yao <xiny@nvidia.com>
* support interleaved
Signed-off-by: Xin Yao <xiny@nvidia.com>
* format
Signed-off-by: Xin Yao <xiny@nvidia.com>
* add interleaved to RotaryPositionEmbedding in test
Signed-off-by: Xin Yao <xiny@nvidia.com>
* update
Signed-off-by: Xin Yao <xiny@nvidia.com>
* merge sbhd/bshd and thd functions
Signed-off-by: Xin Yao <xiny@nvidia.com>
---------
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: Peter Dykas <wdykas@nvidia.com>
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@sudhakarsingh27sudhakarsingh27sudhakarsingh27 left review comments
@cyanguwacyanguwacyanguwa approved these changes
+1 more reviewer
@xrennvidiaxrennvidiaxrennvidia approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

2.3.0

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@yaox12@tomlifu@sudhakarsingh27@cyanguwa@xrennvidia